"""Slices of one cloud: points within half a layer (grey), mesh rings (black), package (blue), GWN variants (orange).
args: object tag n out.png [KEY=VAL;KEY=VAL ...] (one GWN panel per variant); layers: the 4 with most GWN rings."""
import sys
from pathlib import Path
import numpy as np, trimesh
import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
sys.path.insert(0, str(Path(__file__).parent)); sys.path.insert(0, "/home/joao/code/data/meshes/fcp_2710_tests/scripts")
import gwn
from fcp._slice import PointSlicer, MeshSlicer, centres
from parts_report import OBJ, ROOT

name, tag, n, out = sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4]
variants = sys.argv[5:] or ["ISO=local"]
g = trimesh.load(OBJ[name], force="mesh")
P = np.asarray(trimesh.load(ROOT / "pointclouds" / f"{name}_{tag}_{n}.ply").vertices)
c = centres(128)
ps = PointSlicer(P, up=(0, 0, 1), device="cuda", slicer="surf")
Lp = ps.loops(c)
res = []
for v in variants:
    saved = {}
    kw = {}
    for kv in v.split(";"):
        k, val = kv.split("=")
        if k == "c":
            kw["c_eps"] = float(val); continue
        if k == "sig":  # sigma as a fraction of the mesh's diagonal
            kw["sigma"] = float(val) * float(np.linalg.norm(g.extents)); continue
        saved[k] = getattr(gwn, k)
        setattr(gwn, k, val if isinstance(saved[k], str) else type(saved[k])(val))
    q = gwn.slicer_gwn(P, up=(0, 0, 1), **kw)
    res.append((v, q.loops(c)))
    for k, val in saved.items():
        setattr(gwn, k, val)
Xc = ps.pose.to_canonical(P)
mesh_c = trimesh.Trimesh(ps.pose.to_canonical(g.vertices), g.faces)
cnt = np.array([len(x) for x in res[0][1]])
import os
Ls = [int(v) for v in os.environ["LAYERS"].split(",")] if os.environ.get("LAYERS") else sorted(np.argsort(-cnt)[:4])
cols = 2 + len(res)
fig, ax = plt.subplots(len(Ls), cols, figsize=(3.2 * cols, 3.2 * len(Ls)))
half = (c[1] - c[0]) / 2
for r, L in enumerate(Ls):
    sec = mesh_c.section(plane_origin=[0, 0, c[L]], plane_normal=[0, 0, 1])
    near = Xc[np.abs(Xc[:, 2] - c[L]) < half]
    panels = [("mesh", None), ("package", Lp[L])] + [(v, R[L]) for v, R in res]
    for k, (lab, loops) in enumerate(panels):
        a = ax[r, k]
        a.scatter(near[:, 0], near[:, 1], s=2, c="0.6")
        if sec is not None:
            for e in sec.discrete:
                a.plot(e[:, 0], e[:, 1], "k-", lw=0.6)
        if loops is not None:
            for Q in loops:
                Q = np.vstack([Q, Q[:1]]); a.plot(Q[:, 0], Q[:, 1], "-", lw=0.9, color="tab:blue" if k == 1 else "tab:orange")
        a.set_aspect("equal"); a.set_xticks([]); a.set_yticks([])
        a.set_title(f"L{L} {lab} ({len(loops) if loops is not None else len(sec.discrete) if sec is not None else 0})", fontsize=7)
plt.tight_layout(); plt.savefig(out, dpi=110)
print("saved", out, "layers", Ls)