"""Decode a prefix holding 95 / 99 / 99.5 % of the atom energy (sum of squared atom values, atoms in selection order) against the
full descriptor: Chamfer (mean symmetric distance to the clean mesh / diagonal, x1e-3) and SDF-normal median / PGP10.
Descriptors: gwn + 2D closing (CLOSE 0.75; sigma given on pn*), 1024 atoms, 128 layers; and the mesh route on the clean mesh.
argv: object; -> appends ../report/energy_prefix.csv."""
import csv, sys
from pathlib import Path
import numpy as np, trimesh
sys.path.insert(0, str(Path(__file__).parent)); sys.path.insert(0, "/home/joao/code/data/meshes/fcp_2710_tests/scripts")
import gwn, fcp, fcp.tools  # noqa
from parts_report import OBJ, ROOT, stats  # noqa
from bench import sigma_of  # noqa

import os
ATOMS = int(os.environ.get("FCP_ATOMS", 1024))  # atoms encoded (the full descriptor)
OUT = Path(__file__).parent.parent / "report" / ("energy_prefix.csv" if ATOMS == 1024 else f"energy_prefix_{ATOMS}.csv")
FRACS = (0.95, 0.99, 0.995, 0.999, 1.0)
FIXED = (1024,) if ATOMS > 1024 else ()  # plain prefixes of this many atoms too (frac = -T)
CLOUDS = (("clean", 10000), ("clean", 2048), ("pn0p12", 10000), ("pn0p12", 2048), ("pn0p6", 10000), ("pn0p6", 2048))
name = sys.argv[1]
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 route", "clean", 0, 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(("gwn+close", tag, n, gwn.encode_gwn(P, up=(0, 0, 1), sigma=sigma_of(tag, diag), atoms=ATOMS)))
rows = []
for cfg, tag, n, d in runs:
    v = np.asarray(d.atoms[1], float)
    cum = np.cumsum(v * v) / max(float((v * v).sum()), 1e-30)
    for fr in FRACS + tuple(-t for t in FIXED):
        T = (len(v) if fr >= 1 else int(np.searchsorted(cum, fr) + 1)) if fr > 0 else min(-fr, len(v))
        dp = d.prefix(T)
        m = fcp.decode(dp, device="cuda")
        if m is None:
            rows.append(dict(object=name, config=cfg, cloud=tag, points=n, frac=fr, atoms=T, parts=len(d.parts), mean=np.nan, n_med=np.nan, pgp10=np.nan))
            continue
        Ne = fcp.tools.Shape(dp, 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)))
        r = dict(object=name, config=cfg, cloud=tag, points=n, frac=fr, atoms=T, parts=len(d.parts), mean=1e3 * stats(dp, ref, m)[0],
                 n_med=float(np.median(ang)), pgp10=100 * float((ang < 10).mean()))
        rows.append(r)
        print(f"{name:16s} {cfg:10s} {tag:6s} {n:6d} {fr:6.3f} atoms {T:5d} err {r['mean']:7.2f} normal {r['n_med']:5.1f} pgp10 {r['pgp10']:5.1f}", flush=True)
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)