"""Momentum: a long track keeps going through a disturbance that only adds something to one side, and what it adds goes to
the track it touches. After tracking, before the cut rule (swapped in at split_by_fit; tracking, Fourier and the cut rule
unchanged; no package change):

1. Each solid track of at least MIN_TRACK layers is fitted along the height with a low-degree polynomial (DEG, on its
   rings' coefficients, start points by continuity), and refitted on the layers it follows (shape miss under TAU layer
   heights), ITER times. The layers it does not follow are disturbed; the clean body there is interpolated (PCHIP per
   coefficient) from the undisturbed layers on either side (the momentum from below and above).
2. At a disturbed layer, if the ring is the clean ring plus an outward lobe (the clean ring lies inside the ring but for
   OUT_FRAC of its area, and the ring's largest piece outside the clean ring holds LOBE_FRAC of its area), the track's ring
   becomes ring AND clean ring, and the lobe (that piece, widened by OVERLAP layer heights into the body, within the ring)
   is peeled off; a lobe reaches MIN_DEPTH layer heights out and meets the clean ring along at most MAX_CONTACT of its
   perimeter (a band all around is not one). The decoder takes the union of overlapping solid parts.
3. A run of peeled layers is attached to the track it touches: one that starts on the layer after the run (or ends on the
   layer before it) with a ring there that meets the run's last (first) lobe; with none, the lobes are a part of their own.
Meshes only, default N, M, KAPPA. Writes ../report/momentum.csv and renders ../report/img/momentum_<object>_<config>_a<atoms>.png.

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

import numpy as np
import shapely
import trimesh
from numpy.polynomial import chebyshev as npcheb
from scipy.interpolate import PchipInterpolator
from shapely.geometry import Polygon
from shapely.geometry.polygon import orient

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

MIN_TRACK, DEG, TAU, ITER = 30, 10, 0.5, 4
OUT_FRAC, LOBE_FRAC, OVERLAP = 0.03, 0.005, 1.5
MIN_DEPTH, MAX_CONTACT = 2.0, 0.4  # a lobe reaches this many layer heights out, along at most this share of the clean ring
BUDGETS = [256, 1024, 6144]
OUT = REPORT / os.environ.get("FCP_OUT", "momentum.csv")
FIELDS = ["object", "config", "atoms", "parts", "peeled", "worst_lh", "over_025", "chamfer", "p95", "n_mean", "n_over30", "seconds"]
_split = _parts.split_by_fit
LOG = {}


def poly(P):
    g = Polygon(P)
    return g if g.is_valid else g.buffer(0)


def largest(g):
    if g.is_empty:
        return None
    parts = list(g.geoms) if hasattr(g, "geoms") else [g]
    parts = [p for p in parts if p.geom_type == "Polygon"]
    return max(parts, key=lambda p: p.area) if parts else None


def coeffs_of(polys, K, ccw):
    """Fourier coefficients (R, 2, K) of shapely polygons, oriented like the track's rings."""
    arrs = []
    for p in polys:
        p = orient(p, 1.0 if ccw else -1.0)
        arrs.append(np.asarray(p.exterior.coords, float)[:-1])
    return _parts._ring_coeffs(arrs, _parts.N_SAMPLES, (K - 1) // 2).astype(float)


def is_lobe(lb, clean):
    """Something sticking out locally, not a band around the ring: it reaches MIN_DEPTH layer heights beyond the clean
    ring, and it meets the clean ring along at most MAX_CONTACT of its perimeter."""
    depth = max(clean.exterior.distance(shapely.Point(c)) for c in np.asarray(lb.exterior.coords)[::4])
    contact = lb.buffer(0.25 * LH).intersection(clean.exterior).length
    return depth >= MIN_DEPTH * LH and contact <= MAX_CONTACT * clean.exterior.length


def clean_fit(C, A):
    """The track's clean body: its rings' coefficients fitted along the height on the layers the fit follows."""
    n, _, K = C.shape
    Cn = _parts.continuous_start(C, A)
    X = Cn.reshape(n, -1)
    u = _parts.cell_u(n)
    V = npcheb.chebvander(u, min(DEG, n - 1))
    ok = np.ones(n, bool)
    R = rings(C)
    for _ in range(ITER):
        coef, *_ = np.linalg.lstsq(V[ok], X[ok], rcond=None)
        F = (V @ coef).reshape(n, 2, K)
        miss = miss_batch(R, rings(F)) / LH
        new = miss <= TAU
        if new.sum() < DEG + 2 or np.array_equal(new, ok):
            break
        ok = new
    # the clean body at the disturbed layers: interpolated (PCHIP per coefficient) from the undisturbed layers on either
    # side, so a far disturbance (a rim) does not bend it; at the undisturbed layers the rings themselves
    idx = np.arange(n)
    if ok.sum() >= 2 and (~ok).any():
        Fx = X.copy()
        Fx[~ok] = PchipInterpolator(idx[ok], X[ok], axis=0, extrapolate=True)(idx[~ok])
        F = Fx.reshape(n, 2, K)
    return F, ~ok


def layer_peel(ring, clean):
    """(body, lobe, strong) at one disturbed layer, or None: the ring is the clean ring plus an outward lobe; strong when the
    lobe also passes `is_lobe` (a run is peeled from its first strong layer to its last: weak layers only between them)."""
    if clean.difference(ring).area > OUT_FRAC * clean.area:
        return None
    lb = largest(ring.difference(clean))
    if lb is None or lb.area < LOBE_FRAC * ring.area:
        return None
    body = largest(ring.intersection(clean))
    lobe = largest(lb.buffer(OVERLAP * LH).intersection(ring))
    if body is None or lobe is None:
        return None
    return body, lobe, is_lobe(lb, clean)


def momentum_split(parts, coeffs, areas, M, tol=None, follows=None):
    parts = [dict(p) for p in parts]
    coeffs = [np.asarray(C, float).copy() for C in coeffs]
    areas = [list(a) for a in areas]
    K = coeffs[0].shape[2] if coeffs else 1
    peeled_total = 0
    runs = []  # (track, first offset, [(body, lobe)], ccw)
    for t, (p, C, A) in enumerate(zip(list(parts), list(coeffs), list(areas))):
        if p["sign"] < 0 or len(C) < MIN_TRACK:
            continue
        F, bad = clean_fit(C, np.asarray(A))
        R, Rf = rings(C), rings(F)
        ccw = np.sum(R[0, :, 0] * np.roll(R[0, :, 1], -1) - np.roll(R[0, :, 0], -1) * R[0, :, 1]) > 0
        cand = [layer_peel(poly(R[j]), poly(Rf[j])) if bad[j] else None for j in range(len(C))]
        j = 0
        while j < len(C):  # runs: within a stretch of layers that all peel, from its first strong layer to its last
            if cand[j] is None:
                j += 1
                continue
            k = j
            while k < len(C) and cand[k] is not None:
                k += 1
            strong = [i for i in range(j, k) if cand[i][2]]
            if strong:  # from the first strong layer to the last: weak layers only fill the gaps between strong ones
                a, b = strong[0], strong[-1] + 1
                runs.append((t, a, [c[:2] for c in cand[a:b]], ccw))
            j = k
    # attach each run's lobes to the track it touches, or make them a part of their own (runs of 3 layers or more); then the
    # body's rings of the runs kept become ring AND clean ring
    kept = []
    for t, j0, pcs, ccw in runs:
        h0 = parts[t]["h0"] + j0
        h1 = h0 + len(pcs) - 1
        lobes = [lb for _, lb in pcs]
        L = coeffs_of(lobes, K, ccw)
        La = [lb.area for lb in lobes]
        done = False
        for q, (pq, Cq) in enumerate(zip(parts, coeffs)):
            if q == t or pq["sign"] < 0:
                continue
            if pq["h0"] == h1 + 1 and poly(rings(Cq[:1])[0]).intersects(lobes[-1]):
                coeffs[q] = np.concatenate([L, coeffs[q]])
                areas[q] = La + areas[q]
                pq["h0"] = h0
                done = True
                break
            if pq["h1"] == h0 - 1 and poly(rings(Cq[-1:])[0]).intersects(lobes[0]):
                coeffs[q] = np.concatenate([coeffs[q], L])
                areas[q] = areas[q] + La
                pq["h1"] = h1
                done = True
                break
        if not done and len(pcs) >= 3:
            parts.append({"sign": 1, "h0": h0, "h1": h1})
            coeffs.append(L)
            areas.append(La)
            done = True
        if done:
            js = list(range(j0, j0 + len(pcs)))
            coeffs[t][js] = coeffs_of([b for b, _ in pcs], K, ccw)
            for jj, (b, _) in zip(js, pcs):
                areas[t][jj] = b.area
            peeled_total += len(js)
            kept.append((h0, h1))
    LOG["peeled"] = peeled_total
    LOG["runs"] = kept
    return _split(parts, [c.astype(_parts.F) for c in coeffs], areas, M, tol, follows)


def main():
    names = sys.argv[1:] or ["mug"]
    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 (("package", _split), ("momentum", momentum_split)):
                LOG.clear()
                _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), peeled=LOG.get("peeled", 0),
                                    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"momentum{os.environ.get('FCP_TAG', '')}_{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} peeled={LOG.get('peeled', 0):3d} runs={LOG.get('runs', [])} "
                      f"worst={wm.max():.2f} lh  parts: {[(s, a, b) for s, a, b in d.parts]}", flush=True)


if __name__ == "__main__":
    main()