"""Tracking with rings split at their necks: before the package tracks the slices, every solid ring that has a neck is cut
into two rings there, so a lobe growing out of a body (a mug's handle root, a leg joining a torso) is its own ring from
the layer its neck forms, and tracking can carry it on as its own part; the body stays one smooth ring. The package's
tracking, Fourier and cut rule then run unchanged on the new slices (swapped in at parts_system; no package change).

Neck: the ring resampled at NS points by arc length; over all pairs of points, the chord d against the shorter arc s
between them; the pair of least d / s is a neck if d / s < NECK_Q, the chord lies inside the ring (NCHK points on it),
and both pieces hold at least MIN_FRAC of the ring's area. Pieces are split again (a ring with several lobes).
Meshes only, default N, M, KAPPA. Writes ../report/neck_split.csv and renders ../report/img/neck_<object>_a<atoms>.png.

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

import numpy as np
import trimesh
from matplotlib.path import Path as MPath

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, _rings  # noqa: E402
from dp_split import worst_miss  # noqa: E402
from renders import IMG, render  # noqa: E402

NS, NECK_Q, MIN_FRAC, NCHK = 512, 0.3, 0.03, 16
BUDGETS = [256, 1024, 6144]
OUT = REPORT / "neck_split.csv"
FIELDS = ["object", "config", "atoms", "parts", "tracks", "splits", "worst_lh", "over_025", "chamfer", "p95", "n_mean",
          "n_over30", "seconds"]
_system = _parts.parts_system
STATS = {}


def resample(P, n=NS):
    Q = np.vstack([P, P[:1]])
    s = np.r_[0, np.cumsum(np.linalg.norm(np.diff(Q, axis=0), axis=1))]
    t = np.linspace(0, s[-1], n, endpoint=False)
    return np.c_[np.interp(t, s, Q[:, 0]), np.interp(t, s, Q[:, 1])]


def area(P):
    return 0.5 * abs((P[:, 0] * np.roll(P[:, 1], -1) - np.roll(P[:, 0], -1) * P[:, 1]).sum())


def split_neck(P, depth=0):
    """[ring] or the ring cut at its neck into pieces (each split again)."""
    if depth > 4 or len(P) < 16:
        return [P]
    R = resample(P)
    n = len(R)
    seg = np.linalg.norm(R - np.roll(R, -1, 0), axis=1)
    per = seg.sum()
    cum = np.r_[0, np.cumsum(seg)][:n]
    i, j = np.triu_indices(n, k=8)  # pairs at least 8 samples apart along the ring
    arc = cum[j] - cum[i]
    arc = np.minimum(arc, per - arc)
    d = np.linalg.norm(R[i] - R[j], axis=1)
    q = d / np.maximum(arc, 1e-12)
    path = MPath(R)
    A = area(R)
    for k in np.argsort(q)[:64]:  # the narrowest candidates first
        if q[k] >= NECK_Q:
            break
        a, b = int(i[k]), int(j[k])
        chord = R[a] + np.linspace(0.05, 0.95, NCHK)[:, None] * (R[b] - R[a])
        if not path.contains_points(chord).all():
            continue
        lobe, body = R[a + 1:b], np.vstack([R[b:], R[:a + 1]])  # the lobe leaves out the neck's two points
        if min(area(lobe), area(body)) < MIN_FRAC * A or len(lobe) < 8 or len(body) < 8:
            continue
        STATS["splits"] = STATS.get("splits", 0) + 1
        return split_neck(body, depth + 1) + split_neck(lobe, depth + 1)
    return [P]


def system(S, *a, **k):
    """parts_system on the slices with every solid ring split at its necks."""
    loops = [[] for _ in range(S.n_heights)]
    for r in range(len(S.h)):
        P = np.asarray(S.rings.ring(r), float)
        loops[S.h[r]] += [P] if S.hole[r] else split_neck(P)
    S2 = _rings.slices(loops, "cpu")
    out = _system(S2, *a, **k)
    STATS["tracks"] = STATS.get("tracks", 0)
    return out


def tracks_of(g, sysf):
    """The number of tracks (parts with parts='tracked')."""
    _parts.parts_system = sysf
    try:
        return len(fcp.encode(g, atoms=1, up=UP, parts="tracked", device=DEVICE).parts)
    finally:
        _parts.parts_system = _system


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, sysf in (("package", _system), ("neck", system)):
                STATS.clear()
                _parts.parts_system = sysf
                t = time.perf_counter()
                try:
                    d = fcp.encode(g, atoms=6144, up=UP, device=DEVICE)
                finally:
                    _parts.parts_system = _system
                sec = time.perf_counter() - t
                splits = STATS.get("splits", 0)
                ntr = tracks_of(g, sysf)
                _parts.parts_system = sysf
                try:
                    wm = worst_miss(g, d, _parts.split_by_fit)
                finally:
                    _parts.parts_system = _system
                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), tracks=ntr, splits=splits,
                                    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 (1024, 6144):
                        render(dt, m, IMG / f"neck_{name}_{cfg}_a{T}.png",
                               f"{name}: {'rings split at necks' if cfg == 'neck' else 'package'}, {T} atoms, {len(d.parts)} parts "
                               f"({ntr} tracks)")
                f.flush()
                print(f"{name:16s} {cfg:8s} tracks={ntr:3d} parts={len(d.parts):3d} splits={splits:4d} worst={wm.max():.2f} lh  "
                      f"{sec:.1f} s", flush=True)


if __name__ == "__main__":
    main()