"""The package's cut, with its condition made "coefficients AND shape": a piece is cut only when its fit misses a ring by
more than FIT_TOL in coefficients (the points at the same t: the package's test) AND by more than t layer heights in
shape (closest points, each ring's samples to the other's polyline, both ways). Sliding along the ring (large in
coefficients, small in shape) no longer cuts; a real change (large in both) still does. The cut search, its position
rule, FIT_MIN and FIT_GAIN unchanged (split_by_fit's `follows`). Meshes only, default N, M, KAPPA.
Writes ../report/and_test.csv and renders ../report/img/and_<object>_<config>_a<atoms>.png.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2810_partition/scripts/and_test.py [objects...]
"""
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 fcp import _parts  # noqa: E402
from dp_split import LH, miss_batch, worst_miss  # noqa: E402
from renders import IMG, render  # noqa: E402
from shape_test import rings  # noqa: E402

TS = (0.15, 0.2, 0.25, 0.3)  # shape tolerance, layer heights
BUDGETS = [256, 1024, 6144]
OUT = REPORT / "and_test.csv"
FIELDS = ["object", "config", "atoms", "parts", "worst_lh", "over_025", "chamfer", "p95", "n_mean", "n_over30", "seconds"]
_split = _parts.split_by_fit


def make_split(t):
    def split(parts, coeffs, areas, M, tol=None, follows=None):
        tol = _parts.FIT_TOL if tol is None else tol
        K = np.asarray(coeffs[0]).shape[2] if coeffs else 1
        w = np.full(K, 0.5)
        w[0] = 1.0
        Cn = [_parts.continuous_start(np.asarray(C, float), A) for C, A in zip(coeffs, areas)]
        sliced = [rings(np.asarray(C, float)) for C in coeffs]

        def follows_and(i, a, b, Cf):
            coef = np.sqrt((((Cn[i][a:b] - Cf) ** 2) * w).sum((1, 2)))  # the package's miss (Cf: unweighted fit)
            if coef.max() <= tol:
                return True
            return miss_batch(sliced[i][a:b], rings(np.asarray(Cf, float))).max() <= t * LH  # cut only if both miss

        return _split(parts, coeffs, areas, M, tol, follows=follows_and)
    return split


CONFIGS = {"greedy": None, **{f"and-{t:g}": make_split(t) for t in TS}}


def main():
    names = sys.argv[1:] or list(OBJ)
    IMG.mkdir(parents=True, exist_ok=True)
    old = [r for r in csv.DictReader(OUT.open())] if OUT.exists() else []
    keep = [r for r in old if r["object"] not in names]
    with OUT.open("w", newline="") as f:
        w = csv.DictWriter(f, FIELDS)
        w.writeheader()
        w.writerows(keep)
        for name in names:
            g = trimesh.load(OBJ[name], force="mesh")
            ref = Ref(g)
            for cfg, split in CONFIGS.items():
                split = split or _split
                _parts.split_by_fit = split
                t0 = time.perf_counter()
                try:
                    d = fcp.encode(g, atoms=6144, up=UP, device=DEVICE)
                finally:
                    _parts.split_by_fit = _split
                sec = time.perf_counter() - t0
                wm = worst_miss(g, d, split)
                for T in BUDGETS:
                    dt = d.prefix(T)
                    m = fcp.decode(dt, device=DEVICE)
                    w.writerow(dict(object=name, config=cfg, atoms=T, parts=len(d.parts), worst_lh=round(float(wm.max()), 3),
                                    over_025=int((wm > 0.25).sum()), seconds=round(sec, 2),
                                    **dict(zip(["chamfer", "p95", "n_mean", "n_over30"], [round(float(v), 5) for v in ref.measure(m)]))))
                    if T == 1024 and cfg in ("and-0.2", "and-0.25"):
                        render(dt, m, IMG / f"and_{name}_{cfg}_a{T}.png", f"{name}: {cfg}, {T} atoms, {len(d.parts)} parts")
                f.flush()
                print(f"{name:16s} {cfg:9s} parts={len(d.parts):3d} worst={wm.max():.2f} lh  {sec:.1f} s", flush=True)


if __name__ == "__main__":
    main()