"""The bowl_gso texture variants decoded as points (decode(kind="points_filtered"): points on the parts' surface with
exact normals, no grid), drawn with a z-buffer (orthographic along +y, RES pixels, each point a 3x3 splat, the nearest
point per pixel kept), shaded by the normals, and
coloured by distance to GT. GT: points sampled on the mesh with normals interpolated from its vertex normals (its large
triangles do not show). Writes report/img/texture_points.png and report/texture_points.html.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2710_tests/scripts/texture_points.py
"""
import sys
import time
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

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

ENC = sys.modules["fcp.encode"]
VARIANTS = [(128, 16, 24, 6144), (256, 16, 24, 6144), (256, 64, 24, 6144), (256, 64, 48, 60000), (512, 128, 24, 6144),
            (512, 128, 48, 200000)]  # layers, N, M, atom cap (the package's is 6144)
NPTS = 1_200_000
PRES = 1024  # the points' grid: rings at PRES heights (between the layers by the stored polynomial), more per ring
RES = 700  # pixels across [-1.05, 1.05]
EMAX = 0.1  # % of the diagonal
LIGHT = np.array([0.3, -0.8, 0.5]) / np.linalg.norm([0.3, -0.8, 0.5])


def splat(ax, P, Nn, title, err=None, layers=None):
    """z-buffer: the nearest point (smallest y) per pixel, each point covering 3 pixels across and, for decoded points
    (on the layers), one layer's height (else 3 pixels)."""
    shade = 0.15 + 0.85 * np.clip(Nn @ LIGHT, 0, 1)
    if err is None:
        c = np.c_[shade, shade, shade * 1.03]
    else:
        c = plt.get_cmap("magma_r")(np.clip(err / EMAX, 0, 1))[:, :3] * (0.45 + 0.55 * shade)[:, None]
    i = np.floor((P[:, 0] + 1.05) / 2.1 * RES).astype(int)
    j = np.floor((1.05 - P[:, 2]) / 2.1 * RES).astype(int)
    pix, dep, idx = [], [], []
    hj = 1 if layers is None else max(1, int(np.ceil(2.2 / layers / 2.1 * RES / 2)))
    for di in (-1, 0, 1):
        for dj in range(-hj, hj + 1):
            ii, jj = i + di, j + dj
            ok = (ii >= 0) & (ii < RES) & (jj >= 0) & (jj < RES)
            pix.append(jj[ok] * RES + ii[ok]); dep.append(P[ok, 1] + 0.002 * (abs(di) + abs(dj))); idx.append(np.nonzero(ok)[0])
    pix, dep, idx = np.concatenate(pix), np.concatenate(dep), np.concatenate(idx)
    o = np.lexsort((dep, pix))
    first = np.r_[True, pix[o][1:] != pix[o][:-1]]
    img = np.ones((RES * RES, 3))
    img[pix[o][first]] = c[idx[o][first]].clip(0, 1)
    ax.imshow(img.reshape(RES, RES, 3), interpolation="nearest")
    ax.set_axis_off(); ax.set_title(title, fontsize=9)


def main():
    g = trimesh.load(OBJ["bowl_gso"], force="mesh")
    ref = np.asarray(g.sample(50_000, seed=1))
    rows, panels = [], []
    G = None
    for layers, N, M, cap in VARIANTS:
        ENC.MAX_ATOMS = cap
        d = fcp.encode(g, atoms=cap, layers=layers, up=UP.get("bowl_gso", (0, 0, 1)), device="cuda", N=N, M=M)
        if G is None:
            G = trimesh.Trimesh(d.pose.to_canonical(g.vertices), g.faces, process=False)
            diag = float(np.linalg.norm(np.ptp(G.vertices, 0)))
            Pg, fi = trimesh.sample.sample_surface(G, NPTS, seed=0)
            bc = trimesh.triangles.points_to_barycentric(G.triangles[fi], Pg)
            Ng = (trimesh.Trimesh(g.vertices, g.faces, process=False).vertex_normals[G.faces[fi]] * bc[:, :, None]).sum(1)
            Ng = (Ng @ d.pose.rotation) * d.pose.scale  # the world mesh's normals, as the decoded ones
            Ng /= np.linalg.norm(Ng, axis=1, keepdims=True)
            gt = np.mean([band(G, z) for z in (-0.3, 0.0, 0.4)])
            panels.append((Pg, Ng, f"GT\nharmonics 45-55 = {gt:.3f} % of r", None, None))
            rows.append(dict(variant="GT", texture_pct_r=gt))
        t = time.perf_counter()
        Pw, Nw = fcp.decode(d, kind="points_filtered", n_points=NPTS, device="cuda", resolution=PRES)
        tp = time.perf_counter() - t
        P = d.pose.to_canonical(np.asarray(Pw))
        Nn = (np.asarray(Nw) @ d.pose.rotation) * d.pose.scale  # canonical = ((x - c) R) / s: normals by R, times s
        Nn /= np.linalg.norm(Nn, axis=1, keepdims=True)
        m = fcp.decode(d, device="cuda")  # the section band and Chamfer as before (mesh)
        C = trimesh.Trimesh(d.pose.to_canonical(m.vertices), m.faces, process=False)
        b = np.mean([band(C, z) for z in (-0.3, 0.0, 0.4)])
        c = stats(d, ref, m)[0] * 1e3
        err = np.concatenate([np.abs(trimesh.proximity.closest_point(G, P[a:a + 200_000])[1])
                              for a in range(0, len(P), 200_000)]) / diag * 100  # chunks: memory
        lab = f"{layers} layers, N={N}, M={M}, {len(d.atoms[0])} atoms"
        rows.append(dict(variant=lab, atoms=len(d.atoms[0]), chamfer=round(c, 3), texture_pct_r=b,
                         points_err_mean=err.mean(), points_decode_s=round(tp, 3)))
        panels.append((P, Nn, f"{lab}\nChamfer {c:.3f}‰, harmonics 45-55 = {b:.3f} %", err, PRES))
        print(len(P), "points", flush=True)
        print(rows[-1], flush=True)
    ENC.MAX_ATOMS = 6144
    n = len(panels)
    fig, axs = plt.subplots(2, n, figsize=(4.2 * n, 9), dpi=130)
    for k, (P, Nn, t, err, layers) in enumerate(panels):
        splat(axs[0, k], P, Nn, t, layers=layers)
        if err is None:
            axs[1, k].set_axis_off()
            fig.colorbar(plt.cm.ScalarMappable(cmap="magma_r", norm=plt.Normalize(0, EMAX)), ax=axs[1, k],
                         orientation="horizontal", label="distance to GT, % of the diagonal", fraction=0.5)
        else:
            splat(axs[1, k], P, Nn, f"distance to GT, mean {err.mean():.4f} %", err, layers)
    fig.suptitle("bowl_gso clean mesh decoded as points (rings at 1024 heights, exact normals; z-buffer); raking light", fontsize=12)
    img = IMG / "texture_points.png"
    fig.savefig(img, bbox_inches="tight"); plt.close(fig)
    html = [f"<html><head><meta charset='utf-8'><title>Bowl texture, points</title><style>{CSS}</style></head><body>",
            "<h1>bowl_gso: the dimple texture, decoded as points</h1>",
            "<p>The texture is about 50 periods around (harmonics 45-55 of a section). N harmonics per ring store "
            "harmonics 1..N only; layers and M set the resolution along the height. Atom caps above 6144 are raised for "
            "the experiment. Points: decode(kind='points_filtered'), no grid.</p>",
            pd.DataFrame(rows).round(4).to_html(index=False), f"<img src='img/{img.name}' style='max-width:100%'>",
            "</body></html>"]
    (REPORT / "texture_points.html").write_text("\n".join(html))
    print(img)


if __name__ == "__main__":
    main()