"""Renders at the knee (knee.py's descriptors and scores): per noise level one image, a row per object, columns mesh route /
sigma given, full / sigma = 0, full / sigma = 0 at its knee / sigma given at its knee. Decoded at 96 cells per axis for the
picture, coloured by part; titles: atoms (energy), parts, Chamfer x1e-3 and normal median from knee.csv.
-> ../report/img/knee_render_<cloud>.png

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python knee_render.py pn0p6_10000,pn1p2_10000
"""
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import pandas as pd

sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, "/home/joao/code/data/meshes/fcp_2710_tests/scripts")
sys.path.insert(0, "/home/joao/code/fcp_devel/api")
import fcp  # noqa: E402
from gallery_energy import draw  # noqa: E402
from parts_report import OBJ  # noqa: E402

REP = Path(__file__).parent.parent / "report"
K = pd.read_csv(REP / "knee.csv")


def row(name, cl, which):
    g = K[(K.object == name) & (K.cloud == cl)]
    r = g[g.is_knee == 1].iloc[0] if which == "knee" else g[g.atoms == g.stored].iloc[0]
    return r


def main(clouds):
    objs = [o for o in OBJ if (K.object == o).any()]
    for cl in clouds:
        cols = [("mesh route", "mesh", "full"), ("σ given, full", cl, "full"), ("σ = 0, full", f"{cl}_s0", "full"),
                ("σ = 0, knee", f"{cl}_s0", "knee"), ("σ given, knee", cl, "knee")]
        fig = plt.figure(figsize=(len(cols) * 2.7, len(objs) * 2.7), dpi=100)
        for r, name in enumerate(objs):
            for k, (lab, c, which) in enumerate(cols):
                d = fcp.load(REP / "knee" / f"{name}_{c}.fcp")
                s = row(name, c, which)
                dp = d.prefix(int(s.atoms))
                ax = fig.add_subplot(len(objs), len(cols), r * len(cols) + k + 1, projection="3d", proj_type="ortho")
                draw(ax, dp, fcp.decode(dp, device="cuda", resolution=96),
                     f"{name} · {lab}\n{int(s.atoms)} atoms ({100 * s.energy:.2f} %) · {int(s.parts)} parts\n"
                     f"err {s['mean']:.2f} · normal {s.n_med:.1f}°")
            print(cl, name, flush=True)
        fig.suptitle(f"{cl}: full descriptors (6144 atoms) and their knees; σ = 0 keeps the noise in the atoms", fontsize=11, y=0.999)
        fig.tight_layout()
        fig.savefig(REP / "img" / f"knee_render_{cl}.png")
        plt.close(fig)
        print("saved", cl, flush=True)


if __name__ == "__main__":
    main(sys.argv[1].split(","))