"""What happens to a track where something joins it (the mug's handle): along the object's longest solid track, the step
between neighbouring rings in coefficients (the points at the same t, start points by continuity: what the package's fit
sees) and in shape (closest points), in layer heights; the package's cuts; and the rings themselves (canonical x-y)
around the largest steps, coloured by layer, with the slice's other rings (the handle) in grey.
Writes ../report/img/handle_<object>.png.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2810_partition/scripts/mug_handle.py [object]
"""
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

sys.path.insert(0, str(Path(__file__).parent))
from degree_sweep import DEVICE, OBJ, REPORT, UP  # noqa: E402  (puts the package on the path)
import fcp  # noqa: E402
from fcp import _parts  # noqa: E402
from dp_split import LH, miss_batch  # noqa: E402
from shape_test import rings  # noqa: E402

_split = _parts.split_by_fit


def main(name):
    g = trimesh.load(OBJ[name], force="mesh")
    cap = {}
    sysf = _parts.parts_system

    def system(S, *a, **k):
        cap["S"] = S
        return sysf(S, *a, **k)

    def split(parts, coeffs, areas, M, tol=None, follows=None):
        i = max(range(len(parts)), key=lambda j: (parts[j]["sign"] > 0, len(coeffs[j])))
        cap["p"], cap["C"], cap["A"] = parts[i], np.asarray(coeffs[i], float), np.asarray(areas[i], float)
        cap["cuts"] = [q["h0"] for q in _split([parts[i]], [coeffs[i]], [areas[i]], M, tol, follows)[0][1:]]
        cap["all"] = [(q, np.asarray(c, float)) for q, c in zip(parts, coeffs)]
        return _split(parts, coeffs, areas, M, tol, follows)

    _parts.parts_system, _parts.split_by_fit = system, split
    try:
        fcp.encode(g, atoms=1, up=UP, device=DEVICE)
    finally:
        _parts.parts_system, _parts.split_by_fit = sysf, _split
    p, C, A = cap["p"], cap["C"], cap["A"]
    K = C.shape[2]
    w = np.full(K, 0.5)
    w[0] = 1.0
    Cn = _parts.continuous_start(C, A)
    coef = np.sqrt((np.diff(Cn, axis=0) ** 2 * w).sum((1, 2))) / LH  # same t
    R = rings(C)
    shape = miss_batch(R[:-1], R[1:]) / LH
    x = np.arange(p["h0"], p["h1"]) + 0.5
    fig = plt.figure(figsize=(14, 9), dpi=100)
    gs = fig.add_gridspec(2, 4, height_ratios=[1, 1.25])
    ax = fig.add_subplot(gs[0, :])
    ax.plot(x, coef, "C0.-", lw=1, ms=4, label="step in coefficients (points at the same t)")
    ax.plot(x, shape, "C3.-", lw=1, ms=4, label="step in shape (closest points)")
    for c in cap["cuts"]:
        ax.axvline(c - 0.5, color="k", lw=0.8, ls="--", alpha=0.6)
    ax.set_yscale("log")
    ax.set_ylabel("ring-to-ring step (layer heights)")
    ax.set_title(f"{name}: solid track, layers {p['h0']}-{p['h1']}; dashed: the package's cuts {cap['cuts']}", fontsize=10)
    ax.legend(fontsize=8)
    ax.grid(alpha=0.3, which="both")
    big = np.argsort(-coef)[:4]  # the four largest coefficient steps
    for k, i in enumerate(sorted(big)):
        ax2 = fig.add_subplot(gs[1, k])
        hs = range(max(0, i - 2), min(len(R), i + 4))
        cols = plt.cm.viridis(np.linspace(0, 1, len(hs)))
        for col, j in zip(cols, hs):
            h = p["h0"] + j
            for q, cc in cap["all"]:  # the slice's other rings at this layer (the handle), grey
                if q is not p and q["h0"] <= h <= q["h1"] and q["sign"] > 0:
                    o = rings(cc[[h - q["h0"]]])[0]
                    ax2.plot(*np.vstack([o, o[:1]]).T, color="0.75", lw=0.8)
            r = np.vstack([R[j], R[j][:1]])
            ax2.plot(*r.T, color=col, lw=1.1, label=f"layer {h}")
            ax2.plot(*R[j][0], "o", color=col, ms=4)  # the ring's start point (t = 0)
        ax2.set_aspect("equal")
        ax2.set_title(f"step {p['h0'] + i}->{p['h0'] + i + 1}: coef {coef[i]:.2f}, shape {shape[i]:.2f}", fontsize=9)
        ax2.legend(fontsize=6, loc="lower left")
        ax2.set_xticks([]); ax2.set_yticks([])
    fig.tight_layout()
    fig.savefig(REPORT / "img" / f"handle_{name}.png")
    print(name, "track", p["h0"], p["h1"], "cuts", cap["cuts"])
    for i in np.argsort(-coef)[:8]:
        print(f"  step {p['h0'] + i}->{p['h0'] + i + 1}: coef {coef[i]:.2f}  shape {shape[i]:.2f}  ratio {shape[i] / coef[i]:.2f}")


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