"""tracked-fit with one change: a fit follows its part when every ring's shape is within the tolerance (nearest-point RMS
between the fitted ring and the sliced ring, both ways), not its coefficients (which also count the ring's start sliding
along it, a tangential drift that moves no surface; report: "Why the elephant is cut so much"). Tracking, the cut search,
the step rule, FIT_TOL, FIT_MIN and FIT_GAIN unchanged; plugged in through split_by_fit's `follows` (no package change).
Meshes only, default cap (KAPPA 1e2, M 24). Writes ../report/shape_test.csv and renders ../report/img/<object>_shape-fit_*.png.

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

import numpy as np
import trimesh
from scipy.spatial import cKDTree

sys.path.insert(0, str(Path(__file__).parent))
from degree_sweep import BUDGETS, DEVICE, OBJ, REPORT, UP, Ref  # noqa: E402  (puts the package on the path)
import fcp  # noqa: E402
from fcp import _parts  # noqa: E402
from renders import IMG, render  # noqa: E402

OUT = REPORT / "shape_test.csv"
FIELDS = ["object", "test", "atoms", "parts", "mean_len", "chamfer", "p95", "n_mean", "n_over30", "seconds"]
NT = 256  # samples per ring
_split = _parts.split_by_fit
_t = np.linspace(0, 2 * np.pi, NT, endpoint=False)


def rings(C):
    """(R, 2, K) Fourier coefficients -> (R, NT, 2) ring points."""
    K = C.shape[-1]
    k = np.arange(1, (K - 1) // 2 + 1)
    B = np.concatenate([np.ones((1, NT)), np.cos(k[:, None] * _t), np.sin(k[:, None] * _t)])  # (K, NT), reordered below
    order = np.r_[0, 1 + 2 * np.arange(len(k)), 2 + 2 * np.arange(len(k))]  # c[0], cos 1..N, sin 1..N
    return np.einsum("rck,kt->rtc", C[..., order], B)


def to_polyline(P, Q):
    """Distance of each point P to the closed polyline Q: its nearest vertex's two segments (the nearest vertex alone
    leaves a floor of up to half a sample spacing, as large as the tolerance on a big ring)."""
    j = cKDTree(Q).query(P)[1]
    d = np.full(len(P), np.inf)
    for s in (-1, 0):  # the segments (j-1, j) and (j, j+1)
        a, b = Q[(j + s) % len(Q)], Q[(j + s + 1) % len(Q)]
        ab = b - a
        u = np.clip(np.einsum("ij,ij->i", P - a, ab) / np.maximum(np.einsum("ij,ij->i", ab, ab), 1e-30), 0, 1)
        d = np.minimum(d, np.linalg.norm(P - a - u[:, None] * ab, axis=1))
    return d


def shape_miss(A, B):
    """Per ring: RMS distance between rings A and B (R, NT, 2), each one's samples to the other's polyline, both ways."""
    out = np.empty(len(A))
    for r, (P, Q) in enumerate(zip(A, B)):
        out[r] = np.sqrt(np.mean(np.r_[to_polyline(P, Q), to_polyline(Q, P)] ** 2))
    return out


def shape_split(parts, coeffs, areas, M, tol=None, follows=None):
    tol = _parts.FIT_TOL if tol is None else tol
    sliced = [rings(np.asarray(C, float)) for C in coeffs]

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

    return _split(parts, coeffs, areas, M, tol, follows=follows_shape)


TESTS = {"coef": _split, "shape": shape_split}


def main():
    IMG.mkdir(parents=True, exist_ok=True)
    with OUT.open("w", newline="") as f:
        w = csv.writer(f)
        w.writerow(FIELDS)
        for name, path in OBJ.items():
            g = trimesh.load(path, force="mesh")
            ref = Ref(g)
            for test, fn in TESTS.items():
                _parts.split_by_fit = fn
                t = time.perf_counter()
                d = fcp.encode(g, atoms=max(BUDGETS), up=UP, parts="tracked-fit", device=DEVICE)
                sec = time.perf_counter() - t
                _parts.split_by_fit = _split
                lens = [h1 - h0 + 1 for _, h0, h1 in d.parts]
                for T in BUDGETS:
                    m = fcp.decode(d.prefix(T), device=DEVICE)
                    w.writerow([name, test, T, len(d.parts), round(float(np.mean(lens)), 1),
                                *[round(float(v), 5) for v in ref.measure(m)], round(sec, 3)])
                    if test == "shape" and T in (1024, 6144):
                        render(d.prefix(T), m, IMG / f"{name}_shape-fit_k100_M24_a{T}.png",
                               f"{name}: tracked-fit, shape test, κ 100, M 24, {T} atoms, {len(d.parts)} parts")
                f.flush()
                print(f"{name:16s} {test:6s} parts={len(d.parts):3d} {sec:.2f} s", flush=True)


if __name__ == "__main__":
    main()