#!/usr/bin/env python3
"""Render every measured input-size candidate as a PNG from its saved NPZ."""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import ListedColormap
import numpy as np
from scipy.spatial import cKDTree

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]
PART_ROOTS = {"AC28": ROOT, "ME1a_DT": ROOT / "experiments/ME1a_DT"}
IMAGE_SIZES = (64, 128, 256, 512, 1024, 2048)
VOXEL_SIZES = (32, 64, 128, 256, 512, 1024)
GRAPH_SIZES = (3000, 5000, 6000, 12000, 24000, 48000, 96000)
POINT_SIZES = (2048, 4096, 8192, 16384, 32768, 65536)
VIEWS = ("+X", "-X", "+Y", "-Y", "+Z", "-Z")
INK = "#16324f"
TEAL = "#087f8c"


def _style_axis(ax, x: str, y: str) -> None:
    ax.set_xlabel(f"{x} (mm)")
    ax.set_ylabel(f"{y} (mm)")
    ax.grid(alpha=.14, linewidth=.5)
    ax.set_aspect("equal", adjustable="box")


def _image(part: str, size: int, source: Path, target: Path) -> dict:
    with np.load(source, allow_pickle=False) as z:
        depth = z["depth_mm"]
        visible = z["visibility"]
        fov = float(z["field_of_view_mm"][0])
        pitch = float(z["pixel_pitch_mm"][0])
        channels = int(z["x"].shape[1]) if "x" in z else int(z["channel_count"])
        geometry_only = "x" not in z
    with np.load(source.with_name("image_256.npz"), allow_pickle=False) as z:
        reference = z["depth_mm"]
        reference_visible = z["visibility"]
    cmap = plt.get_cmap("viridis").copy()
    cmap.set_bad("#eef3f7")
    fig, axs = plt.subplots(2, 3, figsize=(12, 7), constrained_layout=True)
    for i, ax in enumerate(axs.flat):
        values = reference[i][reference_visible[i]]
        lo, hi = np.quantile(values, [.01, .99]) if len(values) else (0, 1)
        ax.imshow(np.ma.masked_where(~visible[i], depth[i]), cmap=cmap,
                  vmin=lo, vmax=hi, origin="lower", interpolation="nearest")
        ax.set_title(f"View {VIEWS[i]}  |  {int(visible[i].sum()):,} visible pixels", fontsize=10)
        ax.axis("off")
    fig.suptitle(f"{part}  |  IMAGE {size} × {size}  |  {pitch:.2f} mm/pixel  |  {channels}-channel contract, {fov:.0f} mm FOV"
                 + ("  |  geometry diagnostic only" if geometry_only else ""),
                 fontsize=13, weight="bold", color=INK)
    fig.savefig(target, dpi=135, facecolor="white")
    plt.close(fig)
    return {"physical_pitch_mm": round(pitch, 4), "views": 6, "channels": channels,
            "geometry_diagnostic_only": geometry_only}


def _voxel(part: str, size: int, source: Path, target: Path) -> dict:
    with np.load(source, allow_pickle=False) as z:
        coords = z["voxel_coords"]
        dims = tuple(int(x) for x in z["grid_dims"])
        pitch = float(z["voxel_pitch_mm"])
    occ = np.zeros(dims, dtype=np.uint8)
    occ[coords[:, 0], coords[:, 1], coords[:, 2]] = 1
    panels = [
        (occ.max(axis=1).T, dims[0] * pitch, dims[2] * pitch, "X–Z projection", "X", "Z"),
        (occ.max(axis=2).T, dims[0] * pitch, dims[1] * pitch, "X–Y projection", "X", "Y"),
        (occ.max(axis=0).T, dims[1] * pitch, dims[2] * pitch, "Y–Z projection", "Y", "Z"),
    ]
    fig, axs = plt.subplots(1, 3, figsize=(14, 5), constrained_layout=True)
    cmap = ListedColormap(["#f4f7fa", TEAL])
    for ax, (array, width, height, title, x, y) in zip(axs, panels):
        ax.imshow(array, origin="lower", interpolation="nearest", cmap=cmap,
                  extent=(0, width, 0, height), vmin=0, vmax=1)
        ax.set_xlim(0, width); ax.set_ylim(0, height)
        ax.set_title(title, color=INK)
        _style_axis(ax, x, y)
    fig.suptitle(f"{part}  |  SURFACE VOXELS {dims[0]} × {dims[1]} × {dims[2]}  |  {pitch:g} mm pitch  |  {len(coords):,} occupied cells",
                 fontsize=13, weight="bold", color=INK)
    fig.savefig(target, dpi=135, facecolor="white")
    plt.close(fig)
    return {"grid_shape": dims, "physical_pitch_mm": pitch, "occupied_cells": len(coords)}


def _spacing(pos: np.ndarray) -> np.ndarray:
    return cKDTree(pos).query(pos, k=2, workers=-1)[0][:, 1]


def _graph(part: str, size: int, source: Path, target: Path) -> dict:
    with np.load(source, allow_pickle=False) as z:
        pos = z["pos_mm"]
        edges = z["edge_index"]
    unique = edges[:, edges[0] < edges[1]].T
    if not len(unique):
        unique = edges.T[::2]
    # Large meshes are sampled only for the preview; the NPZ retains every edge.
    shown = unique if len(unique) <= 75000 else unique[np.linspace(0, len(unique)-1, 75000, dtype=int)]
    segments_xz = pos[shown][:, :, [0, 2]]
    spacing = _spacing(pos)
    fig, axs = plt.subplots(1, 3, figsize=(14, 4.8), constrained_layout=True)
    axs[0].add_collection(LineCollection(segments_xz, colors="#adbeca", linewidths=.16, alpha=.45))
    axs[0].scatter(pos[:, 0], pos[:, 2], s=.48, c=TEAL, linewidths=0, rasterized=True)
    axs[0].autoscale(); axs[0].set_title("X–Z nodes and mesh edges", color=INK)
    _style_axis(axs[0], "X", "Z")
    axs[1].scatter(pos[:, 0], pos[:, 1], s=.48, c=TEAL, linewidths=0, rasterized=True)
    axs[1].set_title("X–Y node distribution", color=INK)
    _style_axis(axs[1], "X", "Y")
    axs[2].hist(spacing, bins=35, color=TEAL, edgecolor="white")
    axs[2].set_title("Nearest-node spacing", color=INK)
    axs[2].set_xlabel("Distance (mm)"); axs[2].set_ylabel("Nodes")
    axs[2].grid(alpha=.14, axis="y")
    fig.suptitle(f"{part}  |  GRAPH {size:,} requested → {len(pos):,} actual nodes  |  {len(unique):,} undirected edges",
                 fontsize=13, weight="bold", color=INK)
    fig.savefig(target, dpi=135, facecolor="white")
    plt.close(fig)
    return {"actual_nodes": len(pos), "undirected_edges": len(unique),
            "nearest_node_p95_mm": round(float(np.quantile(spacing, .95)), 4)}


def _point(part: str, size: int, source: Path, target: Path) -> dict:
    with np.load(source, allow_pickle=False) as z:
        pos = z["pos_mm"]
    spacing = _spacing(pos)
    fig, axs = plt.subplots(1, 3, figsize=(14, 4.8), constrained_layout=True)
    for ax, j, title, y in ((axs[0], 2, "X–Z surface samples", "Z"),
                             (axs[1], 1, "X–Y surface samples", "Y")):
        ax.scatter(pos[:, 0], pos[:, j], s=1.2, c=TEAL, linewidths=0, rasterized=True)
        ax.set_title(title, color=INK)
        _style_axis(ax, "X", y)
    axs[2].hist(spacing, bins=35, color=TEAL, edgecolor="white")
    axs[2].set_title("Nearest-point spacing", color=INK)
    axs[2].set_xlabel("Distance (mm)"); axs[2].set_ylabel("Points")
    axs[2].grid(alpha=.14, axis="y")
    fig.suptitle(f"{part}  |  CAD SURFACE POINTS {len(pos):,}  |  deterministic area-uniform sample",
                 fontsize=13, weight="bold", color=INK)
    fig.savefig(target, dpi=135, facecolor="white")
    plt.close(fig)
    return {"points": len(pos), "nearest_point_p95_mm": round(float(np.quantile(spacing, .95)), 4)}


def main() -> None:
    manifest = []
    for part, part_root in PART_ROOTS.items():
        folder = HERE / "candidates" / part
        for modality, sizes, render in (("image", IMAGE_SIZES, _image),
                                        ("voxel", VOXEL_SIZES, _voxel),
                                        ("graph", GRAPH_SIZES, _graph),
                                        ("points", POINT_SIZES, _point)):
            for size in sizes:
                source = ((part_root / f"data_processed/graphs/graph_{size}.npz") if size <= 12000 else
                          (folder / "high_graphs" / f"graph_{size}.npz")) if modality == "graph" else (folder / f"{modality}_{size}.npz")
                target = folder / f"{modality}_{size}.png"
                metrics = render(part, size, source, target)
                manifest.append({"part": part, "modality": modality, "requested_size": size,
                                 "source": str(source.relative_to(ROOT)),
                                 "png": str(target.relative_to(HERE)), **metrics})
                print(f"{part} {modality} {size} -> {target.relative_to(HERE)}")
    (HERE / "candidate_samples_manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2))
    print(f"Rendered {len(manifest)} candidate PNGs")


if __name__ == "__main__":
    main()