"""Stop at the noise floor: the atoms are orthonormal, so the noise of the rings spreads evenly over them; its level is
estimated from the values at the top height degrees of the long parts (median |value| / 0.6745, the wavelet rule), and the
descriptor keeps only the atoms above c times that level (a prefix: the values are unchanged). Measured at the kept count
(capped at 6144) against the package's 1024 and 6144 prefixes, on the meshes and the four cloud kinds.
Writes ../report/noise_floor.csv.

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

import numpy as np
import trimesh

sys.path.insert(0, "/home/joao/code/data/meshes/fcp_2810_partition/scripts")  # the shared helpers: objects, measures, renders
from degree_sweep import DEVICE, OBJ, UP, Ref  # noqa: E402  (puts the package on the path)
import fcp  # noqa: E402

REPORT = Path(__file__).resolve().parent.parent / "report"  # fcp_2910_v6/report
IMG = REPORT / "img"
from fcp import _parts  # noqa: E402

CLOUDS = Path("/home/joao/code/data/meshes/fcp_2710_tests/pointclouds")
INPUTS = ["mesh", "clean_10000", "clean_2048", "noise0p6_10000", "pn0p6_10000"]
CS = (2.0, 3.0, 4.0)
MIN_LEN, TOP = 24, 3  # parts of at least MIN_LEN layers; their TOP highest height degrees
_sys = _parts.parts_system


def noise_level(P, Y, Phi, w):
    """The atoms' noise level (score units: |v| sqrt(w)) from the top height degrees of the long parts."""
    vals = []
    for b in range(len(Y)):
        J = int(np.count_nonzero(np.abs(Phi[b]).sum(0) > 0))  # the block's degrees
        n = int(np.count_nonzero(np.abs(Phi[b]).sum(1) > 0))
        if n < MIN_LEN or J < TOP + 2:
            continue
        v = Phi[b][:, J - TOP:J].T @ Y[b]
        vals.append(np.abs(v) * np.sqrt(w[b]))
    return float(np.median(np.concatenate(vals)) / 0.6745) if vals else 0.0


def main():
    with (REPORT / "noise_floor.csv").open("w", newline="") as f:
        wr = csv.writer(f)
        wr.writerow(["object", "input", "config", "atoms", "level", "chamfer", "p95", "n_mean", "n_over30"])
        for name in OBJ:
            g = trimesh.load(OBJ[name], force="mesh")
            ref = Ref(g)
            for inp in INPUTS:
                if inp == "mesh":
                    geom, kw = g, {}
                else:
                    geom = np.asarray(trimesh.load(CLOUDS / f"{name}_{inp}.ply").vertices)
                    kw = {"noise": 0.006 * float(np.linalg.norm(g.extents))} if inp.startswith("pn") else {}
                cap = {}

                def spy(*a, **k):
                    cap["sys"] = _sys(*a, **k)
                    return cap["sys"]

                _parts.parts_system = spy
                try:
                    d = fcp.encode(geom, atoms=6144, up=UP, device=DEVICE, **kw)
                finally:
                    _parts.parts_system = _sys
                P, Y, Phi, w = (np.asarray(x, float) if i else x for i, x in enumerate(cap["sys"]))
                level = noise_level(P, Y, Phi, w)
                # the score of every stored atom, in order
                i, v = d.atoms
                blk = i // (d.M + 1)
                score = np.abs(v.astype(float)) * np.sqrt(np.asarray(w, float)[blk])
                rows = [("package-1024", 1024), ("package-6144", 6144)]
                for c in CS:
                    rows.append((f"floor-{c:g}", int(np.count_nonzero(score > c * level)) if level > 0 else len(i)))
                for cfg, T in rows:
                    m = fcp.decode(d.prefix(T), device=DEVICE)
                    wr.writerow([name, inp, cfg, T, level, *[round(float(x), 5) for x in ref.measure(m)]])
                f.flush()
            print(name, flush=True)


if __name__ == "__main__":
    main()