"""One part's height polynomial against its sliced rings: the object's largest part (most layers x area), tracked-fit with
a given cap. Top: the equivalent radius sqrt(area / pi) of the sliced rings (dots, at the layer centres) and of the rings
the stored polynomial gives at 1024 and 6144 atoms, drawn continuously between the centres (the polynomial itself, 20
samples per layer), and the ground truth: the mesh sectioned at 10 heights per layer (the part's ring: centroid nearest
the sliced ring at the nearest centre). Middle: single coefficients the same way (centroid x and y, the cos / sin of harmonic 1 in x, a
higher harmonic; ground truth for the centroids, which do not depend on the ring's start). Bottom: the shape miss against
the sliced rings at the centres and against the ground truth at every height (layer heights). Writes
../report/img/profile_<object>_k<kappa>_M<M>.png.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2810_partition/scripts/part_profile.py dino 1e4 50 [N]
"""
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 trimesh  # noqa: E402
from numpy.polynomial import chebyshev as npcheb  # noqa: E402

sys.path.insert(0, str(Path(__file__).parent))
from degree_sweep import DEVICE, OBJ, REPORT, UP, set_kappa  # noqa: E402  (puts the package on the path)
import fcp  # noqa: E402
from fcp import _parts  # noqa: E402
from fcp._slice import centres  # noqa: E402
from fit_audit import area  # noqa: E402
from shape_test import rings, shape_miss  # noqa: E402

SUB = 20  # polynomial samples per layer
GT_SUB = 10  # mesh sections per layer (ground truth)
NS = 512  # arc-length samples of a ground-truth ring


def gt_rings(g, d, zs, ref):
    """Ground truth: the mesh sectioned in the canonical frame at heights zs; at each, the closed ring of the part (centroid
    nearest ref(z) = (centroid, area) of the part's sliced ring at the nearest layer centre, area within a factor 3).
    Returns per height (area, centroid (2,), arc-length samples (NS, 2)) or None."""
    cm = trimesh.Trimesh(d.pose.to_canonical(g.vertices), g.faces, process=False)
    out = []
    for z in zs:
        sec = cm.section(plane_origin=[0, 0, z], plane_normal=[0, 0, 1])
        best = None
        for L in (sec.discrete if sec is not None else []):
            Q = np.asarray(L)[:, :2]
            if len(Q) < 4:
                continue
            if np.allclose(Q[0], Q[-1]):
                Q = Q[:-1]
            seg = np.linalg.norm(np.diff(np.vstack([Q, Q[:1]]), axis=0), axis=1)
            s = np.r_[0, np.cumsum(seg)]
            t = np.linspace(0, s[-1], NS, endpoint=False)
            Qc = np.vstack([Q, Q[:1]])
            R = np.c_[np.interp(t, s, Qc[:, 0]), np.interp(t, s, Qc[:, 1])]
            a = abs(0.5 * (Q[:, 0] * np.roll(Q[:, 1], -1) - np.roll(Q[:, 0], -1) * Q[:, 1]).sum())
            c = R.mean(0)
            rc, ra = ref(z)
            if not ra / 3 < a < ra * 3:
                continue
            dist = np.linalg.norm(c - rc)
            if best is None or dist < best[0]:
                best = (dist, a, c, R)
        out.append(None if best is None else best[1:])
    return out


def main(name, kappa, M, N=16):
    set_kappa(kappa)
    cap, sysf = {}, _parts.parts_system

    def capture(*a, **k):
        cap["out"] = sysf(*a, **k)
        return cap["out"]

    _parts.parts_system = capture
    g = trimesh.load(OBJ[name], force="mesh")
    d = fcp.encode(g, atoms=6144, up=UP, N=N, M=M, parts="tracked-fit", device=DEVICE)
    _parts.parts_system = sysf
    P, Y, _, _ = cap["out"]
    K, B = d.K, 2 * d.K
    lh = float(centres(d.layers)[1] - centres(d.layers)[0])
    sl = {}
    for i, p in enumerate(P):
        n = p["h1"] - p["h0"] + 1
        sl[i] = Y[i * B:(i + 1) * B, :n].T.astype(float).reshape(n, 2, K)
    i = max((j for j in sl if P[j]["sign"] > 0), key=lambda j: np.abs(area(rings(sl[j]))).sum())
    p, S = P[i], sl[i]
    n = p["h1"] - p["h0"] + 1
    deg = _parts.height_degree(n, M)
    h = np.arange(p["h0"], p["h1"] + 1)
    uc = _parts.cell_u(n)
    uf = np.linspace(-1, 1, n * SUB + 1)  # the part spans its cells' outer edges
    hf = p["h0"] - 0.5 + (uf + 1) / 2 * n
    poly = {}
    for T in (1024, 6144):
        C = d.prefix(T).coef()[i * B:(i + 1) * B]  # (2K, M + 1) Chebyshev in the part's local height
        kept = int(np.count_nonzero(np.abs(C).sum(0) > 0))
        poly[T] = (npcheb.chebval(uf, C.T).T.reshape(-1, 2, K), npcheb.chebval(uc, C.T).T.reshape(-1, 2, K), kept,
                   int(np.count_nonzero(C)))
    inside = (hf >= p["h0"]) & (hf <= p["h1"])  # between the first and the last layer centre (what a decode reads)
    zc = centres(d.layers)
    hg = np.arange(p["h0"] - 0.5 + 0.5 / GT_SUB, p["h1"] + 0.5, 1 / GT_SUB)  # ground-truth heights, in layers
    zg = zc[0] + hg * lh
    RS = rings(S)
    ref = lambda z: (S[int(np.clip(round((z - zc[0]) / lh) - p["h0"], 0, n - 1)), :, 0],
                     abs(area(RS[[int(np.clip(round((z - zc[0]) / lh) - p["h0"], 0, n - 1))]])[0]))
    gt = gt_rings(g, d, zg, ref)
    ok = np.array([x is not None for x in gt])
    hg, gt = hg[ok], [x for x in gt if x is not None]
    g_r = np.sqrt(np.array([x[0] for x in gt]) / np.pi) / lh
    g_c = np.array([x[1] for x in gt]) / lh
    ug = 2 * (hg - (p["h0"] - 0.5)) / n - 1  # the ground-truth heights in the part's local height
    g_miss = {}
    for T in poly:
        C = d.prefix(T).coef()[i * B:(i + 1) * B]
        Pr = rings(npcheb.chebval(ug, C.T).T.reshape(-1, 2, K))
        g_miss[T] = np.array([shape_miss(x[2][None], Pr[[j]])[0] for j, x in enumerate(gt)]) / lh

    def zoom(ax, data, fits):
        lo, hi = float(np.min(data)), float(np.max(data))
        pad = 0.15 * (hi - lo) + 0.5
        ax.set_ylim(lo - pad, hi + pad)
        out = max(float(np.abs(f[~inside]).max()) for f in fits)
        ax.text(0.01, 0.97, f"beyond the end centres up to |{out:.0f}|", transform=ax.transAxes, fontsize=7, va="top", color="0.4")

    fig = plt.figure(figsize=(14, 13), dpi=100)
    gs = fig.add_gridspec(3, 5, height_ratios=[1.2, 1, 1])
    ax = fig.add_subplot(gs[0, :])
    R0 = np.sqrt(np.abs(area(rings(S))) / np.pi) / lh
    ax.plot(hg, g_r, "k-", lw=2.2, alpha=0.35, label=f"ground truth: mesh sections, {GT_SUB} per layer")
    ax.plot(h, R0, "ko", ms=3.5, label="sliced rings (layer centres)")
    fits = []
    for T, st in ((1024, "C2-"), (6144, "C1-")):
        Rf = np.sqrt(np.abs(area(rings(poly[T][0]))) / np.pi) / lh
        fits.append(Rf)
        ax.plot(hf, Rf, st, lw=1, label=f"polynomial, {T} atoms ({poly[T][3]} values in this part, degrees up to {poly[T][2] - 1})")
    zoom(ax, R0, fits)
    ax.set_ylabel("equivalent radius (layer heights)")
    ax.legend(fontsize=8, loc="lower left")
    ax.set_title(f"{name}: part {i} (solid), layers {p['h0']}-{p['h1']} ({n} layers), height degree {deg} "
                 f"(KAPPA {kappa:g}, N {N}, M {M}); {len(P)} parts", fontsize=10)
    picks = [((0, 0), "x centroid"), ((1, 0), "y centroid"), ((0, 1), "x cos 1"), ((1, 2), "y sin 1"), ((0, 7), "x cos 4")]  # the harmonics depend on the ring's start: their reference is the sliced rings only
    for c, ((xy, k), lab) in enumerate(picks):
        ax = fig.add_subplot(gs[1, c])
        if k == 0:
            ax.plot(hg, g_c[:, xy], "k-", lw=2.2, alpha=0.35, label="ground truth")
        ax.plot(h, S[:, xy, k] / lh, "ko", ms=2.5)
        ax.plot(hf, poly[6144][0][:, xy, k] / lh, "C1-", lw=1, label="6144")
        ax.plot(hf, poly[1024][0][:, xy, k] / lh, "C2-", lw=0.9, label="1024")
        zoom(ax, S[:, xy, k] / lh, [poly[T][0][:, xy, k] / lh for T in poly])
        ax.set_title(lab, fontsize=9)
        ax.set_xlim(p["h0"] - 0.6, p["h1"] + 0.6); ax.grid(alpha=0.3)
        if c == 0:
            ax.set_ylabel("coefficient (layer heights)"); ax.legend(fontsize=7)
    ax2 = fig.add_subplot(gs[2, :])
    for T, st, col in ((1024, "o", "C2"), (6144, "o", "C1")):
        ax2.plot(hg, g_miss[T], "-", color=col, lw=0.8, alpha=0.7, label=f"{T} atoms against the ground truth (between the layers too)")
        ax2.plot(h, shape_miss(RS, rings(poly[T][1])) / lh, st, color=col, ms=3.5, label=f"{T} atoms against the sliced rings (centres)")
    ax2.axhline(_parts.FIT_TOL / lh, color="r", ls="--", lw=0.7, label="cut tolerance (0.01)")
    ax2.set_xlabel("layer")
    ax2.set_ylabel("shape miss (layer heights)")
    ax2.set_ylim(0, 2.0)
    ax2.legend(fontsize=8)
    for a in (fig.axes[0], ax2):
        a.set_xlim(p["h0"] - 0.6, p["h1"] + 0.6)
        a.grid(alpha=0.3)
    fig.tight_layout()
    out = REPORT / "img" / (f"profile_{name}_k{kappa:g}_M{M}.png" if N == 16 else f"profile_{name}_k{kappa:g}_N{N}_M{M}.png")
    fig.savefig(out)
    print(out, f"part {i} layers {p['h0']}-{p['h1']} degree {deg}")
    set_kappa(1e2)


if __name__ == "__main__":
    main(sys.argv[1], float(sys.argv[2]), int(sys.argv[3]), *(int(a) for a in sys.argv[4:5]))