"""The knee of the atom energy (docs/noisy_clouds.md section 3, fcp_devel): atoms in selection order, Parseval energy w v^2
(w = 1 for the k = 0 coefficient, 1/2 otherwise); |v| (running median, 64) on log-log stays high, drops, then decays as a noise
tail. Rule tested here: fit a line to the tail (log |v| against log T over the last half of the atoms); the knee is the first
atom after which the running median stays within KNEE_TOL (x) of that line.

For one object: gwn + 2D closing (CLOSE 0.75, sigma given) at 6144 atoms on its point-noise clouds (pn0p12, pn0p6, pn1p2,
10k) and the mesh route on the clean mesh; Chamfer and SDF normals on a ladder of prefixes and at the knee ->
../report/knee.csv (appends), descriptors ../report/knee/<object>_<cloud>.fcp, energy curves ../report/knee/<object>.npz.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python knee.py horse
"""
import csv
import sys
from pathlib import Path

import numpy as np
import trimesh

sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, "/home/joao/code/data/meshes/fcp_2710_tests/scripts")
import gwn  # noqa: E402
import fcp  # noqa: E402
import fcp.tools  # noqa: E402
from parts_report import OBJ, ROOT, stats  # noqa: E402
from bench import sigma_of  # noqa: E402

REP = Path(__file__).parent.parent / "report"
KD = REP / "knee"
ATOMS = 6144
WIN = 64
KNEE_TOL = 1.25
CLOUDS = (("pn0p12", 10000), ("pn0p6", 10000), ("pn1p2", 10000))
LADDER = np.unique(np.round(np.geomspace(32, ATOMS, 16)).astype(int))


def curves(d):
    """(atom energies, running median of |v|, tail energy fraction) in selection order."""
    i, v = d.atoms
    w = np.where((np.asarray(i) // (d.M + 1)) % d.K == 0, 1.0, 0.5)
    v = np.asarray(v, float)
    e = w * v * v
    rm = np.array([np.median(np.abs(v[max(0, k - WIN):k + 1])) for k in range(len(v))])
    tail = np.cumsum(e[::-1])[::-1] / e.sum()
    return e, rm, tail


def knee(rm):
    """Kneedle on log |v| (running median) against log T over [WIN, 95 % of the atoms] (the end of a descriptor drops):
    the atom farthest below the chord. Also the slope of the curve after the knee (the tail)."""
    T = np.arange(1, len(rm) + 1)
    a, b = WIN - 1, max(WIN + 1, int(0.95 * len(rm)))
    x, y = np.log(T[a:b]), np.log(np.maximum(rm[a:b], 1e-30))
    chord = y[0] + (y[-1] - y[0]) * (x - x[0]) / (x[-1] - x[0])
    k = int(a + np.argmax(chord - y) + 1)
    sl = float(np.polyfit(x[k - a - 1:], y[k - a - 1:], 1)[0]) if b - k > 8 else float("nan")
    return k, sl


def score(d, ref, Ps, Ns):
    m = fcp.decode(d, device="cuda")
    if m is None:
        return float("nan"), float("nan"), float("nan")
    Ne = fcp.tools.Shape(d, device="cuda").normal(Ps, method="sdf")
    Ne /= np.maximum(np.linalg.norm(Ne, axis=1, keepdims=True), 1e-12)
    ang = np.degrees(np.arccos(np.clip((Ne * Ns).sum(1), -1, 1)))
    return 1e3 * stats(d, ref, m)[0], float(np.median(ang)), 100 * float((ang < 10).mean())


def main(name):
    KD.mkdir(parents=True, exist_ok=True)
    g = trimesh.load(OBJ[name], force="mesh")
    diag = float(np.linalg.norm(g.extents))
    ref = np.asarray(g.sample(50_000, seed=1))
    Ps, fi = trimesh.sample.sample_surface(g, 3000, seed=3)
    Ns = g.face_normals[fi]
    gwn.CLOSE = 0.75
    runs = [("mesh", fcp.encode(g, device="cuda", up=(0, 0, 1), atoms=ATOMS))]
    for tag, n in CLOUDS:
        P = np.asarray(trimesh.load(ROOT / "pointclouds" / f"{name}_{tag}_{n}.ply").vertices)
        runs.append((f"{tag}_{n}", gwn.encode_gwn(P, up=(0, 0, 1), sigma=sigma_of(tag, diag), atoms=ATOMS)))
        if tag != "pn0p12":  # sigma = 0: the slicer keeps the noise (the setting of noisy_clouds.md: noise in the atoms)
            runs.append((f"{tag}_{n}_s0", gwn.encode_gwn(P, up=(0, 0, 1), sigma=0.0, atoms=ATOMS)))
    rows, npz = [], {}
    for cl, d in runs:
        fcp.save(d, KD / f"{name}_{cl}.fcp")
        e, rm, tail = curves(d)
        k, slope = knee(rm)
        npz[f"{cl}_rm"], npz[f"{cl}_tail"], npz[f"{cl}_knee"] = rm, tail, np.array([k, slope])
        Ts = sorted(set(int(t) for t in LADDER if t <= d.n_atoms) | {k, d.n_atoms})
        for T in Ts:
            err, nm, p10 = score(d.prefix(T), ref, Ps, Ns)
            rows.append(dict(object=name, cloud=cl, atoms=T, stored=d.n_atoms, is_knee=int(T == k), knee=k, tail_slope=slope,
                             energy=float(1 - tail[T] if T < len(tail) else 1.0), parts=len(d.parts), mean=err, n_med=nm,
                             pgp10=p10))
        kr = next(r for r in rows if r["cloud"] == cl and r["is_knee"])
        fr = next(r for r in rows if r["cloud"] == cl and r["atoms"] == d.n_atoms)
        br = min((r for r in rows if r["cloud"] == cl), key=lambda r: r["mean"] if r["mean"] == r["mean"] else 1e9)
        print(f"{name:16s} {cl:12s} knee {k:5d} ({100 * kr['energy']:.2f} %, tail slope {slope:.2f}) err {kr['mean']:.2f} "
              f"normal {kr['n_med']:.1f} | full {fr['mean']:.2f} / {fr['n_med']:.1f} | best {br['atoms']} {br['mean']:.2f} / {br['n_med']:.1f}",
              flush=True)
    np.savez(KD / f"{name}.npz", **npz)
    out = REP / "knee.csv"
    new = not out.exists()
    with out.open("a", newline="") as fh:
        w = csv.DictWriter(fh, list(rows[0]))
        if new:
            w.writeheader()
        w.writerows(rows)


if __name__ == "__main__":
    main(sys.argv[1])