"""Cut where the shape jumps between neighbouring rings, and nowhere else: along each track, the step s_i between ring i
and ring i + 1 (ring_steps.py: RMS distance, each one's T samples to the other's polyline, both ways; layer heights); a
cut between them where s_i is at least SPIKE_K times the median of the steps within WINDOW layers on either side (itself
left out) and at least SPIKE_ABS layer heights. Pieces down to 2 layers (a single ring is never left alone: of two
neighbouring spikes the larger cuts). "spike": nothing else. "spike+fit": then each piece whose full fit still misses a
ring by more than BACK layer heights in shape is cut by split_by_fit's own search with that shape test (its FIT_MIN).
Against the package ("greedy"). Swapped in for split_by_fit (no package change). Meshes only, default N, M, KAPPA.
Writes ../report/spike_split.csv and renders ../report/img/spike_<object>_<config>_a<atoms>.png.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2810_partition/scripts/spike_split.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

SPIKE_K, SPIKE_ABS, WINDOW = 2.5, 1.0, 4
BACK = 0.25  # layer heights
BUDGETS = [256, 1024, 6144]
RENDER_AT = (1024,)
OUT = REPORT / "spike_split.csv"
FIELDS = ["object", "config", "atoms", "parts", "worst_lh", "over_025", "chamfer", "p95", "n_mean", "n_over30", "seconds"]
_split = _parts.split_by_fit


def spike_cuts(C):
    """Cut layers (offsets in the track) where the ring-to-ring step spikes."""
    R = rings(np.asarray(C, float))
    n = len(R)
    if n < 4:
        return []
    s = miss_batch(R[:-1], R[1:]) / LH  # s[i]: between ring i and ring i + 1
    spikes = []
    for i in range(len(s)):
        nb = np.r_[s[max(0, i - WINDOW):i], s[i + 1:i + 1 + WINDOW]]
        if s[i] >= SPIKE_ABS and len(nb) and s[i] >= SPIKE_K * np.median(nb):
            spikes.append(i)
    cuts = []
    for i in sorted(spikes, key=lambda i: -s[i]):  # largest first; a cut at i + 1 leaves rings 0..i and i + 1..
        c = i + 1
        edges = sorted([0, n, *cuts, c])
        if all(b - a >= 2 for a, b in zip(edges[:-1], edges[1:])):
            cuts.append(c)
    return sorted(cuts)


def shape_follows(sliced, tol):
    def follows(i, a, b, Cf):
        return miss_batch(sliced[i][a:b], rings(np.asarray(Cf, float))).max() <= tol
    return follows


def make_split(back):
    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)
            edges = [0, *spike_cuts(C), len(C)]
            for a, b in zip(edges[:-1], edges[1:]):
                q = {"sign": p["sign"], "h0": p["h0"] + a, "h1": p["h0"] + b - 1}
                if back is None:
                    pp, cc, aa = [q], [C[a:b]], [list(A[a:b])]
                else:
                    sl = [rings(C[a:b])]
                    pp, cc, aa = _split([q], [C[a:b]], [list(A[a:b])], M, follows=shape_follows(sl, back * LH))
                out_p += pp
                out_c += [np.asarray(c, _parts.F) for c in cc]
                out_a += [list(x) for x in aa]
        return out_p, out_c, out_a
    return split


CONFIGS = {"greedy": None, "spike": make_split(None), "spike+fit": make_split(BACK)}


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
                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"spike_{name}_{cfg}_a{T}.png", f"{name}: {cfg}, {T} atoms, {len(d.parts)} parts")
                f.flush()
                print(f"{name:16s} {cfg:10s} parts={len(d.parts):3d} worst={wm.max():.2f} lh  {sec:.1f} s", flush=True)


if __name__ == "__main__":
    main()