"""A smaller dictionary for a small budget: N (harmonics per ring) and M (height degree) lowered, tracked-fit at the
default cap (KAPPA 1e2; the cut test fits with the same basis, so a lower M also changes the parts). Each encoding at 6144
atoms (or its whole dictionary, if smaller), decoded at its 256 / 1024 / 2048 prefixes. Meshes only.
Writes ../report/nm_sweep.csv and renders ../report/img/nm_<object>_N<N>_M<M>_a1024.png for RENDER.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2810_partition/scripts/nm_sweep.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))
from degree_sweep import DEVICE, OBJ, REPORT, UP, Ref  # noqa: E402  (puts the package on the path)
import fcp  # noqa: E402
from renders import IMG, render  # noqa: E402

NM = [(16, 24), (12, 16), (8, 12), (8, 8), (6, 6), (4, 4), (16, 16), (16, 12), (16, 8)]  # the last three: angular detail kept, height degree lowered
ONLY = [tuple(int(v) for v in a.split(',')) for a in sys.argv[1:]]  # run only these (N,M), appending to the CSV
BUDGETS = [256, 1024, 2048]
RENDER = ("dino", "elephant", "horse", "bowl_gso")
OUT = REPORT / "nm_sweep.csv"
FIELDS = ["object", "N", "M", "atoms", "stored", "parts", "chamfer", "p95", "n_mean", "n_over30", "seconds"]


def main():
    IMG.mkdir(parents=True, exist_ok=True)
    new = not ONLY or not OUT.exists()
    with OUT.open("w" if new else "a", newline="") as f:
        w = csv.writer(f)
        if new:
            w.writerow(FIELDS)
        for name, path in OBJ.items():
            g = trimesh.load(path, force="mesh")
            ref = Ref(g)
            for N, M in (ONLY or NM):
                t = time.perf_counter()
                d = fcp.encode(g, atoms=6144, up=UP, N=N, M=M, device=DEVICE)
                sec = time.perf_counter() - t
                stored = len(d.atoms[0])
                for T in BUDGETS:
                    dt = d.prefix(T)
                    m = fcp.decode(dt, device=DEVICE)
                    w.writerow([name, N, M, T, len(dt.atoms[0]), len(d.parts), *[round(float(v), 5) for v in ref.measure(m)],
                                round(sec, 3)])
                    if T == 1024 and name in RENDER:
                        render(dt, m, IMG / f"nm_{name}_N{N}_M{M}_a{T}.png",
                               f"{name}: N {N}, M {M}, {len(dt.atoms[0])} atoms, {len(d.parts)} parts")
                f.flush()
                print(f"{name:16s} N={N:2d} M={M:2d} parts={len(d.parts):3d} stored={stored}", flush=True)


if __name__ == "__main__":
    main()