"""The object in 3D (the source mesh, shaded) with the descriptor's wrong points beyond 0.1 % of the diagonal on it,
coloured by band as in slice_errors.py, from four sides; one row per cached descriptor.

    env -u PYTHONPATH .venv/bin/python scripts/render3d.py <object> <out tag> <label>=<cache .fcp> [...]
"""
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 open3d as o3d  # noqa: E402
from matplotlib.lines import Line2D  # noqa: E402
from mpl_toolkits.mplot3d.art3d import Poly3DCollection  # noqa: E402

sys.path.insert(0, str(Path(__file__).parent))
from common import CACHE, REPORT, RULE, cut, fcp, load_mesh  # noqa: E402
from fcpq import FCPQuery  # noqa: E402
from slice_errors import BANDS, COLORS  # noqa: E402

VIEWS = [("front", 10, -90), ("side", 10, 0), ("back", 10, 90), ("below", -60, -60)]
MAX_FACES = 30000


def main():
    name, tag, descs = sys.argv[1], sys.argv[2], [a.split("=", 1) for a in sys.argv[3:]]
    m = load_mesh(name)
    sc = o3d.t.geometry.RaycastingScene()
    sc.add_triangles(o3d.core.Tensor(np.asarray(m.vertices, np.float32)), o3d.core.Tensor(np.asarray(m.faces, np.uint32)))
    np.random.seed(0)
    rng = np.random.default_rng(0)
    diag = np.linalg.norm(m.bounds[1] - m.bounds[0])
    S, fi = m.sample(1_000_000, return_index=True)
    P = S + m.face_normals[fi] * rng.uniform(-0.02, 0.02, (len(S), 1)) * diag
    T = o3d.core.Tensor(P.astype(np.float32))
    gt = sc.compute_occupancy(T, nsamples=3).numpy() > 0
    dist = np.abs(sc.compute_distance(T).numpy()) / diag * 100

    # the mesh, decimated for drawing, shaded by a fixed light
    mm = m
    if len(m.faces) > MAX_FACES:
        o = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(m.vertices), o3d.utility.Vector3iVector(m.faces))
        o = o.simplify_quadric_decimation(MAX_FACES)
        import trimesh

        mm = trimesh.Trimesh(np.asarray(o.vertices), np.asarray(o.triangles), process=False)
    tri = np.asarray(mm.vertices)[np.asarray(mm.faces)]
    light = np.array([0.3, -0.5, 0.8]) / np.linalg.norm([0.3, -0.5, 0.8])
    shade = 0.45 + 0.5 * np.clip(np.asarray(mm.face_normals) @ light, 0, 1)
    fc = np.column_stack([shade * 0.86, shade * 0.86, shade * 0.88, np.full(len(shade), 0.35)])
    lo, hi = m.bounds
    c, r = (lo + hi) / 2, (hi - lo).max() / 2

    fig = plt.figure(figsize=(5 * len(VIEWS), 5.2 * len(descs)))
    for row, (label, f) in enumerate(descs):
        d = cut(fcp.FCP.from_bytes((CACHE / f).read_bytes()))  # the study's cut
        q = FCPQuery(d, nt=max(1024, 8 * d.N))
        ins = q.inside_cpu(d.pose.to_canonical(P).astype(np.float32), RULE)
        bad = (ins != gt) & (dist > BANDS[0][0])
        counts = ", ".join(f"> {lo_}: {int((bad & (dist > lo_)).sum())}" for lo_, _ in BANDS)
        for col, (vname, elev, azim) in enumerate(VIEWS):
            ax = fig.add_subplot(len(descs), len(VIEWS), row * len(VIEWS) + col + 1, projection="3d")
            ax.add_collection3d(Poly3DCollection(tri, facecolors=fc, edgecolors="none"))
            for (b0, b1), colr in zip(BANDS, COLORS):
                s = bad & (dist > b0) & (dist <= b1)
                ax.scatter(*P[s].T, s=7 + 8 * BANDS.index((b0, b1)), color=colr, depthshade=False)
            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(elev=elev, azim=azim)
            ax.set_axis_off()
            ax.set_title(vname if col else f"{label}: {int(bad.sum())} wrong > 0.1 % diag\n({counts.split(', ', 1)[1]})",
                         fontsize=10, loc="left")
    handles = [Line2D([], [], ls="", marker="o", color=c_, label=f"{a}–{b} % diag" if np.isfinite(b) else f"> {a} % diag")
               for (a, b), c_ in zip(BANDS, COLORS)]
    fig.legend(handles=handles, loc="lower center", ncol=4, fontsize=10, frameon=False)
    fig.tight_layout(rect=(0, 0.05, 1, 1))
    out = REPORT / "slices" / f"3d_{tag}_{name}.png"
    fig.savefig(out, dpi=100, bbox_inches="tight", facecolor="white")
    print(out)


if __name__ == "__main__":
    main()