"""Which parts fit their shape: per part and per layer, the shape distance (RMS, each ring's samples to the other's
polyline, both ways) between the sliced ring and the ring given back at the layer centre by (a) the part's full fit (all
degrees of its height basis: what the cut test judges), (b) the descriptor at 6144 atoms (the most it holds), (c) at 1024.
In layer heights. For the parts that miss most, the profile along the part: the sliced and fitted equivalent radius
(sqrt(area / pi)) and the miss per layer. Cut tests: coef (the package) and shape (shape_test.py).
Writes ../report/fit_audit.csv and ../report/img/audit_<object>_<test>.png.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2810_partition/scripts/fit_audit.py [objects...]
"""
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 _heights, _parts  # noqa: E402
from fcp._slice import centres  # noqa: E402
from shape_test import TESTS, rings, shape_miss  # noqa: E402

OUT = REPORT / "fit_audit.csv"
IMG = REPORT / "img"
FIELDS = ["object", "test", "part", "sign", "h0", "h1", "n", "degree", "full_max", "full_mean", "a6144_max", "a6144_mean",
          "a1024_max", "a1024_mean", "worst_layer"]
SHOW = 4  # parts drawn per object and test (the worst at 6144)


def area(R):
    """Signed shoelace area of rings (R, NT, 2)."""
    x, y = R[..., 0], R[..., 1]
    return 0.5 * (x * np.roll(y, -1, 1) - np.roll(x, -1, 1) * y).sum(1)


def audit(g, test):
    cap = {}
    sysf = _parts.parts_system

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

    _parts.split_by_fit, _parts.parts_system = TESTS[test], capture
    try:
        d = fcp.encode(g, atoms=6144, up=UP, parts="tracked-fit", device=DEVICE)
    finally:
        _parts.split_by_fit, _parts.parts_system = TESTS["coef"], sysf
    P, Y, _, _ = cap["out"]
    K = d.K
    B = 2 * K
    lh = float(centres(d.layers)[1] - centres(d.layers)[0])
    coefs = {T: d.prefix(T).coef() for T in (6144, 1024)}
    res = []
    for i, p in enumerate(P):
        n = p["h1"] - p["h0"] + 1
        Ys = Y[i * B:(i + 1) * B, :n].T.astype(float)  # (n, 2K): the sliced rings' coefficients (start points by continuity)
        Q, _ = _parts.orth_basis(n, d.M)
        sliced = rings(Ys.reshape(n, 2, K))
        fits = {"full": Q @ (Q.T @ Ys)}
        for T, C in coefs.items():
            fits[f"a{T}"] = _heights.centre_values(C[i * B:(i + 1) * B].T, n).astype(float)
        miss = {k: shape_miss(sliced, rings(v.reshape(n, 2, K))) / lh for k, v in fits.items()}
        r_s = np.sqrt(np.abs(area(sliced)) / np.pi)
        r_f = {k: np.sqrt(np.abs(area(rings(v.reshape(n, 2, K)))) / np.pi) for k, v in fits.items()}
        res.append(dict(i=i, p=p, n=n, deg=Q.shape[1] - 1, miss=miss, r_s=r_s, r_f=r_f, lh=lh))
    return d, res


def draw(name, test, res, path):
    worst = sorted(res, key=lambda r: -r["miss"]["a6144"].max())[:SHOW]
    fig, axs = plt.subplots(2, len(worst), figsize=(4.4 * len(worst), 6.4), dpi=100, squeeze=False)
    for c, r in enumerate(worst):
        p, h = r["p"], np.arange(r["p"]["h0"], r["p"]["h1"] + 1)
        ax = axs[0, c]
        ax.plot(h, r["r_s"] / r["lh"], "k-", lw=1.6, label="sliced")
        for k, st in (("full", "C0--"), ("a6144", "C1-"), ("a1024", "C2:")):
            ax.plot(h, r["r_f"][k] / r["lh"], st, lw=1, label=k)
        ax.set_title(f"part {r['i']} ({'hole' if p['sign'] < 0 else 'solid'}), layers {p['h0']}-{p['h1']}, degree {r['deg']}", fontsize=9)
        ax.set_ylabel("equivalent radius (layer heights)", fontsize=8)
        ax.legend(fontsize=7)
        ax = axs[1, c]
        for k, st in (("full", "C0--"), ("a6144", "C1-"), ("a1024", "C2:")):
            ax.plot(h, r["miss"][k], st, lw=1, label=k)
        ax.axhline(_parts.FIT_TOL / r["lh"], color="r", lw=0.6, ls="--", label="cut tolerance")
        ax.set_xlabel("layer"); ax.set_ylabel("shape miss (layer heights)", fontsize=8)
        ax.legend(fontsize=7)
    fig.suptitle(f"{name}, cut test {test}: the {len(worst)} parts that miss most at 6144 atoms", fontsize=11)
    fig.tight_layout()
    fig.savefig(path)
    plt.close(fig)


def main():
    names = sys.argv[1:] or list(OBJ)
    IMG.mkdir(parents=True, exist_ok=True)
    old = [r for r in csv.DictReader(OUT.open())] if OUT.exists() else []
    rows = [r for r in old if r["object"] not in names]
    for name in names:
        g = trimesh.load(OBJ[name], force="mesh")
        for test in ("coef", "shape"):
            d, res = audit(g, test)
            for r in res:
                m = r["miss"]
                rows.append(dict(object=name, test=test, part=r["i"], sign=r["p"]["sign"], h0=r["p"]["h0"], h1=r["p"]["h1"],
                                 n=r["n"], degree=r["deg"],
                                 **{f"{k}_{s}": round(float(f(m[k])), 3) for k in ("full", "a6144", "a1024")
                                    for s, f in (("max", np.max), ("mean", np.mean))},
                                 worst_layer=r["p"]["h0"] + int(np.argmax(m["a6144"]))))
            draw(name, test, res, IMG / f"audit_{name}_{test}.png")
            print(f"{name:16s} {test:5s} parts={len(res):3d}  worst full {max(r['miss']['full'].max() for r in res):.2f}"
                  f"  worst 6144 {max(r['miss']['a6144'].max() for r in res):.2f} layer heights", flush=True)
    with OUT.open("w", newline="") as f:
        w = csv.DictWriter(f, FIELDS)
        w.writeheader()
        w.writerows(rows)


if __name__ == "__main__":
    main()