"""The partition of each mesh under a few of degree_sweep.py's settings (rings by part, the decoded mesh by its nearest
part from the front and at 45 degrees, the layer spans): ../report/img/<object>_<rule>_k<kappa>_M<M>_a<atoms>.png. Decoded at 1024
and 6144 atoms (one encoding; the partition does not depend on the budget). The drawing is fcp_2710_tests' parts_report.render with the side
view turned to 45 degrees.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2810_partition/scripts/renders.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 trimesh  # noqa: E402
from matplotlib.collections import PolyCollection  # noqa: E402
from mpl_toolkits.mplot3d.art3d import Line3DCollection  # noqa: E402
from scipy.spatial import cKDTree  # noqa: E402

sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, "/home/joao/code/data/meshes/fcp_2710_tests/scripts")
from degree_sweep import DEVICE, OBJ, REPORT, UP, set_kappa  # noqa: E402  (puts the package on the path)
import fcp  # noqa: E402
from parts_report import _axes, colors, part_rings  # noqa: E402

IMG = REPORT / "img"
SHOW = [("tracked-fit", 1e2, 24), ("tracked-fit", 1e3, 35), ("tracked-fit", 1e4, 50)]  # the default, and two degree caps
ATOMS = [1024, 6144]  # one encoding at the largest, its prefixes decoded
VIEWS = [((0, -90), "rings by part (x-z)"), ((0, -90), "decoded mesh by nearest part (x-z)"),
         ((25, -45), "decoded mesh by nearest part (45°, raised 25°)")]


def name(obj, rule, kappa, M, T):
    return f"{obj}_{rule}_k{kappa:g}_M{M}_a{T}.png"


def paint(ax, T, fc, elev, azim, P):
    """Triangles T (F, 3, 3) drawn orthographically from (elev, azim) as matplotlib's view_init names them, far ones
    first (painter's order per triangle: the 3D axes' own sort lets back walls show through the caps)."""
    e, a = np.radians(elev), np.radians(azim)
    right = np.array([-np.sin(a), np.cos(a), 0.0])
    up = np.array([-np.cos(a) * np.sin(e), -np.sin(a) * np.sin(e), np.cos(e)])
    toward = np.array([np.cos(e) * np.cos(a), np.cos(e) * np.sin(a), np.sin(e)])
    order = np.argsort(T.mean(1) @ toward)
    xy = np.stack([T[order] @ right, T[order] @ up], -1)
    ax.add_collection(PolyCollection(xy, facecolors=fc[order], edgecolors=fc[order], linewidths=0.1))
    q = np.stack([P @ right, P @ up], -1)
    m, r = (q.min(0) + q.max(0)) / 2, (q.max(0) - q.min(0)).max() / 2 * 1.05
    ax.set_xlim(m[0] - r, m[0] + r); ax.set_ylim(m[1] - r, m[1] + r); ax.set_aspect("equal"); ax.set_axis_off()


def render(d, mesh, path, title):
    """Canonical frame, orthographic, height axis z up: every ring of every part (holes dashed) seen along +y; the decoded
    mesh coloured by its nearest part, seen along +y and at 45 degrees; the parts' layer spans."""
    rings = part_rings(d)
    col = colors(len(d.parts))
    fig = plt.figure(figsize=(18, 5), dpi=110)
    fig.suptitle(title, fontsize=11)
    allp = np.concatenate([r.reshape(-1, 3) for r in rings if r is not None])
    for k, ((elev, azim), label) in enumerate(VIEWS):
        if k == 0:
            ax = fig.add_subplot(1, 4, 1, projection="3d", proj_type="ortho")
            segs, cs, ls = [], [], []
            for i, r in enumerate(rings):
                if r is None:
                    continue
                for ring in r[::max(1, len(r) // 40)] if len(r) > 40 else r:
                    segs.append(np.vstack([ring, ring[:1]])); cs.append(col[i]); ls.append("--" if d.parts[i][0] < 0 else "-")
            ax.add_collection3d(Line3DCollection(segs, colors=cs, linewidths=0.7, linestyles=ls))
        elif mesh is not None:
            V = d.pose.to_canonical(mesh.vertices)
            own = np.full(len(allp), -1)
            o = 0
            for i, r in enumerate(rings):
                if r is not None:
                    own[o:o + r.size // 3] = i; o += r.size // 3
            Z = np.array([1.0, 1.0, 20.0])  # height weighted: a face takes the part at its own height (a cap's centre is far from every ring in x-y)
            part = own[cKDTree(allp * Z).query(V[mesh.faces].mean(1) * Z)[1]]
            n = np.cross(V[mesh.faces[:, 1]] - V[mesh.faces[:, 0]], V[mesh.faces[:, 2]] - V[mesh.faces[:, 0]])
            n /= np.maximum(np.linalg.norm(n, axis=1, keepdims=True), 1e-12)
            light = np.array([0.45, -0.6, 0.65]); light /= np.linalg.norm(light)
            shade = 0.35 + 0.65 * np.abs(n @ light)
            ax = fig.add_subplot(1, 4, k + 1)
            paint(ax, V[mesh.faces], col[part] * shade[:, None], elev, azim, allp)
        ax.set_title(label, fontsize=9)
        if k == 0:
            ax.view_init(elev, azim)
            _axes(ax, allp)
    ax = fig.add_subplot(1, 4, 4)
    for i, (s, h0, h1) in enumerate(d.parts):
        ax.barh(i, h1 - h0 + 1, left=h0, color=col[i], hatch="//" if s < 0 else None, edgecolor="k", linewidth=0.3)
    ax.set_xlim(0, d.layers); ax.set_xlabel("layer"); ax.set_ylabel("part"); ax.invert_yaxis()
    ax.set_title(f"{len(d.parts)} parts ({sum(s < 0 for s, _, _ in d.parts)} holes, hatched): layer spans", fontsize=9)
    fig.tight_layout()
    fig.savefig(path)
    plt.close(fig)


def main():
    IMG.mkdir(parents=True, exist_ok=True)
    for obj, path in OBJ.items():
        g = trimesh.load(path, force="mesh")
        for rule, kappa, M in SHOW:
            set_kappa(kappa)
            d = fcp.encode(g, atoms=max(ATOMS), up=UP, M=M, parts=rule, device=DEVICE)
            for T in ATOMS:
                m = fcp.decode(d.prefix(T), device=DEVICE)
                render(d.prefix(T), m, IMG / name(obj, rule, kappa, M, T),
                       f"{obj}: {rule}, κ {kappa:g}, M {M}, {T} atoms, {len(d.parts)} parts")
            print(obj, rule, kappa, M, len(d.parts), flush=True)
    set_kappa(1e2)


if __name__ == "__main__":
    main()