"""The package's gwn and sink (v5.31) against the prototypes they came from, on the 20-cloud sample of ablation.py
(per object clean 10k and one other cloud type): parts, Chamfer (ablation.py's stats) and encode time, CPU and CUDA.
Prototype rows: report/ablation.csv ("gwn": gwn.py direct sums; "local": surf_local.py + pointfit.py). Writes
report/v531_check.csv.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2710_tests/scripts/v531_check.py
"""
import csv
import sys
import time
from pathlib import Path

import numpy as np
import trimesh

sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, "/home/joao/code/fcp_devel/api")
import fcp  # noqa: E402
from ablation import _sample  # noqa: E402
from make_clouds import PN  # noqa: E402
from make_noisy import label  # noqa: E402
from parts_report import OBJ, REPORT, ROOT, UP, stats  # noqa: E402

OUT = REPORT / "v531_check.csv"
PROTO = {"gwn": "gwn", "sink": "local"}


def main():
    A = {(r["object"], r["cloud"], r["points"], r["config"]): r for r in csv.DictReader((REPORT / "ablation.csv").open())}
    rows = []
    for name, tag, n in sorted(_sample(list(OBJ))):
        g = trimesh.load(OBJ[name], force="mesh")
        ref = np.asarray(g.sample(50_000, seed=1))
        sigma = {f"pn{label(v)}": v for v in PN}.get(tag, 0.0) * float(np.linalg.norm(g.extents))
        P = np.asarray(trimesh.load(ROOT / "pointclouds" / f"{name}_{tag}_{n}.ply").vertices)
        for sl in ("gwn", "sink"):
            p = A[(name, tag, n, PROTO[sl])]
            for dev in ("cuda", "cpu"):
                t = time.perf_counter()
                d = fcp.encode(P, up=UP.get(name, (0, 0, 1)), device=dev, slicer=sl, noise=sigma)
                sec = time.perf_counter() - t
                e = stats(d, ref, fcp.decode(d, device="cuda"))[0]
                rows.append(dict(object=name, cloud=tag, points=n, slicer=sl, device=dev, parts=len(d.parts),
                                 proto_parts=int(p["parts"]), mean=e, proto_mean=float(p["mean"]), seconds=round(sec, 2),
                                 proto_seconds=float(p["seconds"])))
                r = rows[-1]
                print(f"{name:16s} {tag:9s} {n:>5s} {sl:4s} {dev:4s} parts {r['parts']:3d} (proto {r['proto_parts']:3d})  "
                      f"err {e * 1e3:6.2f} (proto {r['proto_mean'] * 1e3:6.2f})  {sec:5.2f}s (proto {r['proto_seconds']:5.2f}s)",
                      flush=True)
    with OUT.open("w", newline="") as f:
        w = csv.DictWriter(f, list(rows[0]))
        w.writeheader()
        w.writerows(rows)


if __name__ == "__main__":
    main()