"""The shape change between neighbouring rings along each track (the tracks split_by_fit receives, before any cut): ring
i and ring i + 1 each sampled at T points from their Fourier series, the RMS distance of each one's samples to the other's
polyline, both ways (no fit, and no dependence on where a ring starts). In layer heights. A smooth stretch changes slowly
(the slope times a layer height); a ledge or something joining is a spike, a corner a change of level. Marked: the
package's cuts (greedy). Tracks of at least MIN_LEN layers. Writes ../report/img/steps_<object>.png.

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

MIN_LEN = 20
_split = _parts.split_by_fit


def tracks_and_cuts(g):
    """The tracks (sign, h0, h1, sliced coefficients) and, per track, the package's cut layers."""
    cap = []

    def split(parts, coeffs, areas, M, tol=None, follows=None):
        for p, C, A in zip(parts, coeffs, areas):  # each track cut on its own: its own cuts only
            pieces = _split([p], [C], [A], M, tol, follows)[0]
            cap.append((p, np.asarray(C, float), [q["h0"] for q in pieces[1:]]))
        return _split(parts, coeffs, areas, M, tol, follows)

    _parts.split_by_fit = split
    try:
        fcp.encode(g, atoms=1, up=UP, device=DEVICE)
    finally:
        _parts.split_by_fit = _split
    return cap


def main():
    names = sys.argv[1:] or ["elephant", "dino", "horse", "shoe", "bowl_gso"]
    for name in names:
        g = trimesh.load(OBJ[name], force="mesh")
        tr = [t for t in tracks_and_cuts(g) if t[0]["h1"] - t[0]["h0"] + 1 >= MIN_LEN]
        fig, axs = plt.subplots(len(tr), 1, figsize=(13, 2.6 * len(tr)), dpi=100, squeeze=False)
        for ax, (p, C, cuts) in zip(axs[:, 0], tr):
            R = rings(C)
            s = miss_batch(R[:-1], R[1:]) / LH  # between ring i and ring i + 1
            x = np.arange(p["h0"], p["h1"]) + 0.5  # the step sits between the two layers
            ax.plot(x, s, "k.-", lw=1, ms=4)
            for c in cuts:
                ax.axvline(c - 0.5, color="C3", lw=1, alpha=0.7)
            ax.set_ylabel("ring-to-ring\n(layer heights)", fontsize=8)
            ax.set_title(f"{'hole' if p['sign'] < 0 else 'solid'} track, layers {p['h0']}-{p['h1']}; red: the package's cuts "
                         f"({len(cuts)})", fontsize=9)
            ax.grid(alpha=0.3)
            ax.set_yscale("log")
        axs[-1, 0].set_xlabel("layer")
        fig.suptitle(f"{name}: shape change between neighbouring rings along each track of {MIN_LEN}+ layers", fontsize=11)
        fig.tight_layout()
        fig.savefig(REPORT / "img" / f"steps_{name}.png")
        plt.close(fig)
        print(name, len(tr), "tracks", flush=True)


if __name__ == "__main__":
    main()