"""The slices where the descriptor's inside-outside answers are wrong, drawn: the exact mesh slice (black), the
descriptor's rings at that layer (blue), and the wrong points of a slab of ±SLAB layers around it, coloured by their
distance to the true surface (bands of 0.1–0.2, 0.2–0.3, 0.3–0.4 and > 0.4 % of the diagonal; filled: false inside,
open: false outside). Per object the ROWS slices with the most wrong points beyond 0.1 % diag, each with the two
densest clusters zoomed in. No decode.

    env -u PYTHONPATH .venv/bin/python scripts/slice_errors.py <object> <cache .fcp> <tag> [rows]
"""
import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import open3d as o3d  # noqa: E402
import trimesh  # noqa: E402
from matplotlib.lines import Line2D  # noqa: E402

sys.path.insert(0, str(Path(__file__).parent))
from common import CACHE, REPORT, RULE, cut, fcp, load_mesh  # noqa: E402
from fcpq import FCPQuery  # noqa: E402

LPAD = 1.1  # fcp._slice.L
SLAB = 2  # layers each side of the drawn slice
BANDS = [(0.1, 0.2), (0.2, 0.3), (0.3, 0.4), (0.4, np.inf)]
COLORS = ["#fc9272", "#ef3b2c", "#a50f15", "#3b0006"]  # one hue, light -> dark: farther is darker
EXACT, RING = "#111111", "#2b6cb0"


def clusters(xy, r, k=2):
    """Centres of the k densest groups of points (greedy: densest point, drop its r-neighbourhood, repeat)."""
    out, left = [], np.arange(len(xy))
    for _ in range(k):
        if len(left) == 0:
            break
        d2 = ((xy[left, None, :] - xy[None, left, :]) ** 2).sum(-1) if len(left) < 4000 else None
        if d2 is None:
            sub = left[np.random.default_rng(0).choice(len(left), 4000, replace=False)]
            cnt = (((xy[left, None, :] - xy[None, sub, :]) ** 2).sum(-1) < r * r).sum(1)
        else:
            cnt = (d2 < r * r).sum(1)
        c = xy[left[int(cnt.argmax())]]
        members = left[((xy[left] - c) ** 2).sum(1) < (2 * r) ** 2]
        out.append((xy[members].mean(0), len(members)))
        left = left[((xy[left] - c) ** 2).sum(1) >= (2 * r) ** 2]
    return out


def draw(ax, sec, rings, pts, ins, dist, diag, win=None):
    for c in sec.discrete if sec is not None else []:
        ax.plot(c[:, 0], c[:, 1], "-", color=EXACT, lw=1.6 if win is None else 2.2, zorder=2)
    for r in rings:
        ax.plot(*np.vstack([r, r[:1]]).T, "-", color=RING, lw=0.9 if win is None else 1.4, zorder=3)
    s = 9 if win is None else 34
    for (lo, hi), col in zip(BANDS, COLORS):
        b = (dist > lo) & (dist <= hi)
        ax.scatter(pts[b & ins, 0], pts[b & ins, 1], s=s, color=col, zorder=4, linewidths=0)
        ax.scatter(pts[b & ~ins, 0], pts[b & ~ins, 1], s=s, facecolors="white", edgecolors=col, linewidths=1.2, zorder=4)
    ax.set_aspect("equal")
    ax.tick_params(labelsize=7)
    if win is not None:
        (x0, x1), (y0, y1) = win
        ax.set_xlim(x0, x1)
        ax.set_ylim(y0, y1)
        bar = 0.001 * diag  # 0.1 % of the diagonal
        ax.plot([x0 + 0.06 * (x1 - x0), x0 + 0.06 * (x1 - x0) + bar], [y0 + 0.06 * (y1 - y0)] * 2, "-", color=EXACT, lw=2.5)
        ax.text(x0 + 0.06 * (x1 - x0), y0 + 0.09 * (y1 - y0), "0.1 % diag", fontsize=7)


def main():
    name, f, tag = sys.argv[1], sys.argv[2], sys.argv[3]
    rows = int(sys.argv[4]) if len(sys.argv) > 4 else 3
    m = load_mesh(name)
    d = cut(fcp.FCP.from_bytes((CACHE / f).read_bytes()))  # the study's cut
    sc = o3d.t.geometry.RaycastingScene()
    sc.add_triangles(o3d.core.Tensor(np.asarray(m.vertices, np.float32)), o3d.core.Tensor(np.asarray(m.faces, np.uint32)))
    np.random.seed(0)
    rng = np.random.default_rng(0)
    diag = np.linalg.norm(m.bounds[1] - m.bounds[0])
    S, fi = m.sample(1_000_000, return_index=True)
    P = S + m.face_normals[fi] * rng.uniform(-0.02, 0.02, (len(S), 1)) * diag
    T = o3d.core.Tensor(P.astype(np.float32))
    gt = sc.compute_occupancy(T, nsamples=3).numpy() > 0
    dist = np.abs(sc.compute_distance(T).numpy()) / diag * 100
    q = FCPQuery(d, nt=max(1024, 8 * d.N))
    Pc = d.pose.to_canonical(P).astype(np.float32)
    ins = q.inside_cpu(Pc, RULE)
    bad = (ins != gt) & (dist > BANDS[0][0])
    Pa = (P - d.pose.center) @ d.pose.rotation  # aligned frame, world units: x, y in the slice, z the height
    ma = trimesh.Trimesh((np.asarray(m.vertices) - d.pose.center) @ d.pose.rotation, m.faces, process=False)
    sz = d.pose.scale
    layer = (Pc[:, 2] + LPAD) / (2 * LPAD) * d.layers - 0.5
    L = np.rint(layer).astype(int)
    counts = np.bincount(L[bad], minlength=d.layers)
    # the busiest slices, at least 4 * SLAB layers apart
    picks = []
    for j in np.argsort(-np.convolve(counts, np.ones(2 * SLAB + 1), "same")):
        if all(abs(j - p) > 4 * SLAB for p in picks):
            picks.append(int(j))
        if len(picks) == rows:
            break
    fig, axs = plt.subplots(len(picks), 3, figsize=(16, 5.2 * len(picks)), gridspec_kw={"width_ratios": [1.4, 1, 1]})
    axs = np.atleast_2d(axs)
    for row, j in zip(axs, picks):
        zc = -LPAD + (j + 0.5) * 2 * LPAD / d.layers
        sec = ma.section(plane_origin=[0, 0, zc * sz[2]], plane_normal=[0, 0, 1])
        rings = [q.rings[q.first[p] + j - q.h0[p]] * sz[:2] for p in range(q.P) if q.h0[p] <= j <= q.h1[p]]
        sel = bad & (np.abs(layer - j) <= SLAB)
        pts, si, sd = Pa[sel, :2], ins[sel], dist[sel]
        n_far = {f"> {lo}": int((sd > lo).sum()) for lo, _ in BANDS}
        draw(row[0], sec, rings, pts, si, sd, diag)
        row[0].set_title(f"{name}, layer {j} ±{SLAB}: {len(pts)} wrong > 0.1 %\n"
                         f"({', '.join(f'{k}: {v}' for k, v in n_far.items() if k != '> 0.1')})", fontsize=9, loc="left")
        r = 0.006 * diag
        for ax, (c, n) in zip(row[1:], clusters(pts, r) if len(pts) else []):
            pad = 0.02 * diag
            draw(ax, sec, rings, pts, si, sd, diag, win=((c[0] - pad, c[0] + pad), (c[1] - pad, c[1] + pad)))
            ax.set_title(f"zoom: {n} points", fontsize=9, loc="left")
            row[0].add_patch(plt.Rectangle((c[0] - pad, c[1] - pad), 2 * pad, 2 * pad, fill=False, ec="#555555", lw=0.8, ls="--"))
        for ax in row[1 + len(clusters(pts, r) if len(pts) else []):]:
            ax.axis("off")
    handles = [Line2D([], [], color=EXACT, lw=2, label="mesh slice"), Line2D([], [], color=RING, lw=2, label="descriptor rings")]
    handles += [Line2D([], [], ls="", marker="o", color=c, label=f"{lo}–{hi} % diag" if np.isfinite(hi) else f"> {lo} % diag")
                for (lo, hi), c in zip(BANDS, COLORS)]
    handles += [Line2D([], [], ls="", marker="o", color="#777777", label="false inside"),
                Line2D([], [], ls="", marker="o", mfc="white", mec="#777777", label="false outside")]
    fig.legend(handles=handles, loc="upper center", ncol=8, fontsize=9, frameon=False)
    fig.tight_layout(rect=(0, 0, 1, 0.97))
    out = REPORT / "slices" / f"{tag}_{name}.png"
    out.parent.mkdir(exist_ok=True)
    fig.savefig(out, dpi=110)
    total = {f"{lo}": int((bad & (dist > lo)).sum()) for lo, _ in BANDS}
    print(f"{name} {tag}: wrong beyond {total} % diag; slices {picks} -> {out}")


if __name__ == "__main__":
    main()