"""Why "coefficients AND shape" changes so little: every piece the package's cut search tests (each track, then the pieces
of each cut), its fit's worst miss in coefficients (the points at the same t; the package's test) and in shape (closest
points), in layer heights. A piece the package cuts (coefficients over FIT_TOL = 0.58 of a layer) is spared by the AND
rule only if its shape miss is under t. Writes ../report/img/and_scatter.png and ../report/and_scatter.csv.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2810_partition/scripts/and_scatter.py
"""
import csv
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():
    rows = []
    for name, path in OBJ.items():
        g = trimesh.load(path, force="mesh")

        def split(parts, coeffs, areas, M, tol=None, follows=None):
            tol = _parts.FIT_TOL if tol is None else tol
            K = np.asarray(coeffs[0]).shape[2]
            w = np.full(K, 0.5)
            w[0] = 1.0
            Cn = [_parts.continuous_start(np.asarray(C, float), A) for C, A in zip(coeffs, areas)]
            sliced = [rings(np.asarray(C, float)) for C in coeffs]

            def log_follows(i, a, b, Cf):  # the package's own decision, both misses recorded
                coef = np.sqrt((((Cn[i][a:b] - Cf) ** 2) * w).sum((1, 2))).max()
                shape = miss_batch(sliced[i][a:b], rings(np.asarray(Cf, float))).max()
                rows.append(dict(object=name, h0=parts[i]["h0"] + a, n=b - a, coef_lh=coef / LH, shape_lh=shape / LH))
                return coef <= tol

            return _split(parts, coeffs, areas, M, tol, follows=log_follows)

        _parts.split_by_fit = split
        try:
            d = fcp.encode(g, atoms=1, up=UP, device=DEVICE)
        finally:
            _parts.split_by_fit = _split
        print(name, len(d.parts), "parts", flush=True)
    with (REPORT / "and_scatter.csv").open("w", newline="") as f:
        wr = csv.DictWriter(f, list(rows[0]))
        wr.writeheader()
        wr.writerows(rows)
    c = np.array([r["coef_lh"] for r in rows])
    s = np.array([r["shape_lh"] for r in rows])
    tol = _parts.FIT_TOL / LH
    cut = c > tol
    print(f"pieces tested {len(rows)}; the package cuts {cut.sum()}; of those, shape miss under t:")
    for t in (0.15, 0.2, 0.25, 0.3, 0.4, 0.58):
        print(f"  t = {t:4.2f}: {int((cut & (s <= t)).sum()):3d}")
    print(f"shape / coefficient miss over the cut pieces: median {np.median(s[cut] / c[cut]):.2f}, "
          f"10-90 % {np.percentile(s[cut] / c[cut], 10):.2f}-{np.percentile(s[cut] / c[cut], 90):.2f}")
    fig, ax = plt.subplots(figsize=(8, 7), dpi=100)
    names = list(OBJ)
    for k, o in enumerate(names):
        sel = np.array([r["object"] == o for r in rows])
        ax.scatter(c[sel], s[sel], s=14, color=f"C{k}", label=o, alpha=0.8)
    ax.axvline(tol, color="r", lw=1, ls="--", label="coefficient tolerance 0.01 (0.58 layer)")
    for t in (0.15, 0.3):
        ax.axhline(t, color="0.4", lw=0.8, ls=":")
        ax.text(0.012, t * 1.07, f"shape t = {t}", fontsize=8, color="0.3")
    lim = [1e-2, max(c.max(), s.max()) * 1.5]  # pieces fitted exactly (short ones: misses near 0) left out of view
    ax.plot(lim, lim, color="0.7", lw=0.8, label="shape = coefficients")
    ax.set_xscale("log"); ax.set_yscale("log")
    ax.set_xlim(lim); ax.set_ylim(lim)
    ax.set_xlabel("worst miss in coefficients (points at the same t), layer heights")
    ax.set_ylabel("worst miss in shape (closest points), layer heights")
    ax.set_title("every piece the package's cut search tests: cut when right of the red line;\n"
                 "AND spares it only below the shape tolerance", fontsize=10)
    ax.legend(fontsize=7, loc="upper left")
    ax.grid(alpha=0.3, which="both")
    fig.tight_layout()
    fig.savefig(REPORT / "img" / "and_scatter.png")


if __name__ == "__main__":
    main()