"""The fewest parts for a tolerance: each track partitioned by dynamic programming into the fewest pieces (at least FIT_MIN
layers; a track shorter than 2 FIT_MIN stays whole, as in split_by_fit) whose full fit (the piece's own height basis) misses
no ring by more than the tolerance; among those, the least total squared miss (which puts the cuts on the events). The miss
is "shape" (RMS distance, each ring's samples to the other's polyline, both ways; the tolerance in layer heights) or "coef"
(split_by_fit's own measure: RMS in coefficients, Parseval, start points by continuity; the tolerance FIT_TOL). Against the
package's greedy split ("greedy"). Swapped in for split_by_fit (no package change). Meshes only, default N, M, KAPPA.
Writes ../report/dp_split.csv and renders ../report/img/dp_<object>_<config>_a<atoms>.png.

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

LH = float(centres(128)[1] - centres(128)[0])  # layer height (canonical), 128 layers
CONFIGS = {"greedy": None, "dp-coef": ("coef", _parts.FIT_TOL)}
CONFIGS.update({f"dp-shape-{t:g}": ("shape", t * LH) for t in (0.1, 0.15, 0.2, 0.3)})
BUDGETS = [256, 1024, 6144]
RENDER_AT = (1024, 6144)
OUT = REPORT / "dp_split.csv"
FIELDS = ["object", "config", "atoms", "parts", "worst_lh", "over_025", "chamfer", "p95", "n_mean", "n_over30", "seconds"]
_split = _parts.split_by_fit
_cache = {}


def miss_batch(A, B):
    """Per ring: RMS distance between rings A and B (R, T, 2), each one's samples to the other's polyline, both ways;
    brute force over the samples, then the nearest vertex's two segments."""
    def one_way(P, Q):
        d2 = (P ** 2).sum(-1)[:, :, None] + (Q ** 2).sum(-1)[:, None, :] - 2 * np.einsum("rtc,rsc->rts", P, Q)
        j = np.argmin(d2, axis=2)  # (R, T)
        T = Q.shape[1]
        r = np.arange(len(P))[:, None]
        best = np.full(j.shape, np.inf)
        for s in (-1, 0):
            a, b = Q[r, (j + s) % T], Q[r, (j + s + 1) % T]
            ab = b - a
            u = np.clip(((P - a) * ab).sum(-1) / np.maximum((ab * ab).sum(-1), 1e-30), 0, 1)
            best = np.minimum(best, np.linalg.norm(P - a - u[..., None] * ab, axis=-1))
        return best
    return np.sqrt(np.mean(np.concatenate([one_way(A, B), one_way(B, A)], 1) ** 2, axis=1))


def piece_misses(C, A, M, kind):
    """miss[(i, j)]: per ring of the piece i..j-1, the full fit's miss (canonical units). Cached per track."""
    key = (hashlib.sha1(np.ascontiguousarray(C).tobytes()).hexdigest(), M, kind)
    if key in _cache:
        return _cache[key]
    n, K = len(C), C.shape[2]
    w = np.full(K, 0.5)
    w[0] = 1.0
    Cn = _parts.continuous_start(C, A)
    sliced = rings(np.asarray(C, float)) if kind == "shape" else None
    out = {}
    lengths = [n] if n < 2 * _parts.FIT_MIN else range(_parts.FIT_MIN, n + 1)
    for L in lengths:
        Q, _ = _parts.orth_basis(L, M)
        for i in range(0, n - L + 1):
            j = i + L
            if n >= 2 * _parts.FIT_MIN and (i and i < _parts.FIT_MIN or j < n and n - j < _parts.FIT_MIN):
                continue  # the rest of the track before / after could not be a piece
            X = Cn[i:j].reshape(L, -1)
            Xf = Q @ (Q.T @ X)
            if kind == "coef":
                out[i, j] = np.sqrt((((X - Xf).reshape(L, 2, K) ** 2) * w).sum((1, 2)))
            else:
                out[i, j] = miss_batch(sliced[i:j], rings(Xf.reshape(L, 2, K)))
    _cache[key] = out
    return out


def dp_cuts(n, miss, tol):
    """The fewest pieces whose misses are all within tol; ties: the least total squared miss."""
    best = {0: (0, 0.0, None)}
    for j in range(1, n + 1):
        cand = []
        for i in best:
            m = miss.get((i, j))
            if m is not None and m.max() <= tol:
                cand.append((best[i][0] + 1, best[i][1] + float((m ** 2).sum()), i))
        if not cand:  # nothing passes: the piece that misses least (only when no piece of any length fits)
            for i in best:
                m = miss.get((i, j))
                if m is not None:
                    cand.append((best[i][0] + 1, best[i][1] + 1e6 * float(m.max()), i))
        if cand:
            best[j] = min(cand)
    cuts, j = [], n
    while j:
        i = best[j][2]
        cuts.append(i)
        j = i
    return sorted(cuts)[1:]


def make_split(kind, tol):
    def split(parts, coeffs, areas, M, tol_=None, follows=None):
        out_p, out_c, out_a = [], [], []
        for p, C, A in zip(parts, coeffs, areas):
            C = np.asarray(C, float)
            cuts = dp_cuts(len(C), piece_misses(C, A, M, kind), tol)
            edges = [0, *cuts, len(C)]
            for a, b in zip(edges[:-1], edges[1:]):
                out_p.append({"sign": p["sign"], "h0": p["h0"] + a, "h1": p["h0"] + b - 1})
                out_c.append(np.asarray(C[a:b], _parts.F))
                out_a.append(list(A[a:b]))
        return out_p, out_c, out_a
    return split


def worst_miss(g, d, split):
    """Per part of d: its full fit's worst shape miss (layer heights), from the sliced rings the encoder saw."""
    cap, sysf = {}, _parts.parts_system

    def capture(*a, **k):
        cap["out"] = sysf(*a, **k)
        return cap["out"]

    _parts.parts_system, _parts.split_by_fit = capture, split
    try:
        fcp.encode(g, atoms=1, up=UP, device=DEVICE)
    finally:
        _parts.parts_system, _parts.split_by_fit = sysf, _split
    P, Y, _, _ = cap["out"]
    K = d.K
    B = 2 * K
    out = []
    for i, p in enumerate(P):
        n = p["h1"] - p["h0"] + 1
        Ys = Y[i * B:(i + 1) * B, :n].T.astype(float)
        Q, _ = _parts.orth_basis(n, d.M)
        out.append(miss_batch(rings(Ys.reshape(n, 2, K)), rings((Q @ (Q.T @ Ys)).reshape(n, 2, K))).max() / LH)
    return np.array(out)


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, spec in CONFIGS.items():
                split = _split if spec is None else make_split(*spec)
                _parts.split_by_fit = split
                t = time.perf_counter()
                try:
                    d = fcp.encode(g, atoms=6144, up=UP, device=DEVICE)
                finally:
                    _parts.split_by_fit = _split
                sec = time.perf_counter() - t
                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 in RENDER_AT and cfg != "greedy":
                        render(dt, m, IMG / f"dp_{name}_{cfg}_a{T}.png", f"{name}: {cfg}, {T} atoms, {len(d.parts)} parts")
                f.flush()
                print(f"{name:16s} {cfg:16s} parts={len(d.parts):3d} worst={wm.max():.2f} lh  {sec:.1f} s", flush=True)


if __name__ == "__main__":
    main()