"""Objects (gwn, noise0p6 10k clouds) decoded at a share of the stored energy and with all atoms, beside the clean mesh
(GT): shaded meshes, the decoded meshes coloured by distance to GT, and the middle section (x-z, canonical frame) of GT
and both decodes overlaid. The shares: the Chamfer minimum of the 0.001 % sweeps (gwn_noise0p6_*_energy_fine.csv) for
the bowls; for the dinosaur (no minimum below 100 %) the bowls' 99.856 %. Writes report/img/energy_<name>.png and
report/bowl_energy.html.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2710_tests/scripts/bowl_energy_render.py
"""
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 pandas as pd  # noqa: E402
import trimesh  # noqa: E402
from mpl_toolkits.mplot3d.art3d import Poly3DCollection  # noqa: E402

sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, "/home/joao/code/fcp_devel/api")
import fcp  # noqa: E402
from energy_eval import prefix_at  # noqa: E402
from normals_cmp import normal_err  # noqa: E402
from parts_report import CSS, IMG, OBJ, REPORT, ROOT, UP, stats  # noqa: E402

ITEMS = [("bowl_gso", 0.99856, "the Chamfer minimum"), ("bowl_of", 0.99859, "the Chamfer minimum"),
         ("dino", 0.99856, "the bowls' minimum (the dinosaur has none below 100 %)"),
         ("open_box", 0.99909, "the Chamfer minimum (shallow)")]
ONLY = sys.argv[1:]  # object names to render (default: all)
FINE = {"bowl_gso": "gwn_noise0p6_bowl_energy_fine.csv", "bowl_of": "gwn_noise0p6_bowl_energy_fine.csv",
        "dino": "gwn_noise0p6_dino_energy_fine.csv", "open_box": "gwn_noise0p6_open_box_energy_fine.csv"}
LIGHT = np.array([0.45, -0.6, 0.65]) / np.linalg.norm([0.45, -0.6, 0.65])


def canon(d, m):
    return trimesh.Trimesh(d.pose.to_canonical(m.vertices), m.faces, process=False)


def shaded(ax, m, colors=None, title=""):
    V, Fc = m.vertices, m.faces
    n = m.face_normals
    shade = 0.3 + 0.7 * np.abs(n @ LIGHT)
    fc = (np.array([0.75, 0.75, 0.78]) if colors is None else colors) * shade[:, None]
    ax.add_collection3d(Poly3DCollection(V[Fc], facecolors=np.clip(fc, 0, 1), edgecolors="none", linewidths=0))
    lo, hi = V.min(0), V.max(0)
    c, r = (lo + hi) / 2, (hi - lo).max() / 2
    ax.set_xlim(c[0] - r, c[0] + r); ax.set_ylim(c[1] - r, c[1] + r); ax.set_zlim(c[2] - r, c[2] + r)
    ax.set_box_aspect((1, 1, 1)); ax.view_init(25, -60); ax.set_axis_off(); ax.set_title(title, fontsize=9)


def section(m, ax, color, label, lw):
    s = m.section(plane_origin=(0, 0, 0), plane_normal=(0, 1, 0))
    if s is None:
        return
    for e in s.entities:
        P = s.vertices[e.points]
        ax.plot(P[:, 0], P[:, 2], color=color, lw=lw, label=label)
        label = None


def main():
    rows = []
    for name, keep, why in ITEMS:
        if ONLY and name not in ONLY:
            continue
        g = trimesh.load(OBJ[name], force="mesh")
        ref = np.asarray(g.sample(50_000, seed=1))
        P = np.asarray(trimesh.load(ROOT / "pointclouds" / f"{name}_noise0p6_10000.ply").vertices)
        d = fcp.encode(P, atoms=6144, up=UP.get(name, (0, 0, 1)), device="cuda", slicer="gwn")
        T, Tf = prefix_at(d, keep), len(d.atoms[0])
        wb, wf = (fcp.decode(d, atoms=a, device="cuda") for a in (T, Tf))
        (cb, _), (cf, _) = stats(d, ref, wb), stats(d, ref, wf)
        nb, nf = normal_err(g, wb)[0], normal_err(g, wf)[0]
        mb, mf, G = canon(d, wb), canon(d, wf), canon(d, g)
        scale = float(np.linalg.norm(np.ptp(G.vertices, 0)))
        cmap = plt.get_cmap("magma_r")
        vmax = 0.015  # 1.5 % of the diagonal

        def err_colors(m):
            dist = np.abs(trimesh.proximity.closest_point(G, m.triangles_center)[1]) / scale
            return cmap(np.clip(dist / vmax, 0, 1))[:, :3]

        pct = f"{keep * 100:.3f} %"
        fig = plt.figure(figsize=(18, 10), dpi=100)
        fig.suptitle(f"{name}: gwn, 10k points from the 0.6 % noisy mesh; {pct} = {why}; error against the clean mesh (GT)",
                     fontsize=12)
        shaded(fig.add_subplot(2, 3, 1, projection="3d", proj_type="ortho"), G, title="GT (clean mesh)")
        shaded(fig.add_subplot(2, 3, 2, projection="3d", proj_type="ortho"), mb,
               title=f"{pct} energy, {T} atoms: Chamfer {cb * 1e3:.3f}‰, normal mean {nb:.1f}°")
        shaded(fig.add_subplot(2, 3, 3, projection="3d", proj_type="ortho"), mf,
               title=f"all {Tf} atoms: Chamfer {cf * 1e3:.3f}‰, normal mean {nf:.1f}°")
        shaded(fig.add_subplot(2, 3, 5, projection="3d", proj_type="ortho"), mb, err_colors(mb), f"distance to GT, {pct}")
        shaded(fig.add_subplot(2, 3, 6, projection="3d", proj_type="ortho"), mf, err_colors(mf), "distance to GT, all atoms")
        ax = fig.add_subplot(2, 3, 4)
        section(G, ax, "black", "GT", 2.0)
        section(mf, ax, "tab:red", "all atoms", 0.9)
        section(mb, ax, "tab:blue", pct, 0.9)
        ax.set_aspect("equal"); ax.legend(fontsize=8, loc="best"); ax.set_title("middle section (x-z, canonical)", fontsize=9)
        sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(0, vmax * 100))
        fig.colorbar(sm, ax=fig.axes[-2:], shrink=0.6, label="distance to GT, % of diagonal")
        out = IMG / f"energy_{name}.png"
        fig.savefig(out, bbox_inches="tight"); plt.close(fig)
        F = pd.read_csv(REPORT / FINE[name])
        F = F[F.object == name].groupby("atoms").agg(keep_pct=("keep_pct", "min"), chamfer=("chamfer", "first"))
        rows.append((name, out.name, pct, why, T, cb, cf, nb, nf, F.sort_values("chamfer").head(5).reset_index()))
        print(out)
    html = [f"<html><head><meta charset='utf-8'><title>Energy cut</title><style>{CSS}</style></head><body>",
            "<h1>Where to stop the atoms on noisy clouds (gwn, 0.6 % noisy meshes, 10k points)</h1>",
            "<p>Encoded with gwn (v5.31) at 6144 atoms, decoded at the shortest prefix keeping a share of the stored energy. "
            "The share was swept in steps of 0.001 %: the bowls have a Chamfer minimum at 99.856 % (bowl_gso, 286 atoms, "
            "4.452 ‰; all atoms 5.32) and 99.859 % (bowl_of, 156 atoms, 4.619 ‰; all atoms 5.96); the dinosaur has none "
            "below 100 % (its Chamfer falls until about 99.99 % and stays: best 99.998 %, 3569 atoms, 4.793 ‰; all atoms "
            "4.799), so it is shown at the bowls' share. The open box: a shallow minimum at 99.909 % (810 atoms, 5.392 ‰; "
            "all atoms 5.538). Errors against the clean mesh: Chamfer in ‰ of the diagonal, "
            "normal error the mean angle to the closest clean face's normal. Renders in the canonical frame.</p>"]
    for name, img, pct, why, T, cb, cf, nb, nf, top in rows:
        html.append(f"<h2>{name}: {pct} ({why})</h2><p>{T} atoms: Chamfer {cb * 1e3:.3f} ‰, normals {nb:.1f}°; "
                    f"all atoms: {cf * 1e3:.3f} ‰, {nf:.1f}°. Best five prefixes of the fine sweep:</p>")
        html.append(top.round(3).to_html(index=False))
        html.append(f"<img src='img/{img}' style='max-width:100%'>")
    html.append("</body></html>")
    (REPORT / "bowl_energy.html").write_text("\n".join(html))
    print(REPORT / "bowl_energy.html")


if __name__ == "__main__":
    main()