"""Parts of the mesh and point-cloud encodings, rendered per part, into ../report/index.html.

Run from the fcp checkout:
    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2710_tests/scripts/parts_report.py
"""
import colorsys
import os
import html
import sys
import time
from pathlib import Path

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

sys.path.insert(0, "/home/joao/code/fcp_devel/api")
import fcp
from fcp.decode import _axis_rings, grid

sys.path.insert(0, str(Path(__file__).parent))
import pointfit  # noqa: E402  the bench prototype: cuts decided by a two-sided point test
import neighbour  # noqa: E402  the bench prototype: slices refined from their neighbours (two sweeps, merged)
import ringfit  # noqa: E402  the bench prototypes A (slicer topology, fitted contours) and B (point sheets)
import hybrid  # noqa: E402  the bench prototype: slicer topology repaired from the neighbouring layers, all fitted

ROOT = Path(__file__).resolve().parent.parent
REPORT = ROOT / "report"
IMG = REPORT / "img"
D = Path("/home/joao/code/data")
M = D / "meshes/data"
OBJ = {  # all watertight meshes (the scans without one were replaced by similar shapes, 2026-10-27)
    "battery_stepped": ROOT / "meshes/battery_stepped.ply",  # synthetic, scripts/make_battery.py (for the DeWalt scan)
    "dress": Path("/home/joao/code/fousd-3d/outputs/heavy/MANNEQUIN_SUMMER_DRESS_-_ORANGE_envelope_solid_target.ply"),
    "elephant": M / "meshes_clean/gso__Elephant.ply",  # for the Stanford bunny (open scan)
    "bowl_of": M / "meshes_clean/objectfolder__2.ply",
    "bowl_gso": M / "meshes_clean/gso__Threshold_Bead_Cereal_Bowl_White.ply",
    "mug": M / "meshes_clean/gso__Room_Essentials_Mug_White_Yellow.ply",
    "dino": M / "meshes_clean/gso__Schleich_Therizinosaurus_ln9cruulPqc.ply",
    "shoe": M / "meshes_clean/gso__Timberland_Mens_Classic_2Eye_Boat_Shoe.ply",
    "open_box": ROOT / "meshes/open_box.ply",  # synthetic, scripts/make_openbox.py (replaces the ObjectFolder loaf pan,
    # objectfolder__38.ply, whose mesh is full of slivers: its noisy versions reached only a fraction of the noise)
    "horse": M / "meshes_clean/gso__Breyer_Horse_Of_The_Year_2015.ply",
}
if os.environ.get("FCP_OBJECTS"):  # restrict every script to these objects (comma-separated)
    OBJ = {k: v for k, v in OBJ.items() if k in os.environ["FCP_OBJECTS"].split(",")}
N_POINTS = 200_000
UP = {}  # world axis an object stands upright on, when not +z
DEVICE = "cuda"


def colors(n):
    """n distinct colours: golden-ratio hues, alternating lightness."""
    return np.array([colorsys.hls_to_rgb((i * 0.618034) % 1, 0.45 + 0.15 * (i % 2), 0.75) for i in range(n)])


def part_rings(d):
    """Per part: canonical ring points (R, n_t, 3) at the decode grid heights (None for a part without rings)."""
    c = grid(d.layers)
    C = d.coef()
    B = C.shape[0] // max(len(d.parts), 1)
    out = []
    for i, p in enumerate(d.parts):
        r = _axis_rings((p,), C[i * B:(i + 1) * B], c, d.layers, np, interp=d.interp)
        if not len(r.h):
            out.append(None)
            continue
        XY = np.asarray(r.XY)  # (R, 2, n_t)
        z = np.broadcast_to(c[r.h][:, None], XY[:, 0].shape)
        out.append(np.stack([XY[:, 0], XY[:, 1], z], -1))
    return out


def _axes(ax, P):
    lo, hi = P.min(0), P.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)); ax.set_axis_off()


def render(d, mesh, path, title):
    """Canonical frame, orthographic, height axis z up: every ring of every part (holes dashed) seen along +y; the decoded
    mesh coloured by its nearest part, seen along +y (x right) and along -x (y right); the parts' layer spans."""
    rings = part_rings(d)
    col = colors(len(d.parts))
    fig = plt.figure(figsize=(18, 5), dpi=110)
    fig.suptitle(title, fontsize=11)
    allp = np.concatenate([r.reshape(-1, 3) for r in rings if r is not None])
    for k, (elev, azim) in enumerate(((0, -90), (0, -90), (0, 0))):
        ax = fig.add_subplot(1, 4, k + 1, projection="3d", proj_type="ortho")
        if k == 0:
            segs, cs, ls = [], [], []
            for i, r in enumerate(rings):
                if r is None:
                    continue
                for ring in r[::max(1, len(r) // 40)] if len(r) > 40 else r:
                    closed = np.vstack([ring, ring[:1]])
                    segs.append(closed); cs.append(col[i]); ls.append("--" if d.parts[i][0] < 0 else "-")
            ax.add_collection3d(Line3DCollection(segs, colors=cs, linewidths=0.7, linestyles=ls))
            ax.set_title("rings by part (x-z)", fontsize=9)
        else:
            if mesh is not None:
                V = d.pose.to_canonical(mesh.vertices)
                own = np.full(len(allp), -1)
                o = 0
                for i, r in enumerate(rings):
                    if r is not None:
                        own[o:o + r.size // 3] = i; o += r.size // 3
                fc = V[mesh.faces].mean(1)
                part = own[cKDTree(allp).query(fc)[1]]
                n = np.cross(V[mesh.faces[:, 1]] - V[mesh.faces[:, 0]], V[mesh.faces[:, 2]] - V[mesh.faces[:, 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[mesh.faces], facecolors=col[part] * shade[:, None],
                                                     edgecolors="none", linewidths=0))
            ax.set_title("decoded mesh by nearest part (" + ("y-z" if k == 2 else "x-z") + ")", fontsize=9)
        ax.view_init(elev, azim)
        _axes(ax, allp)
    ax = fig.add_subplot(1, 4, 4)
    for i, (s, h0, h1) in enumerate(d.parts):
        ax.barh(i, h1 - h0 + 1, left=h0, color=col[i], hatch="//" if s < 0 else None, edgecolor="k", linewidth=0.3)
    ax.set_xlim(0, d.layers); ax.set_xlabel("layer"); ax.set_ylabel("part"); ax.invert_yaxis()
    ax.set_title(f"{len(d.parts)} parts ({sum(s < 0 for s, _, _ in d.parts)} holes, hatched): layer spans", fontsize=9)
    fig.tight_layout()
    fig.savefig(path)
    plt.close(fig)


def load(p):
    g = trimesh.load(p, force="mesh", process=True)
    return g, bool(len(g.faces)) and g.is_watertight


def stats(d, ref_pts, mesh):
    q = np.asarray(mesh.sample(50_000, seed=2)) if mesh is not None else None
    diag = np.linalg.norm(np.ptp(ref_pts, axis=0))
    if q is None:
        return np.nan, np.nan
    e = np.concatenate([cKDTree(q).query(ref_pts)[0], cKDTree(ref_pts).query(q)[0]]) / diag
    return e.mean(), np.percentile(e, 95)


def main(cloud=None):
    """cloud: None (200k points sampled here, ../report/index.html) or a cloud of ../pointclouds, e.g. "clean_10000"
    (../pointclouds/<name>_clean_10000.ply, ../report/parts_clean_10000.html)."""
    IMG.mkdir(parents=True, exist_ok=True)
    rows = []
    for name, p in OBJ.items():
        g, watertight = load(p)
        if cloud is None:
            pts, what, sfx = np.asarray(g.sample(N_POINTS, seed=0)), f"points {N_POINTS // 1000}k", ""
        else:
            pts = np.asarray(trimesh.load(ROOT / "pointclouds" / f"{name}_{cloud}.ply").vertices)
            what, sfx = f"points {cloud.replace('_', ' ')}", f"_{cloud}"
        ref = np.asarray(g.sample(50_000, seed=1))
        runs = (([("mesh", "mesh", g)] if watertight else []) + [(what, f"points{sfx}", pts),
                (f"{what}, point test", f"pointtest{sfx}", pts)])
        if cloud is not None:
            runs += [(f"{what}, neighbour slices", f"neighbour{sfx}", pts), (f"{what}, A ringfit", f"ringfitA{sfx}", pts),
                     (f"{what}, B graph", f"graphB{sfx}", pts), (f"{what}, hybrid", f"hybrid{sfx}", pts)]
        for label, key, geom in runs:
            t = time.perf_counter()
            up = UP.get(name, (0, 0, 1))
            if key.startswith("pointtest"):
                d = pointfit.encode_pointtest(geom, up=up, device=DEVICE)
            elif key.startswith("neighbour"):
                d = neighbour.encode_neighbour(geom, up=up, device=DEVICE)
            elif key.startswith("ringfitA"):
                d = ringfit.encode_ringfit(geom, up=up, device=DEVICE)
            elif key.startswith("graphB"):
                d = ringfit.encode_graph(geom, up=up, device=DEVICE)
            elif key.startswith("hybrid"):
                d = hybrid.encode_hybrid(geom, up=up, device=DEVICE)
            else:
                d = fcp.encode(geom, device=DEVICE, up=up)
            te = time.perf_counter() - t
            m = fcp.decode(d, device=DEVICE)
            mean, p95 = stats(d, ref, m)
            tag = f"{name}_{key}"
            render(d, m, IMG / f"{tag}.png", f"{name} — {label}: {len(d.parts)} parts, {d.nbytes} B")
            ns = sum(s < 0 for s, _, _ in d.parts)
            spans = np.array([h1 - h0 + 1 for _, h0, h1 in d.parts])
            rows.append(dict(name=name, input=label, up=UP.get(name, (0, 0, 1)), watertight=watertight, parts=len(d.parts), holes=ns,
                             short=int((spans <= 3).sum()), median_span=float(np.median(spans)), nbytes=d.nbytes,
                             enc=te, mean=mean, p95=p95, spacing=d.provenance.get("spacing"),
                             fit_tol=d.provenance.get("fit_tol"), img=f"img/{tag}.png"))
            print(f"{name:16s} {label:12s} parts {len(d.parts):4d} holes {ns:3d} short {rows[-1]['short']:3d} "
                  f"mean/diag {mean:.5f}", flush=True)
    write_html(rows, cloud)


CSS = """
:root{--bg:#fafaf8;--fg:#1d1d1b;--mut:#6b6b66;--line:#ddd9d0;--acc:#2f6f8f;--bad:#b4442a}
body{background:var(--bg);color:var(--fg);font:15px/1.5 system-ui,sans-serif;margin:0 auto;max-width:1180px;padding:24px 16px}
h1{font-size:26px;margin:0 0 4px}h2{font-size:19px;margin:36px 0 8px;border-bottom:1px solid var(--line);padding-bottom:4px}
.mut{color:var(--mut)}table{border-collapse:collapse;width:100%;font-variant-numeric:tabular-nums;font-size:13.5px}
th,td{padding:5px 8px;border-bottom:1px solid var(--line);text-align:right}th{background:#efece5}td:first-child,th:first-child,td:nth-child(2),th:nth-child(2){text-align:left}
td.bad{color:var(--bad);font-weight:600}img{max-width:100%;border:1px solid var(--line);background:#fff;margin:6px 0}
nav a{margin-right:12px;color:var(--acc)}code{background:#efece5;padding:1px 4px;border-radius:3px}
"""


def write_html(rows, cloud=None):
    def f(x, fmt):
        return "" if x is None or (isinstance(x, float) and np.isnan(x)) else format(x, fmt)

    mesh_parts = {r["name"]: r["parts"] for r in rows if r["input"] == "mesh"}
    tr = []
    for r in rows:
        ref = mesh_parts.get(r["name"])
        bad = ref is not None and r["input"] != "mesh" and r["parts"] > 1.5 * ref
        tr.append(f"<tr><td><a href='#{r['name']}'>{r['name']}</a></td><td>{html.escape(r['input'])}</td>"
                  f"<td>{'+' + 'xyz'[int(np.argmax(r['up']))]}</td><td>{'yes' if r['watertight'] else 'no'}</td><td class='{'bad' if bad else ''}'>{r['parts']}</td>"
                  f"<td>{r['holes']}</td><td>{r['short']}</td><td>{r['median_span']:.0f}</td><td>{r['nbytes']}</td>"
                  f"<td>{f(r['spacing'], '.4g')}</td><td>{f(r['fit_tol'], '.4g')}</td><td>{r['enc']:.2f}</td>"
                  f"<td>{f(r['mean'], '.5f')}</td><td>{f(r['p95'], '.5f')}</td></tr>")
    sec, last = [], None
    for r in rows:
        if r["name"] != last:
            sec.append(f"<h2 id='{r['name']}'>{r['name']}</h2>")
            last = r["name"]
        sec.append(f"<p class='mut'>{html.escape(r['input'])}: {r['parts']} parts, {r['holes']} holes, "
                   f"{r['short']} parts of ≤ 3 layers</p><img src='{r['img']}' loading='lazy'>")
    if cloud is None:
        title, source = "FCP parts: mesh vs point cloud", (
            f"Point clouds: {N_POINTS:,} points sampled uniformly on the source mesh (seed 0, no noise).")
        studies = sorted([*REPORT.glob("synthesis.html"), *REPORT.glob("gallery_final/index.html"), *REPORT.glob("gallery/index.html"), *REPORT.glob("analysis_*.html"), *REPORT.glob("middle_*.html"), *REPORT.glob("compare_*.html"), *REPORT.glob("neighbour_*.html"), *REPORT.glob("parts_*.html"), *REPORT.glob("noisy.html"), *REPORT.glob("slices_*.html"), *REPORT.glob("splits_*.html")])
        nav = "<p>Studies: " + " ".join(f"<a href='{q.relative_to(REPORT)}'>{q.stem.replace('_', ' ')}</a>" for q in studies) + "</p>"
    else:
        title, source = f"FCP parts: mesh vs {cloud.replace('_', ' ')} points", (
            f"Point clouds: <code>pointclouds/&lt;name&gt;_{cloud}.ply</code> (<code>scripts/make_clouds.py</code>: area-uniform, "
            "seeded; noise, if any, is the mesh's, <a href='noisy.html'>noisy meshes</a>).")
        nav = "<p><a href='index.html'>← report</a></p>"
    page = f"""<!doctype html><html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1"><title>{title}</title>
<style>{CSS}</style></head><body>
{nav if cloud else ""}
<h1>{title}</h1>
<p class="mut">fcp v5.2 ({DEVICE}), defaults: 1024 atoms, 128 layers, <code>parts="tracked-fit"</code>, <code>fit_tol="auto"</code>; sliced upright: <code>up</code> = the world axis each object stands on (+z for all).
{source} Errors: symmetric nearest-point distance
between 50k samples of the source and of the decoded mesh, over the source's bounding-box diagonal.
Generated {time.strftime('%Y-%m-%d %H:%M')} by <code>scripts/parts_report.py{' ' + cloud if cloud else ''}</code>.</p>
{"" if cloud else nav}
<h2>1. Parts per encoding</h2>
<p class="mut">"point test" (prototype, <code>scripts/pointfit.py</code>, package untouched): a piece of a track stays whole when, at every layer,
one height fit's ring is within {pointfit.K_SPACING:g} spacings (p95) of the slicer's ring in explaining the layer's points, and the fitted ring lies on
the points within the same margin (both directions). Where to cut and the minimum piece length are the package's.
"neighbour slices" (prototype, <code>scripts/neighbour.py</code>, point clouds of <code>pointclouds/</code> only): the slicer's rings refined from the
neighbouring slice in two independent sweeps (up, down), merged per layer, then the package's tracking and cuts.
"A ringfit" / "B graph" (prototypes, <code>scripts/ringfit.py</code>): contours fitted to the points with the slicer's topology (A) or the
points' own sheets (B); "hybrid" (<code>scripts/hybrid.py</code>): the slicer's topology by default, repaired from the neighbouring
layers, every contour fitted; <a href="analysis_ringfit.html">analysis</a>.</p>
<p class="mut">Red: a point-cloud encoding with more than 1.5× the parts of the same object's mesh encoding. "short": parts spanning ≤ 3 layers.
spacing and fit_tol are in canonical units, from the descriptor's provenance.</p>
<table><tr><th>object</th><th>input</th><th>up</th><th>watertight</th><th>parts</th><th>holes</th><th>short</th><th>median span</th>
<th>bytes</th><th>spacing</th><th>fit_tol</th><th>encode s</th><th>mean/diag</th><th>p95/diag</th></tr>
{''.join(tr)}</table>
<h2>2. Renders</h2>
<p class="mut">Canonical frame, orthographic, height axis z up (layer 0 at the bottom). Left: rings of every part seen along +y (x right; one colour per part, holes dashed; tall parts thinned to ~40 rings).
Then the decoded mesh along +y (x-z) and along -x (y-z), each face coloured by the part whose rings are nearest (matplotlib: expect some depth-sorting artefacts). Right: each part's span of layers (holes hatched).</p>
{''.join(sec)}
</body></html>"""
    (REPORT / ("index.html" if cloud is None else f"parts_{cloud}.html")).write_text(page)


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