"""Run: env -u PYTHONPATH .venv/bin/python ~/code/data/meshes/fcp_2810_partition/scripts/body_track.py elephant > ../report/elephant_body.txt
The elephant's body track: per layer, rings in the slice, the ring's area and perimeter, the step of its coefficients
from the layer below (Parseval-weighted, start points by continuity), and where tracked-fit cuts."""
import sys
import numpy as np
import trimesh
sys.path.insert(0, "/home/joao/code/fcp_devel/api")
sys.path.insert(0, "/home/joao/code/data/meshes/fcp_2810_partition/scripts")
from degree_sweep import OBJ
import fcp
from fcp import _parts

CAP = {}
_sys, _split = _parts.parts_system, _parts.split_by_fit
def system(geoms, *a, **k):
    CAP["S"] = geoms
    return _sys(geoms, *a, **k)
def split(parts, coeffs, areas, M, tol=None, follows=None):
    i = int(np.argmax([len(c) for c in coeffs]))
    CAP["p"], CAP["C"], CAP["A"] = parts[i], np.asarray(coeffs[i], float), np.asarray(areas[i], float)
    out = _split(parts, coeffs, areas, M, tol, follows)
    CAP["cuts"] = sorted(q["h0"] for q in out[0] if q["sign"] > 0 and parts[i]["h0"] < q["h0"] <= parts[i]["h1"])
    return out
_parts.parts_system, _parts.split_by_fit = system, split
import fcp.encode as E
if hasattr(E, "parts_system"): E.parts_system = system
g = trimesh.load(OBJ[sys.argv[1] if len(sys.argv) > 1 else "elephant"], force="mesh")
d = fcp.encode(g, up=(0, 0, 1), device="cuda")
S, p, C, A = CAP.get("S"), CAP["p"], CAP["C"], CAP["A"]
K = C.shape[2]; w = np.full(K, 0.5); w[0] = 1
Cw = (_parts.continuous_start(C, A) * np.sqrt(w)).reshape(len(C), -1)
step = np.r_[0, np.linalg.norm(np.diff(Cw, axis=0), axis=1)]
Cn = _parts.continuous_start(C, A)
# perimeter from coefficients: sample the ring
t = np.linspace(0, 2 * np.pi, 512, endpoint=False)
k = np.arange(1, (K - 1) // 2 + 1)
def ring(c):
    x = c[0, 0] + (c[0, 1::2][:, None] * np.cos(k[:, None] * t) + c[0, 2::2][:, None] * np.sin(k[:, None] * t)).sum(0)
    y = c[1, 0] + (c[1, 1::2][:, None] * np.cos(k[:, None] * t) + c[1, 2::2][:, None] * np.sin(k[:, None] * t)).sum(0)
    return x, y
per = np.array([np.linalg.norm(np.diff(np.stack(ring(c), 1)[np.r_[0:512, 0]], axis=0), axis=1).sum() for c in Cn])
nr = [len(S.per_height(~S.hole)[h]) for h in range(p["h0"], p["h1"] + 1)] if S is not None else [0] * len(C)
print(f"track layers {p['h0']}-{p['h1']}, cuts at {CAP['cuts']}")
print(" layer rings   area  perim   step  cut")
for j, h in enumerate(range(p["h0"], p["h1"] + 1)):
    print(f"{h:6d} {nr[j]:5d} {A[j]:6.3f} {per[j]:6.3f} {step[j]:6.3f}  {'<--' if h in CAP['cuts'] else ''}")

# per stretch between the body's real events: the fit's worst miss in coefficients (what split_by_fit tests) against the
# worst shape distance (symmetric nearest-point RMS between the fitted ring and the sliced ring)
from scipy.spatial import cKDTree
from shape_test import to_polyline  # noqa: E402  (each ring's samples to the other's polyline: no sampling floor)
def shape_rms(a, b):
    P, Q = np.stack(ring(a), 1), np.stack(ring(b), 1)
    return np.sqrt(np.mean(np.r_[to_polyline(P, Q), to_polyline(Q, P)] ** 2))
h0 = p["h0"]
for a, b in [(14, 26), (31, 50), (36, 50), (52, 89), (52, 70), (61, 89), (36, 89)]:
    i, j = a - h0, b - h0 + 1
    for kap in (1e2, 1e4):
        _parts.KAPPA = kap; _parts.height_degree.cache_clear(); _parts.orth_basis.cache_clear()
        M = 24 if kap == 1e2 else 50
        Q, _ = _parts.orth_basis(j - i, M)
        X = Cw[i:j]; Xf = Q @ (Q.T @ X)
        coef_miss = np.sqrt(((X - Xf) ** 2).sum(1))
        Cf = (Xf.reshape(-1, 2, K) / np.sqrt(w))
        shp = np.array([shape_rms(Cn[i + r], Cf[r]) for r in range(j - i)])
        print(f"layers {a:3d}-{b:3d} ({j - i:3d}) kappa {kap:g} deg {Q.shape[1] - 1:2d}: coef miss max {coef_miss.max():.4f}  shape rms max {shp.max():.4f}")