"""The largest dictionary that fits a budget whole: per object and N, the largest M whose dictionary (every part's
2 (2N + 1) (d + 1) values, d = its height degree) holds at most BUDGET atoms, the parts encoded at that N and M (they
depend on both). Writes ../report/budget_fit.csv.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2810_partition/scripts/budget_fit.py [budget]
"""
import csv
import sys
from pathlib import Path

import trimesh

sys.path.insert(0, str(Path(__file__).parent))
from degree_sweep import DEVICE, OBJ, REPORT, UP  # noqa: E402  (puts the package on the path)
import fcp  # noqa: E402
from fcp import _parts  # noqa: E402

BUDGET = int(sys.argv[1]) if len(sys.argv) > 1 else 2048
NS = [16, 12, 8, 6, 4, 2, 1]


def size(d):
    return sum(2 * d.K * (_parts.height_degree(h1 - h0 + 1, d.M) + 1) for _, h0, h1 in d.parts)


def main():
    rows = []
    for name, path in OBJ.items():
        g = trimesh.load(path, force="mesh")
        d0 = fcp.encode(g, atoms=1, up=UP, device=DEVICE)
        line = [f"{name:16s} default: {len(d0.parts):3d} parts, dictionary {size(d0):6d} |"]
        for N in NS:
            best = None
            for M in range(0, 41):
                d = fcp.encode(g, atoms=1, up=UP, N=N, M=M, device=DEVICE)
                if size(d) <= BUDGET:
                    best = (M, size(d), len(d.parts))
                elif best is not None and M > best[0] + 4:
                    break
            rows.append(dict(object=name, budget=BUDGET, N=N, M=best[0] if best else "", dictionary=best[1] if best else "",
                             parts=best[2] if best else ""))
            line.append(f"N{N}: M {best[0]:2d} ({best[2]}p)" if best else f"N{N}: -")
        print(" ".join(line), flush=True)
    with (REPORT / f"budget_fit_{BUDGET}.csv").open("w", newline="") as f:
        w = csv.DictWriter(f, list(rows[0]))
        w.writeheader()
        w.writerows(rows)


if __name__ == "__main__":
    main()