"""The gallery behind the synthesis: for every object, cloud and configuration, the parts (render) and the final rings
(slices) that went into tracking, against GT and the points. Object by object (every configuration on
every cloud type, then the next object): ../report/gallery/<cloud>/<name>.html, images in ../report/img/gallery/, the
per-cloud ../report/gallery/<cloud>/index.html and ../report/gallery/index.html, updated as each object is done.

Configurations (ablation.py's): package, pointtest, A, hybrid-gated, hybrid-gated+pointtest (CONFIGS below).
Slices: the rings the encoder tracked (captured at parts_system), at the layers LAYERS and the layer where the package
has the most rings; GT is the exact section of the mesh the cloud came from (clean or noisy).

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2710_tests/scripts/gallery.py clean_10000 noise0p6_10000 ...
"""
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

sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, "/home/joao/code/fcp_devel/api")
import splits as sp  # noqa: E402  (its patch of parts_system captures the rings every encode tracks)
import fcp  # noqa: E402
import hybrid as hy  # noqa: E402
import pointfit  # noqa: E402
import ringfit as rf  # noqa: E402
import thinwall as tw  # noqa: E402
import eps as epsmod  # noqa: E402
import surf_local  # noqa: E402
sys.path.insert(0, "/home/joao/code/data/meshes/fcp_2810_gwn/scripts")  # the GWN slicer (another agent's, used as is)
import gwn  # noqa: E402
from make_clouds import PN  # noqa: E402
from make_noisy import label  # noqa: E402
from fcp._slice import centres  # noqa: E402
from parts_report import CSS, OBJ, REPORT, ROOT, UP, render, stats  # noqa: E402

ENC = {}  # extra encode arguments for every configuration and the mesh route (--layers)
CONFIGS = {
    "package": lambda P, up: fcp.encode(P, device="cuda", up=up, **ENC),
    "pointtest": lambda P, up: pointfit.encode_pointtest(P, up=up, **ENC),
    "A": lambda P, up: rf.encode_ringfit(P, up=up, **ENC),
    "hybrid-gated": lambda P, up: hy.encode_hybrid(P, up=up, gate=True, **ENC),
    "hybrid-gated+pointtest": lambda P, up: hy.encode_hybrid(P, up=up, cuts="pointtest", gate=True, **ENC),
    "thinwall+pointtest": lambda P, up: tw.encode_thinwall(P, up=up, cuts="pointtest", **ENC),
    "eps": lambda P, up: epsmod.encode_eps(P, up=up, sigma=CTX["sigma"], **ENC),
    "local": lambda P, up: surf_local.encode_local(P, up=up, sigma=CTX["sigma"], **ENC),
    "gwn": lambda P, up: gwn.encode_gwn(P, up=up, sigma=CTX["sigma"], **ENC),
    "gwn+pointtest": lambda P, up: gwn.encode_gwn(P, up=up, sigma=CTX["sigma"], cuts="pointtest", **ENC),
}
CTX = {"sigma": 0.0}  # the cloud's per-point noise (world units), given to "eps"
LAYERS = (16, 40, 64, 88, 112)
GAL = REPORT / "gallery"
IMG = REPORT / "img" / "gallery"


def closed(R):
    return np.vstack([R, R[:1]])


def tracked_rings(S, sc):
    """Per layer, (metric rings, hole flags) of the Slices S the encoder tracked."""
    out = {}
    for r in range(len(S.h)):
        out.setdefault(int(S.h[r]), []).append((np.asarray(S.rings.ring(r), float) * sc[:2], bool(S.hole[r])))
    return out


def one(name, cloud):
    tag, n = cloud.rsplit("_", 1)
    g = trimesh.load(OBJ[name], force="mesh")
    src = trimesh.load(ROOT / "meshes" / f"{name}_{tag}.ply", force="mesh") if tag.startswith("noise") else g
    P = np.asarray(trimesh.load(ROOT / "pointclouds" / f"{name}_{cloud}.ply").vertices)
    up = UP.get(name, (0, 0, 1))
    ref = np.asarray(g.sample(50_000, seed=1))
    CTX["sigma"] = ({f"pn{label(v)}": v for v in PN}.get(tag, 0.0)) * float(np.linalg.norm(g.extents))
    t = time.perf_counter()
    mesh_d = fcp.encode(src, device="cuda", up=up, **ENC)
    msec = time.perf_counter() - t
    mesh_S = sp.CAP["S"]
    mm = fcp.decode(mesh_d, device="cuda")
    mmean, mp95 = stats(mesh_d, ref, mm) if mm is not None else (float("nan"), float("nan"))
    mimg = f"{cloud}_{name}_mesh.png"
    render(mesh_d, mm, IMG / mimg, f"{name} · mesh route ({'clean' if tag == 'clean' else tag} mesh): {len(mesh_d.parts)} parts "
           f"({sum(s < 0 for s, _, _ in mesh_d.parts)} holes), mean/diag {mmean:.4f}")
    res = {"mesh": dict(d=mesh_d, mean=mmean, p95=mp95, sec=msec, img=mimg)}
    for cfg, f in CONFIGS.items():
        t = time.perf_counter()
        d = f(P, up)
        sec = time.perf_counter() - t
        S = sp.CAP["S"]
        m = fcp.decode(d, device="cuda")
        mean, p95 = stats(d, ref, m) if m is not None else (float("nan"), float("nan"))
        img = f"{cloud}_{name}_{cfg}.png"
        render(d, m, IMG / img, f"{name} · {cloud.replace('_', ' ')} · {cfg}: {len(d.parts)} parts "
               f"({sum(s < 0 for s, _, _ in d.parts)} holes), mesh {len(mesh_d.parts)}, mean/diag {mean:.4f}")
        sc = d.pose.scale.astype(float)
        res[cfg] = dict(d=d, mean=mean, p95=p95, sec=sec, img=img, rings=tracked_rings(S, sc), sc=sc)
    # slices: the configured layers and the package's most-ringed layer
    pk = res["package"]["rings"]
    worst = max(pk, key=lambda k: len(pk[k])) if pk else len(centres(res["package"]["d"].layers)) // 2
    layers = sorted(set(LAYERS) | {worst})
    d0 = res["package"]["d"]
    sc = res["package"]["sc"]
    X = (P - d0.pose.center) @ d0.pose.rotation
    s = float(d0.provenance["spacing"])
    c = centres(d0.layers)
    # the mesh route's rings at the mesh layer nearest in height to each cloud layer, moved into the cloud's frame
    cm = centres(mesh_d.layers)
    zw = lambda pose, z: (pose.to_world(np.array([[0.0, 0.0, z]]))[0] @ np.asarray(up, float))  # noqa: E731
    mesh_rings = {}
    for k in layers:
        km = int(np.argmin([abs(zw(mesh_d.pose, z) - zw(d0.pose, c[k])) for z in cm]))
        rs = []
        for r in np.nonzero(mesh_S.h == km)[0]:
            Rc = np.asarray(mesh_S.rings.ring(r), float)
            W = mesh_d.pose.to_world(np.c_[Rc, np.full(len(Rc), cm[km])])
            rs.append((((W - d0.pose.center) @ d0.pose.rotation)[:, :2], bool(mesh_S.hole[r])))
        mesh_rings[k] = (km, rs)
    cols = ["mesh"] + list(CONFIGS)
    fig, axs = plt.subplots(len(layers), len(cols), figsize=(4.2 * len(cols), 4.2 * len(layers)), dpi=75, squeeze=False)
    for row, k in zip(axs, layers):
        gt = [L * sc[:2] for L in sp.gt_section(src, d0.pose, c[k])]
        near = np.abs(X[:, 2] - c[k] * sc[2]) <= s
        allxy = np.vstack(gt) if gt else (X[near, :2] if near.any() else np.zeros((1, 2)))
        lo, hi = allxy.min(0), allxy.max(0)
        mg = 0.08 * (hi - lo).max() + 2 * s
        for ax, cfg in zip(row, cols):
            ax.scatter(X[near, 0], X[near, 1], s=2, c="#a8a8a8", lw=0, zorder=1)
            for L in gt:
                ax.plot(L[:, 0], L[:, 1], c="k", lw=1.5, zorder=2)
            rings = mesh_rings[k][1] if cfg == "mesh" else res[cfg]["rings"].get(k, [])
            for R, h in rings:
                ax.plot(*closed(R).T, c="#2e86ab" if h else "#d1495b", lw=1.1, ls="--" if h else "-", zorder=3)
            ax.set_xlim(lo[0] - mg, hi[0] + mg); ax.set_ylim(lo[1] - mg, hi[1] + mg); ax.set_aspect("equal")
            ax.tick_params(labelsize=5)
            lab = f"mesh route · its layer {mesh_rings[k][0]}" if cfg == "mesh" else f"{cfg} · layer {k}"
            ax.set_title(f"{lab}{' (package: most rings)' if k == worst else ''}: "
                         f"{sum(not h for _, h in rings)} solid, {sum(h for _, h in rings)} holes", fontsize=7)
    fig.tight_layout()
    simg = f"{cloud}_{name}_slices.png"
    fig.savefig(IMG / simg)
    plt.close(fig)
    rows = "".join(f"<tr><td>{cfg}</td><td>{len(r['d'].parts)}</td><td>{sum(x < 0 for x, _, _ in r['d'].parts)}</td>"
                   f"<td>{r['mean']:.5f}</td><td>{r['p95']:.5f}</td><td>{r['sec']:.1f}</td></tr>" for cfg, r in res.items())
    imgs = "".join(f"<h3>{cfg}</h3><img src='../../img/{IMG.name}/{r['img']}' loading='lazy'>" for cfg, r in res.items())
    page = f"""<!doctype html><html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1"><title>{name}, {cloud}</title><style>{CSS}</style></head><body>
<p><a href="../index.html">← gallery</a> · <a href="index.html">{cloud.replace('_', ' ')}</a> · <a href="../../synthesis.html">synthesis</a></p>
<h1>{html.escape(name)} · {cloud.replace('_', ' ')}</h1>
<p class="mut">"mesh": the mesh route on the mesh the cloud came from (clean or noisy), the reference for parts and rings. Errors against
the clean mesh (mean and p95 of the symmetric nearest-point distance over the diagonal).</p>
<table><tr><th>configuration</th><th>parts</th><th>holes</th><th>mean/diag</th><th>p95/diag</th><th>encode s</th></tr>{rows}</table>
<h2>Slices: the rings each configuration tracked</h2>
<p class="mut">First column: the mesh route's rings (at the mesh layer nearest in height). Black: GT (exact section of the source mesh at
the cloud's layer). Grey: the points within a spacing of the plane. Red: solid rings, blue dashed: holes. Metric canonical frame of the
cloud (x right, y up).</p>
<img src="../../img/{IMG.name}/{simg}" loading="lazy">
<h2>Parts</h2>{imgs}
</body></html>"""
    (GAL / cloud / f"{name}.html").write_text(page)
    return {cfg: (len(r["d"].parts), r["mean"]) for cfg, r in res.items() if cfg != "mesh"}, len(mesh_d.parts)


def write_index(cloud, done):
    rows = "".join(f"<tr><td><a href='{name}.html'>{name}</a></td><td>{mp}</td>"
                   + "".join(f"<td>{r[c][0]} · {r[c][1]:.4f}</td>" for c in CONFIGS) + "</tr>" for name, (r, mp) in done.items())
    (GAL / cloud / "index.html").write_text(f"""<!doctype html><html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1"><title>Gallery: {cloud}</title><style>{CSS}</style></head><body>
<p><a href="../index.html">← gallery</a> · <a href="../../synthesis.html">synthesis</a></p><h1>Gallery: {cloud.replace('_', ' ')}</h1>
<p class="mut">parts · mean/diag per configuration; mesh: the mesh route's parts. Each object's page has its slices and part renders.
Objects appear as they are done.</p>
<table><tr><th>object</th><th>mesh</th>{''.join(f'<th>{c}</th>' for c in CONFIGS)}</tr>{rows}</table></body></html>""")


def write_top(clouds, names):
    """The gallery's front page: every cloud type built so far (not only this run's), every object done in it."""
    clouds = sorted(p.name for p in GAL.iterdir() if (p / "index.html").exists())
    names = [n for n in OBJ if any((GAL / c / f"{n}.html").exists() for c in clouds)]
    cell = lambda c, n: f"<td><a href='{c}/{n}.html'>{c.replace('_', ' ')}</a></td>" if (GAL / c / f"{n}.html").exists() else "<td></td>"  # noqa: E731
    links = "".join(f"<tr><td>{name}</td>" + "".join(cell(c, name) for c in clouds) + "</tr>" for name in names)
    (GAL / "index.html").write_text(f"""<!doctype html><html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1"><title>Gallery</title><style>{CSS}</style></head><body>
<p><a href="../synthesis.html">← synthesis</a> · <a href="../index.html">report</a></p><h1>Gallery: every method on the same object</h1>
<p class="mut">Each page: one object on one cloud type with all configurations ({', '.join(CONFIGS)}): a table, the slices (all
methods side by side, layer by layer) and the part renders. Objects appear as they are done. By cloud type:
{' · '.join(f"<a href='{c}/index.html'>{c.replace('_', ' ')}</a>" for c in clouds)}.</p>
<table><tr><th>object</th>{''.join(f'<th>{c.replace("_", " ")}</th>' for c in clouds)}</tr>{links}</table></body></html>""")


def main(clouds):
    """Object by object: every configuration on every cloud type, then the next object (pages written as they are done)."""
    IMG.mkdir(parents=True, exist_ok=True)
    for cloud in clouds:
        (GAL / cloud).mkdir(parents=True, exist_ok=True)
    done = {c: {} for c in clouds}
    for name in OBJ:
        for cloud in clouds:
            done[cloud][name] = one(name, cloud)
            write_index(cloud, done[cloud])
            print(name, cloud, flush=True)
        write_top(clouds, [n for n in OBJ if n in done[clouds[-1]]])


FINAL = ("package", "pointtest", "local", "gwn", "gwn+pointtest")  # --final: the recipes that matter, to report/gallery_final


if __name__ == "__main__":
    if "--final" in sys.argv:
        sys.argv.remove("--final")
        CONFIGS = {k: CONFIGS[k] for k in FINAL}
        GAL = REPORT / "gallery_final"
        IMG = REPORT / "img" / "gallery_final"
    if "--layers" in sys.argv:  # e.g. --layers 256: every encode at that many layers, slices at the same heights
        i = sys.argv.index("--layers")
        NL = int(sys.argv[i + 1])
        del sys.argv[i:i + 2]
        ENC["layers"] = NL
        LAYERS = tuple(int(round(k * NL / 128)) for k in LAYERS)
        GAL = REPORT / f"gallery_{NL}"
        IMG = REPORT / "img" / f"gallery_{NL}"
    if "--objects" in sys.argv:  # e.g. --objects loaf_pan,mug
        i = sys.argv.index("--objects")
        ONLY = sys.argv[i + 1].split(",")
        del sys.argv[i:i + 2]
        OBJ = {k: v for k, v in OBJ.items() if k in ONLY}
    main(sys.argv[1:] or ["clean_10000", "clean_2048", "noise0p6_10000", "noise1_10000", "pn0p6_10000", "pn1p2_10000",
                          "dgrad_10000", "dstripe_10000"])