"""Decoded meshes, coloured by part: one image per cloud, a row per object, columns mesh route / package / eps / local / gwn.
../report/img/gallery_<cloud>.png. Titles: parts and mean error (x 1e-3 of the diagonal).

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python gallery_gwn.py clean_2048 pn1p2_10000 ...
"""
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 eps as epsmod  # noqa: E402
import surf_local  # 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"
CONFIGS = {
    "package": lambda P, sig: fcp.encode(P, device="cuda", up=(0, 0, 1)),
    "eps": lambda P, sig: epsmod.encode_eps(P, up=(0, 0, 1), sigma=sig),
    "local": lambda P, sig: surf_local.encode_local(P, up=(0, 0, 1), sigma=sig),
    "gwn": lambda P, sig: gwn.encode_gwn(P, up=(0, 0, 1), sigma=sig),
}


def draw(ax, d, mesh, title):
    rings = part_rings(d)
    col = colors(len(d.parts))
    ax.set_title(title, fontsize=8)
    ax.set_axis_off()
    if mesh is None:
        return
    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
    if len(F) > 60000:
        F = F[np.random.default_rng(0).choice(len(F), 60000, replace=False)]
    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.45, -0.6, 0.65]); light /= np.linalg.norm(light)
    shade = 0.35 + 0.65 * np.abs(n @ light)
    ax.add_collection3d(Poly3DCollection(V[F], facecolors=col[part] * shade[:, None], edgecolors="none", linewidths=0))
    ax.view_init(20, -60)
    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 main(clouds, objs=None):
    IMG.mkdir(parents=True, exist_ok=True)
    objs = objs or list(OBJ)
    for cl in clouds:
        tag, n = cl.rsplit("_", 1)
        fig = plt.figure(figsize=(5 * 2.6, len(objs) * 2.6), dpi=100)
        for r, name in enumerate(objs):
            g = trimesh.load(OBJ[name], force="mesh")
            src = trimesh.load(ROOT / "meshes" / f"{name}_{tag}.ply", force="mesh") if tag.startswith("noise") else g
            ref = np.asarray(g.sample(50_000, seed=1))
            diag = float(np.linalg.norm(g.extents))
            P = np.asarray(trimesh.load(ROOT / "pointclouds" / f"{name}_{tag}_{n}.ply").vertices)
            runs = [("mesh route", lambda: fcp.encode(src, device="cuda", up=(0, 0, 1)))]
            runs += [(k, (lambda f=f: f(P, sigma_of(tag, diag)))) for k, f in CONFIGS.items()]
            for k, (lab, f) in enumerate(runs):
                d = f()
                m = fcp.decode(d, device="cuda")
                mean = stats(d, ref, m)[0] if m is not None else float("nan")
                ax = fig.add_subplot(len(objs), len(runs), r * len(runs) + k + 1, projection="3d", proj_type="ortho")
                draw(ax, d, m, f"{name} · {lab}\n{len(d.parts)} parts · err {1e3 * mean:.2f}")
            print(cl, name, flush=True)
        fig.suptitle(f"{cl}: decoded meshes coloured by part", fontsize=11)
        fig.tight_layout()
        out = IMG / f"gallery_{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)