"""The rate-distortion partition on the 251-scan set (api/bench/scans.py's objects, loaded and repaired as there): the v6
package's cut rule against rd_split's "rdT256", each scan scored against itself (Chamfer and normals over 30°, Ref of
degree_sweep), parts, encode time, at 256 / 1024 / 6144 atoms. Scans that do not encode are recorded as skipped.
Writes ../report/scans_rd.csv (resumes).

    FCP_API=/home/joao/code/fcp_devel_v6/api env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python \\
        ~/code/data/meshes/fcp_2910_v6/scripts/scans_rd.py [--shard i k]
"""
import csv
import os
import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, "/home/joao/code/data/meshes/fcp_2810_partition/scripts")  # the shared helpers: objects, measures, renders
from degree_sweep import DEVICE, Ref  # noqa: E402  (puts the package on the path)
import fcp  # noqa: E402

sys.path.insert(0, str(Path(__file__).parent))
import rd_split  # noqa: E402

sys.path.insert(0, os.path.join(os.environ.get("FCP_API", "/home/joao/code/fcp_devel/api"), "bench"))
from common import UP, load_file, scan_files  # noqa: E402

REPORT = Path(__file__).resolve().parent.parent / "report"
OUT = REPORT / os.environ.get("FCP_OUT", "scans_rd.csv")
CONFIGS = os.environ.get("FCP_CONFIGS", "package,rdT256").split(",")  # rd_split modes, or parts= values ("rd", "tracked-fit")
BUDGETS = [256, 1024, 6144]
FIELDS = ["i", "name", "config", "atoms", "parts", "chamfer", "p95", "n_over30", "seconds", "status"]


def main():
    shard = [0, 1]
    if "--shard" in sys.argv:
        k = sys.argv.index("--shard")
        shard = [int(sys.argv[k + 1]), int(sys.argv[k + 2])]
    done = set()
    if OUT.exists():
        done = {(int(r["i"]), r["config"]) for r in csv.DictReader(OUT.open())}
    new = not OUT.exists()
    with OUT.open("a", newline="") as f:
        w = csv.DictWriter(f, FIELDS)
        if new:
            w.writeheader()
        for i, name, path in scan_files()[shard[0]::shard[1]]:
            if all((i, c) in done for c in CONFIGS):
                continue
            try:
                m = load_file(path)
                ref = Ref(m)
            except Exception as e:
                w.writerow(dict(i=i, name=name, config="-", status=f"load: {type(e).__name__}"))
                continue
            for cfg in CONFIGS:
                if (i, cfg) in done:
                    continue
                kw = {"parts": cfg} if cfg in ("rd", "tracked-fit") else {}
                rd_split.install("package" if kw else cfg)
                try:
                    t = time.perf_counter()
                    d = fcp.encode(m, atoms=max(BUDGETS), up=UP, device=DEVICE, **kw)
                    sec = time.perf_counter() - t
                    for T in BUDGETS:
                        v = ref.measure(fcp.decode(d.prefix(T), device=DEVICE))
                        w.writerow(dict(i=i, name=name, config=cfg, atoms=T, parts=len(d.parts), chamfer=round(float(v[0]), 5),
                                        p95=round(float(v[1]), 5), n_over30=round(float(v[3]), 3), seconds=round(sec, 3), status="ok"))
                except Exception as e:
                    w.writerow(dict(i=i, name=name, config=cfg, status=f"{type(e).__name__}: {str(e)[:60]}"))
                finally:
                    rd_split.install("package")
            f.flush()
            print(i, name, flush=True)


if __name__ == "__main__":
    main()