"""Longer parts with a higher height degree, meshes only: the cut rule (tracked-fit / tracked) against the degree cap
(M, and KAPPA, the largest condition number of a part's height basis on its layer centres). Each encoding at 6144 atoms,
its prefixes at 256 / 1024 / 6144, decoded by the SDF default at 128. Writes ../report/degree_sweep.csv.

The KAPPA patch changes the degree the encoder and the decoder derive from a part's layer count (`_parts.height_degree`),
in this process only: a descriptor written with KAPPA != 1e2 does not decode right in the package.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2810_partition/scripts/degree_sweep.py
    FCP_OBJECTS=bowl_gso,dino ... (a subset)
"""
import csv
import os
import sys
import time
from pathlib import Path

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

sys.path.insert(0, os.environ.get("FCP_API", "/home/joao/code/fcp_devel/api"))  # FCP_API: another checkout (a worktree)
import fcp  # noqa: E402
from fcp import _parts  # noqa: E402

ROOT = Path(__file__).resolve().parent.parent
REPORT = ROOT / "report"
D = Path("/home/joao/code/data/meshes")
OBJ = {  # the ten watertight meshes of fcp_2710_tests (scripts/parts_report.py)
    "battery_stepped": D / "fcp_2710_tests/meshes/battery_stepped.ply",
    "dress": Path("/home/joao/code/fousd-3d/outputs/heavy/MANNEQUIN_SUMMER_DRESS_-_ORANGE_envelope_solid_target.ply"),
    "elephant": D / "data/meshes_clean/gso__Elephant.ply",
    "bowl_of": D / "data/meshes_clean/objectfolder__2.ply",
    "bowl_gso": D / "data/meshes_clean/gso__Threshold_Bead_Cereal_Bowl_White.ply",
    "mug": D / "data/meshes_clean/gso__Room_Essentials_Mug_White_Yellow.ply",
    "dino": D / "data/meshes_clean/gso__Schleich_Therizinosaurus_ln9cruulPqc.ply",
    "shoe": D / "data/meshes_clean/gso__Timberland_Mens_Classic_2Eye_Boat_Shoe.ply",
    "open_box": D / "fcp_2710_tests/meshes/open_box.ply",
    "horse": D / "data/meshes_clean/gso__Breyer_Horse_Of_The_Year_2015.ply",
}
if os.environ.get("FCP_OBJECTS"):
    OBJ = {k: v for k, v in OBJ.items() if k in os.environ["FCP_OBJECTS"].split(",")}
UP = (0, 0, 1)  # every object stands upright on world +z (as in fcp_2710_tests)
KAPPAS = [1e2, 1e3, 1e4]
MS = [24, 35, 50, 80]
RULES = ["tracked-fit", "tracked"]
BUDGETS = [256, 1024, 6144]
DEVICE = os.environ.get("FCP_DEVICE", "cuda")
OUT = REPORT / os.environ.get("FCP_OUT", "degree_sweep.csv")
FIELDS = ["object", "rule", "kappa", "M", "atoms", "parts", "max_deg", "mean_len", "chamfer", "p95", "n_mean", "n_over30",
          "seconds"]


def set_kappa(k):
    _parts.KAPPA = k
    _parts.height_degree.cache_clear()
    _parts.orth_basis.cache_clear()


class Ref:
    """The clean mesh as 200k surface samples with their face normals; Chamfer over 50k of them (in % of the diagonal)."""

    def __init__(self, g):
        self.pts, fi = trimesh.sample.sample_surface(g, 200_000, seed=1)
        self.nrm = g.face_normals[fi]
        self.tree = cKDTree(self.pts)
        self.chk = self.pts[:50_000]
        self.diag = float(np.linalg.norm(np.ptp(self.pts, axis=0)))

    def measure(self, m):
        if m is None or len(m.faces) == 0:
            return [np.nan] * 4
        q, fi = trimesh.sample.sample_surface(m, 50_000, seed=2)
        a, j = self.tree.query(q)
        e = np.concatenate([cKDTree(q).query(self.chk)[0], a]) / self.diag * 100
        c = np.abs(np.einsum("ij,ij->i", m.face_normals[fi], self.nrm[j]))  # unsigned: a flipped normal is a separate failure
        ang = np.degrees(np.arccos(np.clip(c, -1, 1)))
        return [e.mean(), np.percentile(e, 95), ang.mean(), (ang > 30).mean() * 100]


def main():
    REPORT.mkdir(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 kappa in KAPPAS:
                set_kappa(kappa)
                for M in MS:
                    for rule in RULES:
                        t = time.perf_counter()
                        d = fcp.encode(g, atoms=max(BUDGETS), up=UP, M=M, parts=rule, device=DEVICE)
                        sec = time.perf_counter() - t
                        lens = [h1 - h0 + 1 for _, h0, h1 in d.parts]
                        deg = max(_parts.height_degree(n, M) for n in lens)
                        for T in BUDGETS:
                            m = fcp.decode(d.prefix(T), device=DEVICE)
                            w.writerow([name, rule, f"{kappa:g}", M, T, len(d.parts), deg, round(float(np.mean(lens)), 1),
                                        *[round(float(v), 5) for v in ref.measure(m)], round(sec, 3)])
                        f.flush()
                        print(f"{name:16s} k={kappa:g} M={M:3d} {rule:12s} parts={len(d.parts):3d} deg={deg:3d}", flush=True)
    set_kappa(1e2)


if __name__ == "__main__":
    main()