"""The slices on either side of every fit cut of a point-cloud encoding: GT (exact mesh section), the points within
half a layer of the plane and the point slicer's rings, in the canonical frame. Writes ../report/splits_<name>.html.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2710_tests/scripts/splits.py dino
"""
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, "/home/joao/code/fcp_devel/api")
sys.path.insert(0, str(Path(__file__).parent))
import fcp
from fcp import _parts
from fcp._slice import centres
from parts_report import CSS, N_POINTS, OBJ, REPORT, UP

CAP = {}
_orig_system, _orig_split = _parts.parts_system, _parts.split_by_fit


def system(geoms, *a, **k):
    CAP["S"] = geoms
    CAP["log"] = []
    return _orig_system(geoms, *a, **k)


def split(parts, coeffs, areas, M, tol=None):
    for p, C, A in zip(parts, coeffs, areas):  # the pieces of every track, cut alone
        CAP["log"].append((p["sign"], p["h0"], p["h1"], [(q["h0"], q["h1"]) for q in _orig_split([p], [C], [A], M, tol)[0]]))
    return _orig_split(parts, coeffs, areas, M, tol)


_parts.parts_system, _parts.split_by_fit = system, split


def encode(geom, up):
    d = fcp.encode(geom, device="cuda", up=up)
    return d, CAP["S"], CAP["log"]


def cuts(log):
    """(sign, track h0, h1, cut layer c: the first layer of the piece above, piece below, piece above)."""
    out = []
    for s, h0, h1, pieces in log:
        for lo, hi in zip(pieces[:-1], pieces[1:]):
            out.append((s, h0, h1, hi[0], lo, hi))
    return out


def track_rings(S, sign, h0, h1):
    """The ring ids of the track (sign, h0, h1), by height."""
    for tr in _parts._track(S, S.per_height(S.hole if sign < 0 else ~S.hole)):
        hs = sorted(tr)
        if hs[0] == h0 and hs[-1] == h1:
            return tr
    return {}


def gt_section(mesh, pose, z):
    o = pose.to_world(np.array([[0.0, 0.0, z]]))[0]
    sec = mesh.section(plane_origin=o, plane_normal=pose.rotation[:, 2])
    return [] if sec is None else [pose.to_canonical(np.asarray(L))[:, :2] for L in sec.discrete]


def main(name):
    REPORT.mkdir(exist_ok=True)
    (REPORT / "img").mkdir(exist_ok=True)
    g = trimesh.load(OBJ[name], force="mesh")
    up = UP.get(name, (0, 0, 1))
    P = np.asarray(g.sample(N_POINTS, seed=0))
    _, _, mlog = encode(g, up)
    mesh_cuts = {c for *_, c, _, _ in cuts(mlog)}
    d, S, log = encode(P, up)
    c = centres(d.layers)
    half = (c[1] - c[0]) / 2
    Pc = d.pose.to_canonical(P)
    rows = []
    for k, (s, h0, h1, cut, lo, hi) in enumerate(cuts(log)):
        tr = track_rings(S, s, h0, h1)
        own = {tr[h]: (0 if h < cut else 1) for h in (cut - 1, cut) if h in tr}
        box = np.concatenate([S.rings.ring(r) for r in own]) if own else np.zeros((1, 2))
        lo_xy, hi_xy = box.min(0), box.max(0)
        m = 0.25 * (hi_xy - lo_xy).max() + 0.02
        lo_xy, hi_xy = lo_xy - m, hi_xy + m
        fig, axs = plt.subplots(1, 2, figsize=(10, 5), dpi=100)
        for ax, h in zip(axs, (cut - 1, cut)):
            near = np.abs(Pc[:, 2] - c[h]) <= half
            ax.scatter(Pc[near, 0], Pc[near, 1], s=1.5, c="#9a9a9a", lw=0, label="points (±½ layer)", zorder=1)
            for i, L in enumerate(gt_section(g, d.pose, c[h])):
                ax.plot(L[:, 0], L[:, 1], c="k", lw=1.6, label="GT section" if i == 0 else None, zorder=2)
            first = True
            for r in np.nonzero(S.h == h)[0]:
                R = np.vstack([S.rings.ring(r), S.rings.ring(r)[:1]])
                if r in own:
                    col = ("#d1495b", "#2e86ab")[own[r]]
                    lab = f"slice: part {lo[0]}-{lo[1]}" if own[r] == 0 else f"slice: part {hi[0]}-{hi[1]}"
                else:
                    col, lab = "#e0a100", "slice: other rings" if first else None
                    first = False
                ax.plot(R[:, 0], R[:, 1], c=col, lw=1.1, ls="--" if S.hole[r] else "-", label=lab, zorder=3)
            ax.set_xlim(lo_xy[0], hi_xy[0]); ax.set_ylim(lo_xy[1], hi_xy[1]); ax.set_aspect("equal")
            ax.set_title(f"layer {h} ({'below' if h < cut else 'above'} the cut)", fontsize=10)
            ax.tick_params(labelsize=7)
        h_, l_ = [], []
        for ax in axs:
            for hh, ll in zip(*ax.get_legend_handles_labels()):
                if ll not in l_:
                    h_.append(hh); l_.append(ll)
        fig.legend(h_, l_, loc="lower center", ncol=5, fontsize=8, frameon=False)
        fig.tight_layout(rect=(0, 0.06, 1, 1))
        f = f"img/splits_{name}_{k:02d}.png"
        fig.savefig(REPORT / f)
        plt.close(fig)
        both = any(abs(cut - x) <= 1 for x in mesh_cuts)
        rows.append((k, h0, h1, cut, lo, hi, both, f))
        print(f"cut {k:2d}: track {h0}-{h1}, {lo[0]}-{lo[1]} | {hi[0]}-{hi[1]}  mesh cuts here: {both}", flush=True)
    sec = "".join(
        f"<h2 id='c{k}'>Cut {k}: layers {cut - 1} | {cut} <span class='mut'>(track {h0}-{h1}; parts {lo[0]}-{lo[1]} | "
        f"{hi[0]}-{hi[1]}; the mesh {'also cuts here (±1 layer)' if both else 'does not cut here'})</span></h2>"
        f"<img src='{f}' loading='lazy'>" for k, h0, h1, cut, lo, hi, both, f in rows)
    idx = " ".join(f"<a href='#c{k}'>{cut - 1}|{cut}{'' if both else '*'}</a>" for k, _, _, cut, _, _, both, _ in rows)
    page = f"""<!doctype html><html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1"><title>{name}: slices at the cuts</title>
<style>{CSS}</style></head><body>
<p><a href="index.html">← report</a></p>
<h1>{html.escape(name)}: slices at the fit cuts</h1>
<p class="mut">{N_POINTS:,} points (seed 0), up {up}, fcp defaults; {len(d.parts)} parts, {len(rows)} fit cuts
({sum(not r[6] for r in rows)} where the mesh encoding does not cut, marked *). Each cut: the layer below and the layer above,
canonical frame (x right, y up), zoomed on the cut track's rings. Black: GT, the exact section of the mesh at the layer centre.
Grey: the points within half a layer of the plane. Coloured: the point slicer's rings (red: the part below the cut, blue: the
part above, yellow: other rings; dashed: holes). Generated {time.strftime('%Y-%m-%d %H:%M')} by <code>scripts/splits.py</code>.</p>
<nav>{idx}</nav>
{sec}
</body></html>"""
    (REPORT / f"splits_{name}.html").write_text(page)


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