"""How much of the contour energy each budget holds. Energy: the Parseval-weighted squared distance summed over all rings
(sum of w Y^2; the selection's own measure). Split into: held by the first T atoms; in the dictionary but dropped; outside
the dictionary (the sliced rings minus each part's full fit). Also the RMS contour error per ring left at T, in layer
heights, and the same without the rings' means (centroids) and first harmonics (radius / ellipse), which dominate the
energy. Writes ../report/energy.csv.

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

import numpy as np
import trimesh

sys.path.insert(0, str(Path(__file__).parent))
from degree_sweep import DEVICE, OBJ, REPORT, UP  # noqa: E402  (puts the package on the path)
import fcp  # noqa: E402
from fcp import _parts  # noqa: E402
from fcp._select import Selection  # noqa: E402
from fcp._slice import centres  # noqa: E402

BUDGETS = [256, 1024, 2048, 6144]


def main():
    rows = []
    print(f"{'object':16s} {'dict':>6s} {'outside':>8s} " + " ".join(f"{'held@' + str(T):>10s}" for T in BUDGETS)
          + "   RMS error per ring (layer heights) @" + "/".join(map(str, BUDGETS)) + "   same, harmonics >= 2 only")
    for name, path in OBJ.items():
        g = trimesh.load(path, force="mesh")
        cap, sysf = {}, _parts.parts_system

        def capture(*a, **k):
            cap["out"] = sysf(*a, **k)
            return cap["out"]

        _parts.parts_system = capture
        d = fcp.encode(g, atoms=1, up=UP, device=DEVICE)
        _parts.parts_system = sysf
        P, Y, Phi, w = cap["out"]
        Y, Phi, w = np.asarray(Y, float), np.asarray(Phi, float), np.asarray(w, float)
        K = d.K
        lh = float(centres(d.layers)[1] - centres(d.layers)[0])
        rings = sum(p["h1"] - p["h0"] + 1 for p in P)
        total = float((w[:, None] * Y ** 2).sum())
        proj = np.einsum("bhj,bh->bj", Phi, Y)  # every atom's value (orthonormal atoms)
        score2 = w[:, None] * proj ** 2  # the energy each atom holds
        inside = float(score2.sum())
        harm = np.arange(len(w)) % K  # block -> harmonic slot (0 mean, 1-2 harmonic 1, ...)
        hi = harm >= 3  # harmonics 2 and up
        order = np.argsort(-score2.ravel(), kind="stable")
        s_sorted = score2.ravel()[order]
        hi_sorted = np.repeat(hi, score2.shape[1])[order]
        total_hi = float((w[hi, None] * Y[hi] ** 2).sum())
        dictsize = int(np.count_nonzero(np.abs(Phi).sum(1)))
        line = f"{name:16s} {dictsize:6d} {100 * (total - inside) / total:7.4f}% "
        held, rms, rms_hi = [], [], []
        for T in BUDGETS:
            h = float(s_sorted[:T].sum())
            held.append(100 * h / total)
            rms.append(np.sqrt((total - h) / rings) / lh)
            rms_hi.append(np.sqrt((total_hi - float(s_sorted[:T][hi_sorted[:T]].sum())) / rings) / lh)
            rows.append(dict(object=name, atoms=T, dictionary=dictsize, outside_pct=100 * (total - inside) / total,
                             held_pct=100 * h / total, rms_lh=rms[-1], rms_hi_lh=rms_hi[-1]))
        line += " ".join(f"{v:9.4f}%" for v in held) + "   " + " / ".join(f"{v:.2f}" for v in rms) + "   " + " / ".join(f"{v:.2f}" for v in rms_hi)
        print(line, flush=True)
    with (REPORT / "energy.csv").open("w", newline="") as f:
        wr = csv.DictWriter(f, list(rows[0]))
        wr.writeheader()
        wr.writerows(rows)


if __name__ == "__main__":
    main()