"""The surf slicer's field at chosen layers: the signed field (RIMLS within BAND spacings of the points, the flood beyond),
its zero line (the rings), the points within a spacing of the plane and GT. Writes ../report/field_<name>_<cloud>.html.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2710_tests/scripts/field_view.py bowl_gso clean_10000 114,116,117,118
"""
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
from fcp import _surf  # noqa: E402
from fcp._rings import host  # noqa: E402
from fcp._slice import PointSlicer, centres  # noqa: E402
from parts_report import CSS, OBJ, REPORT, ROOT, UP  # noqa: E402

CAP = {}


def capture(B):
    orig = B.field

    def field(total, nL, off_d, HW_d, geo, pix, *rest):
        f = orig(total, nL, off_d, HW_d, geo, pix, *rest)
        CAP.update(fld=host(f), off=host(off_d), HW=host(HW_d), geo=host(geo), pix=float(pix))
        return f
    B.field = field
    return orig


def main(name, cloud, layers, device="cuda"):
    g = trimesh.load(OBJ[name], force="mesh")
    P = np.asarray(trimesh.load(ROOT / "pointclouds" / f"{name}_{cloud}.ply").vertices)
    up = UP.get(name, (0, 0, 1))
    sl = PointSlicer(P, up=up, device=device)
    B = sl._surf.B
    orig = capture(B)
    try:
        c = centres(128)
        loops = sl.loops(c)
    finally:
        B.field = orig
    sc = sl.pose.scale.astype(float)
    s = float(sl.spacing)
    X = (P - sl.pose.center) @ sl.pose.rotation
    fig, axs = plt.subplots(len(layers), 2, figsize=(14, 6.4 * len(layers)), dpi=80, squeeze=False)
    for (axf, axz), k in zip(axs, layers):
        a, b = CAP["off"][k], CAP["off"][k + 1]
        Hh, Ww = CAP["HW"][k]
        if b - a == 0:
            continue
        F = CAP["fld"][a:b].reshape(Hh, Ww)
        lo = CAP["geo"][k, :2]
        pix = CAP["pix"]
        ext = (lo[0] - 0.5 * pix, lo[0] + (Hh - 0.5) * pix, lo[1] - 0.5 * pix, lo[1] + (Ww - 0.5) * pix)
        near = np.abs(X[:, 2] - c[k] * sc[2]) <= s
        gt = [L * sc[:2] for L in sp.gt_section(g, sl.pose, c[k])]
        rings = [np.asarray(L, float) * sc[:2] for L in loops[k]]
        lim = np.nanpercentile(np.abs(F), 95) or 1.0
        for ax in (axf, axz):
            ax.imshow(F.T, origin="lower", extent=ext, cmap="RdBu_r", vmin=-lim, vmax=lim, interpolation="nearest")
            ax.scatter(X[near, 0], X[near, 1], s=5, c="k", lw=0, zorder=3)
            for L in gt:
                ax.plot(L[:, 0], L[:, 1], c="#00a000", lw=1.2, zorder=4)
            for R in rings:
                R = np.vstack([R, R[:1]])
                ax.plot(R[:, 0], R[:, 1], c="#ffcc00", lw=1.4, zorder=5)
            ax.set_aspect("equal"); ax.tick_params(labelsize=7)
        axf.set_title(f"layer {k}: the field (red outside, blue inside), {len(rings)} rings (yellow), GT green", fontsize=9)
        # zoom: the first gap of a GT contour without a ring nearby, else the slice centre
        focus = None
        if gt and rings:
            from scipy.spatial import cKDTree
            T = cKDTree(np.vstack(rings))
            G = np.vstack(gt)
            d = T.query(G)[0]
            focus = G[np.argmax(d)]
        if focus is None:
            focus = np.vstack(gt).mean(0) if gt else lo
        w = 5 * s
        axz.set_xlim(focus[0] - w, focus[0] + w); axz.set_ylim(focus[1] - w, focus[1] + w)
        axz.set_title(f"zoom (±5 spacings) where GT is farthest from a ring; pixel {pix / s:.1f} s, spacing {s:.2f}", fontsize=9)
    fig.tight_layout()
    tag = f"field_{name}_{cloud}"
    fig.savefig(REPORT / f"img/{tag}.png")
    plt.close(fig)
    (REPORT / f"{tag}.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>{name}: surf field</title><style>{CSS}</style></head><body>
<p><a href="index.html">← report</a> · <a href="synthesis.html">synthesis</a></p>
<h1>{name} ({cloud.replace('_', ' ')}): the surf slicer's field</h1>
<p class="mut">The field the package's slicer contours (RIMLS within {_surf.BAND} spacings of the points, support {_surf.H_SUPPORT} spacings;
the flood beyond), red outside, blue inside; yellow: its zero line (the rings); green: GT; black: the points within a spacing of the
plane. Metric canonical frame. Generated {time.strftime('%Y-%m-%d %H:%M')} by <code>scripts/field_view.py</code>.</p>
<img src="img/{tag}.png"></body></html>""")
    print(tag)


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2], [int(x) for x in sys.argv[3].split(",")])