"""Renders of the energy prefixes (energy_prefix.py at 6144 atoms): per cloud one image, a row per object, columns mesh route
(6144 atoms) / gwn + closing full 6144 / 99.9 % / 99 % of the atom energy / the first 1024 atoms. The decoded mesh at 96 cells per
axis for the render (every face drawn, light from the upper left), coloured by part; titles: atoms, parts, error (x 1e-3, from the
decode at the descriptor's layers) and the SDF-normal median. -> ../report/img/energy_<cloud>.png

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python gallery_energy.py clean_10000 [objects]
"""
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import numpy as np
import trimesh
from scipy.spatial import cKDTree

sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, "/home/joao/code/data/meshes/fcp_2710_tests/scripts")
import gwn  # noqa: E402
import fcp  # noqa: E402
import fcp.tools  # noqa: E402
from parts_report import OBJ, ROOT, colors, part_rings, stats  # noqa: E402
from bench import sigma_of  # noqa: E402

IMG = Path(__file__).parent.parent / "report" / "img"
ATOMS = 6144
RES = 96  # render decode, cells per axis
COLS = (("full", 1.0), ("99.9 %", 0.999), ("99 %", 0.99), ("first 1024", -1024))


def prefix_atoms(d, fr):
    v = np.asarray(d.atoms[1], float)
    if fr < 0:
        return min(int(-fr), len(v))
    if fr >= 1:
        return len(v)
    return int(np.searchsorted(np.cumsum(v * v) / (v * v).sum(), fr) + 1)


def draw(ax, d, mesh, title):
    ax.set_title(title, fontsize=7.5)
    ax.set_axis_off()
    if mesh is None or not len(mesh.faces):
        return
    rings = part_rings(d)
    col = colors(len(d.parts))
    allp = np.concatenate([r.reshape(-1, 3) for r in rings if r is not None])
    own = np.concatenate([np.full(r.size // 3, i) for i, r in enumerate(rings) if r is not None])
    V = d.pose.to_canonical(mesh.vertices)
    F = mesh.faces
    part = own[cKDTree(allp).query(V[F].mean(1))[1]]
    n = np.cross(V[F[:, 1]] - V[F[:, 0]], V[F[:, 2]] - V[F[:, 0]])
    n /= np.maximum(np.linalg.norm(n, axis=1, keepdims=True), 1e-12)
    light = np.array([-0.5, -0.6, 0.65]); light /= np.linalg.norm(light)
    shade = 0.3 + 0.7 * np.clip(n @ light, 0, 1)
    ax.add_collection3d(Poly3DCollection(V[F], facecolors=np.clip(col[part] * shade[:, None], 0, 1), edgecolors="none",
                                         linewidths=0))
    ax.view_init(22, -58)
    lo, hi = V.min(0), V.max(0)
    m, r = (lo + hi) / 2, (hi - lo).max() / 2
    ax.set_xlim(m[0] - r, m[0] + r); ax.set_ylim(m[1] - r, m[1] + r); ax.set_zlim(m[2] - r, m[2] + r)
    ax.set_box_aspect((1, 1, 1))


def scores(d, ref, Ps, Ns):
    m = fcp.decode(d, device="cuda")
    if m is None:
        return float("nan"), float("nan")
    Ne = fcp.tools.Shape(d, device="cuda").normal(Ps, method="sdf")
    Ne /= np.maximum(np.linalg.norm(Ne, axis=1, keepdims=True), 1e-12)
    return 1e3 * stats(d, ref, m)[0], float(np.median(np.degrees(np.arccos(np.clip((Ne * Ns).sum(1), -1, 1)))))


def main(clouds, objs=None):
    IMG.mkdir(parents=True, exist_ok=True)
    objs = objs or list(OBJ)
    gwn.CLOSE = 0.75
    for cl in clouds:
        tag, n = cl.rsplit("_", 1)
        ncol = 1 + len(COLS)
        fig = plt.figure(figsize=(ncol * 2.7, len(objs) * 2.7), dpi=100)
        for r, name in enumerate(objs):
            g = trimesh.load(OBJ[name], force="mesh")
            ref = np.asarray(g.sample(50_000, seed=1))
            Ps, fi = trimesh.sample.sample_surface(g, 3000, seed=3)
            Ns = g.face_normals[fi]
            diag = float(np.linalg.norm(g.extents))
            P = np.asarray(trimesh.load(ROOT / "pointclouds" / f"{name}_{tag}_{n}.ply").vertices)
            dm = fcp.encode(g, device="cuda", up=(0, 0, 1), atoms=ATOMS)
            d = gwn.encode_gwn(P, up=(0, 0, 1), sigma=sigma_of(tag, diag), atoms=ATOMS)
            panels = [("mesh route", dm)] + [(lab, d.prefix(prefix_atoms(d, fr))) for lab, fr in COLS]
            for k, (lab, dp) in enumerate(panels):
                e, nm = scores(dp, ref, Ps, Ns)
                ax = fig.add_subplot(len(objs), ncol, r * ncol + k + 1, projection="3d", proj_type="ortho")
                draw(ax, dp, fcp.decode(dp, device="cuda", resolution=RES),
                     f"{name} · {lab}\n{dp.n_atoms} atoms · {len(dp.parts)} parts\nerr {e:.2f} · normal {nm:.1f}°")
            print(cl, name, flush=True)
        fig.suptitle(f"{cl}: gwn + closing at {ATOMS} atoms and its energy prefixes (mesh route: the clean mesh)", fontsize=11)
        fig.tight_layout()
        out = IMG / f"energy_{cl}.png"
        fig.savefig(out)
        plt.close(fig)
        print("saved", out, flush=True)


if __name__ == "__main__":
    main(sys.argv[1].split(","), sys.argv[2].split(",") if len(sys.argv) > 2 else None)