#!/usr/bin/env python3
"""Honest CAD/scan representations for the Moltex POC.

This module deliberately keeps CAD-only inputs separate from post-process scans.
It uses the existing triangulated CAD graph as a mesh: rasterized pixels retain a
triangle id and barycentric coordinates, so all image fields are from one visible
front-most surface point.  It does not infer thickness, stress, defect labels, or
thermal histories.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Iterable

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.ndimage import distance_transform_edt
from scipy.spatial import cKDTree


CAD_FEATURE_NAMES = np.asarray([
    "x_mm", "y_mm", "z_mm", "normal_x", "normal_y", "normal_z",
    "normal_turn_deg", "curvature_magnitude_per_mm", "gaussian_curvature_per_mm2",
    "curvature_valid", "thickness_mm", "thermal_max_c", "stress_mpa",
])
CAD_FEATURE_UNITS = np.asarray([
    "mm", "mm", "mm", "unit", "unit", "unit", "deg", "1/mm", "1/mm2",
    "bool", "mm", "C", "MPa",
])
SCAN_FEATURE_NAMES = np.asarray([
    "x_native", "y_native", "z_native", "normal_x", "normal_y", "normal_z",
    "nn_spacing_native", "local_surface_variation",
])
IMAGE_FEATURE_NAMES = np.asarray([
    "x_mm", "y_mm", "z_mm", "depth_mm", "normal_x", "normal_y", "normal_z",
    "curvature_magnitude_per_mm", "gaussian_curvature_per_mm2", "curvature_valid",
    "silhouette", "distance_to_boundary_px", "normal_orientation_to_view", "visibility",
    "thickness_mm", "thermal_max_c", "stress_mpa",
])
IMAGE_FEATURE_UNITS = np.asarray([
    "mm", "mm", "mm", "mm", "unit", "unit", "unit", "1/mm", "1/mm2",
    "bool", "bool", "px", "unit", "bool", "mm", "C", "MPa",
])


def repo_root() -> Path:
    return Path(__file__).resolve().parents[1]


def json_dump(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        json.dump(value, f, ensure_ascii=False, indent=2, sort_keys=True)


def finite32(a: np.ndarray) -> np.ndarray:
    return np.nan_to_num(np.asarray(a), nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)


def load_cad_graph(path: Path) -> dict[str, np.ndarray]:
    """Load the pre-existing CAD graph; no graph generator is used."""
    with np.load(path, allow_pickle=False) as d:
        out = {k: d[k].copy() for k in d.files}
    required = {"pos_mm", "normal", "face", "k1_per_mm", "k2_per_mm", "curvature_valid"}
    missing = required - out.keys()
    if missing:
        raise ValueError(f"CAD graph missing {sorted(missing)}")
    return out


def _view_basis(view: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Stable screen basis (u right, w up) for a camera-to-object axis."""
    reference = np.array([0.0, 0.0, 1.0]) if abs(view[2]) < .9 else np.array([0.0, 1.0, 0.0])
    u = np.cross(reference, view); u /= np.linalg.norm(u)
    w = np.cross(view, u); w /= np.linalg.norm(w)
    return u, w


def rasterize_view(
    pos_mm: np.ndarray, faces: np.ndarray, normal: np.ndarray,
    curvature_magnitude: np.ndarray, gaussian: np.ndarray, curvature_valid: np.ndarray,
    view: Iterable[float], resolution: int, field_of_view_mm: float | None = None,
) -> dict[str, np.ndarray]:
    """Orthographically rasterize triangles with a largest-depth (front-most) z-buffer.

    ``view`` points from object toward camera.  The returned source_triangle and
    barycentric arrays enable exact backprojection to the source CAD triangle.
    """
    p = np.asarray(pos_mm, dtype=np.float64)
    f = np.asarray(faces, dtype=np.int64).T if np.asarray(faces).shape[0] == 3 else np.asarray(faces, dtype=np.int64)
    n = np.asarray(normal, dtype=np.float64)
    v = np.asarray(view, dtype=np.float64); v /= np.linalg.norm(v)
    u, w = _view_basis(v)
    pu, pw, depth = p @ u, p @ w, p @ v
    # A square physical FOV gives a single mm/pixel in both screen directions.
    # Do not stretch a short projected axis to fill the image.
    u0, u1 = float(pu.min()), float(pu.max()); w0, w1 = float(pw.min()), float(pw.max())
    side = float(field_of_view_mm or max(u1-u0, w1-w0, 1.0)*1.02)
    if side < max(u1-u0, w1-w0):
        raise ValueError("field_of_view_mm does not contain the projected CAD")
    uc, wc = (u0+u1)/2, (w0+w1)/2
    u0, u1, w0, w1 = uc-side/2, uc+side/2, wc-side/2, wc+side/2
    pitch = side / (resolution-1)
    px = (pu-u0)/pitch; py = (w1-pw)/pitch
    shape = (resolution, resolution)
    source = np.full(shape, -1, dtype=np.int32)
    zbuf = np.full(shape, -np.inf, dtype=np.float64)
    bary = np.zeros((*shape, 3), dtype=np.float32)
    # Rasterization is per-triangle, not point splatting: it intentionally fills
    # triangles and provides interpolation correspondence at every filled pixel.
    for ti, tri in enumerate(f):
        qx, qy, qz = px[tri], py[tri], depth[tri]
        xmin, xmax = max(0, int(np.ceil(qx.min()-.5))), min(resolution-1, int(np.floor(qx.max()-.5)))
        ymin, ymax = max(0, int(np.ceil(qy.min()-.5))), min(resolution-1, int(np.floor(qy.max()-.5)))
        if xmin > xmax or ymin > ymax:
            continue
        denom = (qy[1]-qy[2])*(qx[0]-qx[2]) + (qx[2]-qx[1])*(qy[0]-qy[2])
        if abs(denom) < 1e-12:
            continue
        yy, xx = np.mgrid[ymin:ymax+1, xmin:xmax+1]
        xx = xx + .5; yy = yy + .5
        b0 = ((qy[1]-qy[2])*(xx-qx[2]) + (qx[2]-qx[1])*(yy-qy[2])) / denom
        b1 = ((qy[2]-qy[0])*(xx-qx[2]) + (qx[0]-qx[2])*(yy-qy[2])) / denom
        b2 = 1.0-b0-b1
        inside = (b0 >= -1e-7) & (b1 >= -1e-7) & (b2 >= -1e-7)
        zz = b0*qz[0] + b1*qz[1] + b2*qz[2]
        current = zbuf[ymin:ymax+1, xmin:xmax+1]
        take = inside & (zz > current)
        if not np.any(take):
            continue
        current[take] = zz[take]
        ss = source[ymin:ymax+1, xmin:xmax+1]; ss[take] = ti
        bb = bary[ymin:ymax+1, xmin:xmax+1]
        bb[..., 0][take], bb[..., 1][take], bb[..., 2][take] = b0[take], b1[take], b2[take]
    visible = source >= 0
    xyz = np.zeros((*shape, 3), dtype=np.float32)
    nor = np.zeros((*shape, 3), dtype=np.float32)
    curv_mag = np.zeros(shape, dtype=np.float32); gauss = np.zeros(shape, dtype=np.float32)
    curv_ok = np.zeros(shape, dtype=bool)
    rows, cols = np.where(visible)
    for y, x in zip(rows, cols):
        tri = f[source[y, x]]; b = bary[y, x].astype(np.float64)
        xyz[y, x] = p[tri].T @ b
        ni = n[tri].T @ b; norm = np.linalg.norm(ni)
        nor[y, x] = ni / norm if norm else 0.0
        curv_mag[y, x] = np.dot(b, curvature_magnitude[tri])
        gauss[y, x] = np.dot(b, gaussian[tri])
        # Require all vertices to be mapped, avoiding invented curvature in blends.
        curv_ok[y, x] = bool(np.all(curvature_valid[tri]))
    # Unknown curvature is explicitly zero-imputed only with a false validity mask.
    curv_mag[~curv_ok] = 0.0; gauss[~curv_ok] = 0.0
    # Distance to the outside is a silhouette-distance map (pixels); boundary is
    # implicit at distance near zero and remains valid only on visible pixels.
    silhouette_distance = distance_transform_edt(visible).astype(np.float32) * visible
    orientation = np.sum(nor*v, axis=-1).astype(np.float32) * visible
    return {
        "xyz_mm": xyz, "depth_mm": finite32(zbuf) * visible, "normal": nor,
        "curvature_magnitude_per_mm": curv_mag, "gaussian_curvature_per_mm2": gauss,
        "curvature_valid": curv_ok, "silhouette": visible, "distance_to_boundary_px": silhouette_distance,
        "normal_orientation_to_view": orientation, "visibility": visible,
        "source_triangle": source, "barycentric": bary,
        "view_object_to_camera": v.astype(np.float32), "screen_u": u.astype(np.float32), "screen_w": w.astype(np.float32),
        "orthographic_bounds_uw_mm": np.asarray([u0, u1, w0, w1], dtype=np.float32),
        "pixel_pitch_mm": np.asarray(pitch, dtype=np.float32),
    }


def make_multiview(cad: dict[str, np.ndarray], resolution: int,
                   field_of_view_mm: float | None = None) -> dict[str, np.ndarray]:
    k1, k2 = cad["k1_per_mm"], cad["k2_per_mm"]
    mag = np.sqrt(k1*k1 + k2*k2)
    # Principal curvature sign may flip with local orientation; magnitude and K=k1*k2
    # are the exported orientation-invariant choices supplied by the analytic mapping.
    gaussian = cad.get("gaussian_signed_per_mm2", k1*k2)
    directions = np.asarray([[1,0,0],[-1,0,0],[0,1,0],[0,-1,0],[0,0,1],[0,0,-1]], dtype=np.float32)
    # Same CAD-wide physical FOV for all six views; this is an isotropic camera
    # scale, not a per-axis normalization.
    if field_of_view_mm is None:
        field_of_view_mm = max(np.ptp(cad["pos_mm"], axis=0).max(), 1.0) * 1.02
    renders = [rasterize_view(cad["pos_mm"], cad["face"], cad["normal"], mag, gaussian,
                              cad["curvature_valid"], v, resolution, field_of_view_mm) for v in directions]
    out = {k: np.stack([r[k] for r in renders]) for k in renders[0]}
    # Generic model contract alongside the richer raw correspondence arrays.
    x = np.stack([
        out["xyz_mm"][..., 0], out["xyz_mm"][..., 1], out["xyz_mm"][..., 2], out["depth_mm"],
        out["normal"][..., 0], out["normal"][..., 1], out["normal"][..., 2],
        out["curvature_magnitude_per_mm"], out["gaussian_curvature_per_mm2"], out["curvature_valid"],
        out["silhouette"], out["distance_to_boundary_px"], out["normal_orientation_to_view"], out["visibility"],
        np.zeros_like(out["depth_mm"]), np.zeros_like(out["depth_mm"]), np.zeros_like(out["depth_mm"]),
    ], axis=1).astype(np.float32)
    visible = out["visibility"]
    valid = np.stack([
        visible, visible, visible, visible, visible, visible, visible,
        out["curvature_valid"], out["curvature_valid"], np.ones_like(visible),
        np.ones_like(visible), visible, visible, np.ones_like(visible),
        np.zeros_like(visible), np.zeros_like(visible), np.zeros_like(visible),
    ], axis=1).astype(bool)
    out.update({"x": x, "feature_valid": valid, "feature_names": IMAGE_FEATURE_NAMES,
                "feature_units": IMAGE_FEATURE_UNITS,
                "training_valid": np.zeros_like(visible, dtype=bool),
                "field_of_view_mm": np.full(6, field_of_view_mm, dtype=np.float32)})
    return out


def deterministic_sample_indices(n: int, count: int, seed: int = 20260914) -> np.ndarray:
    """Stable without-replacement sample independent of dictionary/hash ordering."""
    count = min(int(count), int(n))
    return np.sort(np.random.default_rng(seed).choice(n, size=count, replace=False))


def area_uniform_surface_samples(cad: dict[str, np.ndarray], count: int, seed: int = 20260914) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Deterministic area-uniform samples with triangle/barycentric correspondence."""
    p = np.asarray(cad["pos_mm"], dtype=np.float64)
    face = np.asarray(cad["face"], dtype=np.int64).T
    tri = p[face]
    area = .5*np.linalg.norm(np.cross(tri[:,1]-tri[:,0], tri[:,2]-tri[:,0]), axis=1)
    if not np.any(area > 0):
        raise ValueError("CAD mesh has no non-degenerate triangles")
    rng = np.random.default_rng(seed)
    source = rng.choice(len(face), size=int(count), replace=True, p=area/area.sum())
    r = rng.random((int(count), 2)); sr = np.sqrt(r[:,0])
    bary = np.column_stack([1-sr, sr*(1-r[:,1]), sr*r[:,1]])
    return np.einsum("ni,nij->nj", bary, tri[source]), source.astype(np.int32), bary.astype(np.float32)


def cad_points(cad: dict[str, np.ndarray], count: int = 8192) -> dict[str, np.ndarray]:
    if "face" not in cad:
        # Unit-test/small utility fallback only; production CAD graph has triangles.
        ix = deterministic_sample_indices(len(cad["pos_mm"]), count)
        sample_pos = cad["pos_mm"][ix]; tri = None; bary = None
        source = ix.astype(np.int32); normal = cad["normal"][ix]; turn = cad["normal_turn_deg"][ix]
        mag = np.sqrt(cad["k1_per_mm"][ix]**2+cad["k2_per_mm"][ix]**2); gaussian = cad.get("gaussian_signed_per_mm2", cad["k1_per_mm"]*cad["k2_per_mm"])[ix]; curv_ok = cad["curvature_valid"][ix]
        method = "deterministic CAD graph node fallback (no face input)"
    else:
        sample_pos, source, bary = area_uniform_surface_samples(cad, count)
        tri = np.asarray(cad["face"], dtype=np.int64).T[source]
        normal = np.einsum("ni,nij->nj", bary, cad["normal"][tri]); normal /= np.maximum(np.linalg.norm(normal, axis=1, keepdims=True), 1e-12)
        turn = np.einsum("ni,ni->n", bary, cad["normal_turn_deg"][tri])
        vertex_mag = np.sqrt(cad["k1_per_mm"]**2+cad["k2_per_mm"]**2); vertex_gaussian = cad.get("gaussian_signed_per_mm2", cad["k1_per_mm"]*cad["k2_per_mm"])
        mag = np.einsum("ni,ni->n", bary, vertex_mag[tri]); gaussian = np.einsum("ni,ni->n", bary, vertex_gaussian[tri])
        curv_ok = np.all(cad["curvature_valid"][tri], axis=1); method = "deterministic area-uniform triangle sampling"
    mag = np.where(curv_ok, mag, 0.0); gaussian = np.where(curv_ok, gaussian, 0.0)
    x = finite32(np.column_stack([sample_pos, normal, turn, mag, gaussian, curv_ok, np.zeros((len(source), 3))]))
    valid = np.isfinite(x); valid[:, 7:9] = curv_ok[:, None]; valid[:, 9] = True; valid[:, 10:] = False
    result = {"pos_mm": finite32(sample_pos), "normal": finite32(normal), "x": x, "feature_valid": valid,
              "feature_names": CAD_FEATURE_NAMES, "feature_units": CAD_FEATURE_UNITS,
              "source_triangle": source, "training_valid": np.zeros(len(source), dtype=bool),
              "sampling_seed": np.asarray(20260914, dtype=np.int64), "sampling_method": np.asarray(method)}
    if bary is not None: result["barycentric"] = bary
    return result


def estimate_scan_features(points: np.ndarray, neighbors: int = 16) -> tuple[np.ndarray, np.ndarray]:
    """Local PCA descriptors in the original scan coordinate frame; normals are unoriented."""
    p = np.asarray(points, dtype=np.float64)
    if p.ndim != 2 or p.shape[1] != 3 or len(p) == 0:
        raise ValueError("scan points must be a non-empty [N,3] array")
    if not np.isfinite(p).all():
        raise ValueError("scan points contain non-finite coordinates")
    tree = cKDTree(p)
    k = min(max(1, int(neighbors)), len(p))
    d, nn = tree.query(p, k=k)
    if k == 1: d, nn = d[:,None], nn[:,None]
    normals = np.zeros_like(p); variation = np.zeros(len(p)); spacing = np.zeros(len(p))
    normal_valid = np.zeros(len(p), dtype=bool); spacing_valid = np.zeros(len(p), dtype=bool)
    for i in range(len(p)):
        q = p[nn[i]] - p[i]
        vals, vecs = np.linalg.eigh(q.T @ q / max(len(q), 1))
        vals = np.clip(vals, 0, None)
        if len(q) >= 2 and d[i,1] > 1e-12:
            spacing[i], spacing_valid[i] = d[i,1], True
        # A normal is not meaningful for a single point, duplicate cloud, or
        # collinear neighbourhood.  Its zero value is then explicitly masked.
        if len(q) >= 3 and np.linalg.matrix_rank(q, tol=1e-10) >= 2 and vals.sum() > 1e-12:
            normals[i] = vecs[:, 0]; variation[i] = vals[0] / vals.sum(); normal_valid[i] = True
    x = np.column_stack([p, normals, spacing, variation]).astype(np.float32)
    valid = np.zeros_like(x, dtype=bool); valid[:,:3] = True; valid[:,3:6] = normal_valid[:,None]
    valid[:,6] = spacing_valid; valid[:,7] = normal_valid
    return finite32(x), valid


def scan_points(cache: Path, run_id: str, count: int = 8192) -> dict[str, np.ndarray]:
    with np.load(cache, allow_pickle=False) as d:
        raw = np.asarray(d["points"], dtype=np.float64)
    ix = deterministic_sample_indices(len(raw), count, seed=20260914 + sum(map(ord, run_id)))
    x, valid = estimate_scan_features(raw[ix])
    return {"pos_native": finite32(raw[ix]), "x": x, "feature_valid": valid,
            "feature_names": SCAN_FEATURE_NAMES, "feature_units": np.asarray(["native"]*3+["unit"]*3+["native", "unit"]),
            "source_point": ix.astype(np.int32), "frame": np.asarray("raw_scan_native_unregistered"),
            "diagnostic_post_process": np.asarray(True), "training_valid": np.zeros(len(ix), dtype=bool),
            "sampling_seed": np.asarray(20260914 + sum(map(ord, run_id)), dtype=np.int64)}


def triangle_surface_samples(cad: dict[str, np.ndarray], pitch_mm: float) -> dict[str, np.ndarray]:
    """Deterministic interior triangle samples at approximately one per pitch².

    This prevents sparse voxel occupancy from being merely CAD graph vertex
    occupancy.  Every sample retains its source triangle and barycentric weights.
    """
    p = np.asarray(cad["pos_mm"], dtype=np.float64); face = np.asarray(cad["face"], dtype=np.int64).T
    points: list[np.ndarray] = []; triangles: list[np.ndarray] = []; barys: list[np.ndarray] = []
    for ti, ids in enumerate(face):
        q = p[ids]; n = max(1, int(np.ceil(np.max(np.linalg.norm(q[:,None]-q[None,:], axis=-1))/pitch_mm)))
        # Centroids of the regular barycentric subtriangles: no shared vertex-only samples.
        ii, jj = np.indices((n, n)); ii, jj = ii.ravel(), jj.ravel(); keep = ii + jj < n; ii, jj = ii[keep], jj[keep]
        b = np.column_stack([(ii+1/3)/n, (jj+1/3)/n, 1-(ii+jj+2/3)/n])
        points.append(b @ q); triangles.append(np.full(len(b), ti, np.int32)); barys.append(b)
    bary = np.concatenate(barys).astype(np.float32); tri_id = np.concatenate(triangles)
    return {"pos_mm": np.concatenate(points).astype(np.float32), "source_triangle": tri_id, "barycentric": bary}


def surface_voxels(cad: dict[str, np.ndarray], longest_axis_resolution: int,
                   voxel_pitch_mm: float | None = None,
                   grid_shape: tuple[int, int, int] | None = None) -> dict[str, np.ndarray]:
    p = np.asarray(cad["pos_mm"], dtype=np.float64)
    bbox_lo, bbox_hi = p.min(0), p.max(0)
    span = np.maximum(bbox_hi-bbox_lo, 1e-8)
    pitch = float(voxel_pitch_mm) if voxel_pitch_mm is not None else span.max()/int(longest_axis_resolution)
    if pitch <= 0: raise ValueError("voxel pitch must be positive")
    if grid_shape is None:
        dims = np.maximum(1, np.ceil(span/pitch).astype(np.int32))
        origin = bbox_lo
    else:
        dims = np.asarray(grid_shape, dtype=np.int32)
        if dims.shape != (3,) or np.any(dims < 1): raise ValueError("grid_shape must contain three positive dimensions")
        if np.any(span > dims*pitch + 1e-6): raise ValueError("fixed voxel grid does not contain the CAD")
        origin = (bbox_lo+bbox_hi-dims*pitch)/2
    samples = triangle_surface_samples(cad, pitch)
    sample_pos = samples["pos_mm"].astype(np.float64)
    # One physical pitch applies to x/y/z.  Short axes get fewer cells instead of
    # a distorted per-axis coordinate normalization.
    coord = np.floor((sample_pos-origin)/pitch).astype(np.int32)
    if np.any(coord < 0) or np.any(coord > dims): raise ValueError("surface sample lies outside the voxel grid")
    coord = np.minimum(coord, dims-1)
    linear = coord[:, 0] + dims[0]*(coord[:, 1] + dims[1]*coord[:, 2])
    unique, inv, count = np.unique(linear, return_inverse=True, return_counts=True)
    vox = np.column_stack([unique % dims[0], (unique//dims[0]) % dims[1], unique//(dims[0]*dims[1])]).astype(np.int32)
    tri = np.asarray(cad["face"], dtype=np.int64).T[samples["source_triangle"]]; bary = samples["barycentric"].astype(np.float64)
    sample_normal = np.einsum("ni,nij->nj", bary, cad["normal"][tri]); sample_normal /= np.maximum(np.linalg.norm(sample_normal, axis=1, keepdims=True), 1e-12)
    sample_turn = np.einsum("ni,ni->n", bary, cad["normal_turn_deg"][tri])
    vertex_mag = np.sqrt(cad["k1_per_mm"]**2+cad["k2_per_mm"]**2); vertex_gaussian = cad.get("gaussian_signed_per_mm2", cad["k1_per_mm"]*cad["k2_per_mm"])
    sample_valid = np.all(cad["curvature_valid"][tri], axis=1)
    sample_mag = np.where(sample_valid, np.einsum("ni,ni->n", bary, vertex_mag[tri]), 0.0)
    sample_gaussian = np.where(sample_valid, np.einsum("ni,ni->n", bary, vertex_gaussian[tri]), 0.0)
    nvox = len(unique); norm = np.zeros((nvox, 3)); turn = np.zeros(nvox); mag = np.zeros(nvox); gaussian = np.zeros(nvox); curv_count = np.zeros(nvox)
    np.add.at(norm, inv, sample_normal); np.add.at(turn, inv, sample_turn)
    np.add.at(mag, inv, sample_mag); np.add.at(gaussian, inv, sample_gaussian); np.add.at(curv_count, inv, sample_valid.astype(float))
    norm /= np.maximum(np.linalg.norm(norm, axis=1, keepdims=True), 1e-12); turn /= count
    curv_valid = curv_count > 0; mag[curv_valid] /= curv_count[curv_valid]; gaussian[curv_valid] /= curv_count[curv_valid]
    x = finite32(np.column_stack([np.ones(nvox), count, norm, turn, mag, gaussian, curv_valid, np.zeros((nvox, 3))]))
    valid = np.ones_like(x, dtype=bool); valid[:, 6:8] = curv_valid[:,None]; valid[:, 8] = True; valid[:, 9:] = False
    return {"voxel_coords": vox, "x": x, "feature_valid": valid,
            "feature_names": np.asarray(["surface_occupancy", "surface_sample_count", "normal_x", "normal_y", "normal_z", "normal_turn_deg", "curvature_magnitude_per_mm", "gaussian_curvature_per_mm2", "curvature_valid", "thickness_mm", "thermal_max_c", "stress_mpa"]),
            "grid_dims": dims, "bbox_min_mm": bbox_lo.astype(np.float32), "bbox_max_mm": bbox_hi.astype(np.float32),
            "grid_origin_mm": origin.astype(np.float32), "voxel_pitch_mm": np.asarray(pitch, dtype=np.float32),
            "grid_padding_mm": (dims*pitch-span).astype(np.float32),
            "surface_sample_source_triangle": samples["source_triangle"], "surface_sample_barycentric": samples["barycentric"], "surface_sample_voxel_index": inv.astype(np.int32),
            "longest_axis_resolution": np.asarray(longest_axis_resolution, dtype=np.int32),
            "surface_only": np.asarray(True), "solid_occupancy": np.asarray(False), "occupancy_method": np.asarray("deterministic triangle interior samples at voxel pitch"), "training_valid": np.zeros(nvox, dtype=bool)}


def dense_surface_tensor(sparse: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
    dims = tuple(map(int, sparse["grid_dims"])); channels = sparse["x"].shape[1]
    x = np.zeros((channels, *dims), dtype=np.float32); valid = np.zeros((channels, *dims), dtype=bool)
    c = sparse["voxel_coords"]
    x[:, c[:, 0], c[:, 1], c[:, 2]] = sparse["x"].T
    valid[:, c[:, 0], c[:, 1], c[:, 2]] = sparse["feature_valid"].T
    return {"x": x, "feature_valid": valid, "feature_names": sparse["feature_names"], "grid_dims": sparse["grid_dims"],
            "surface_only": np.asarray(True), "solid_occupancy": np.asarray(False)}


def process_tabular(process_path: Path, cad: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
    raw = json.loads(process_path.read_text(encoding="utf-8"))
    run_ids = sorted(raw)
    channels = sorted({(r["channel"], r.get("unit", "")) for run in raw.values() for r in run["rows"]})
    values = np.zeros((len(run_ids), len(channels)), dtype=np.float32); present = np.zeros_like(values, dtype=bool)
    lookup = {c: i for i, c in enumerate(channels)}
    for i, run in enumerate(run_ids):
        for r in raw[run]["rows"]:
            j = lookup[(r["channel"], r.get("unit", ""))]
            if r.get("value") is not None:
                values[i, j], present[i, j] = float(r["value"]), True
    p = cad["pos_mm"]; dims = p.max(0)-p.min(0)
    tri = cad["pos_mm"][cad["face"].T]; area = .5*np.linalg.norm(np.cross(tri[:,1]-tri[:,0], tri[:,2]-tri[:,0]), axis=1).sum()
    geo_names = np.asarray(["cad_bbox_x_mm", "cad_bbox_y_mm", "cad_bbox_z_mm", "cad_surface_area_mm2", "cad_node_count", "cad_triangle_count"])
    geo = np.asarray([*dims, area, len(p), cad["face"].shape[1]], dtype=np.float32)
    return {"run_id": np.asarray(run_ids), "process_x": values, "process_feature_valid": present,
            "process_feature_names": np.asarray([f"{n} [{u}]" for n,u in channels]), "geometry_x": np.repeat(geo[None], len(run_ids), 0),
            "geometry_feature_names": geo_names, "x": np.concatenate([values, np.repeat(geo[None], len(run_ids), 0)], 1),
            "feature_valid": np.concatenate([present, np.ones((len(run_ids),len(geo)), bool)], 1),
            "label_available": np.zeros(len(run_ids), dtype=bool), "training_valid": np.zeros(len(run_ids), dtype=bool),
            "split_metadata": np.asarray(json.dumps({"design_ids": 1, "runs": len(run_ids), "split": "none; one design/two runs; labels unavailable"}))}


def thermal_history_features(time_s: np.ndarray, fields: np.ndarray) -> dict[str, np.ndarray]:
    """Extract real history features only when aligned time and field samples are supplied.

    fields is [T,...]. cooling_rate_per_s is the most negative d(field)/dt,
    and temporal_gradient_max_per_s is max absolute d(field)/dt.  No thermal
    values are produced by this pipeline because input histories are absent.
    """
    t = np.asarray(time_s, dtype=float); a = np.asarray(fields, dtype=float)
    if (t.ndim != 1 or a.ndim < 1 or a.shape[0] != len(t) or len(t) < 2 or
            not np.isfinite(t).all() or not np.isfinite(a).all() or np.any(np.diff(t) <= 0)):
        raise ValueError("time_s must be finite/strictly increasing and fields must be finite [T,...], T>=2")
    rate = np.gradient(a, t, axis=0)
    min_derivative = np.min(rate, axis=0)
    return {"max": np.max(a, axis=0),
            "cooling_rate_per_s": np.minimum(min_derivative, 0.0),
            "max_cooling_rate_per_s": np.max(np.maximum(-rate, 0.0), axis=0),
            "minimum_temporal_derivative_per_s": min_derivative,
            "temporal_gradient_max_per_s": np.max(np.abs(rate), axis=0)}


def _schema(path: Path, arrays: dict[str, np.ndarray], provenance: dict[str, Any]) -> None:
    json_dump(path, {"arrays": {k: {"shape": list(np.asarray(v).shape), "dtype": str(np.asarray(v).dtype)} for k,v in arrays.items()}, "provenance": provenance})


def _save_npz(path: Path, arrays: dict[str, np.ndarray], provenance: dict[str, Any]) -> dict[str, Any]:
    path.parent.mkdir(parents=True, exist_ok=True); np.savez_compressed(path, **arrays)
    schema = path.with_suffix(".schema.json"); _schema(schema, arrays, provenance)
    return {"npz": str(path), "schema": str(schema), "arrays": {k: list(np.asarray(v).shape) for k,v in arrays.items()}}


def _plot_images(mv: dict[str, np.ndarray], path: Path) -> None:
    fig, ax = plt.subplots(2, 6, figsize=(15, 5)); names = ["+X", "-X", "+Y", "-Y", "+Z", "-Z"]
    good_curv = mv["curvature_magnitude_per_mm"][mv["curvature_valid"]]
    cmax = float(np.percentile(good_curv, 99)) if len(good_curv) else 1.0
    for i in range(6):
        d = np.ma.masked_where(~mv["visibility"][i], mv["depth_mm"][i]); c = np.ma.masked_where(~mv["curvature_valid"][i], mv["curvature_magnitude_per_mm"][i])
        ax[0,i].imshow(d, cmap="viridis"); ax[1,i].imshow(c, cmap="magma", vmin=0, vmax=max(cmax, 1e-12)); ax[0,i].set_title(names[i]); ax[0,i].axis("off"); ax[1,i].axis("off")
    ax[0,0].set_ylabel("depth (mm)"); ax[1,0].set_ylabel("|curvature| (1/mm)"); fig.tight_layout(); fig.savefig(path, dpi=150); plt.close(fig)


def _plot_channels(mv: dict[str, np.ndarray], path: Path) -> None:
    keys = [("depth_mm", "depth"), ("curvature_magnitude_per_mm", "|curvature|"), ("gaussian_curvature_per_mm2", "Gaussian K"), ("normal_orientation_to_view", "normal·view"), ("distance_to_boundary_px", "silhouette distance"), ("curvature_valid", "curvature valid")]
    fig, ax = plt.subplots(2, 3, figsize=(10, 7)); mask = mv["visibility"][0]
    for a,(k,title) in zip(ax.flat, keys):
        valid = mask & (mv["curvature_valid"][0] if k in {"curvature_magnitude_per_mm", "gaussian_curvature_per_mm2"} else True)
        values = mv[k][0]; shown = np.ma.masked_where(~valid, values)
        if k in {"curvature_magnitude_per_mm", "gaussian_curvature_per_mm2"} and np.any(valid):
            q = np.percentile(np.abs(values[valid]), 99); limits = {"vmin": -q, "vmax": q} if k.startswith("gaussian") else {"vmin": 0, "vmax": q}
        else: limits = {}
        a.imshow(shown, cmap="viridis", **limits); a.set_title(title); a.axis("off")
    fig.tight_layout(); fig.savefig(path, dpi=150); plt.close(fig)


def _plot_voxels(voxels: list[dict[str, np.ndarray]], path: Path) -> None:
    fig = plt.figure(figsize=(12,4))
    for i,v in enumerate(voxels, 1):
        a = fig.add_subplot(1,3,i, projection="3d"); c = v["voxel_coords"]
        a.scatter(c[:,0],c[:,1],c[:,2],s=1,c=v["x"][:,6],cmap="magma"); a.set_title(f"surface voxels {int(v['longest_axis_resolution'])}"); a.set_box_aspect(v["grid_dims"])
    fig.tight_layout(); fig.savefig(path, dpi=150); plt.close(fig)


def _plot_points(cadp: dict[str, np.ndarray], scans: list[dict[str, np.ndarray]], path: Path) -> None:
    sets = [(cadp["pos_mm"], "CAD (mm)")] + [(s["pos_native"], str(s["frame"])[:13]) for s in scans]
    fig = plt.figure(figsize=(4 * len(sets), 4))
    for i,(p,title) in enumerate(sets,1):
        a=fig.add_subplot(1,len(sets),i,projection="3d"); a.scatter(p[:,0],p[:,1],p[:,2],s=.5); a.set_title(title); a.set_box_aspect(np.maximum(np.ptp(p,axis=0),1e-6))
    fig.tight_layout(); fig.savefig(path,dpi=150); plt.close(fig)


def prepare(root: Path, resolutions: Iterable[int] = (128, 256), point_count: int = 8192) -> dict[str, list[dict[str, Any]]]:
    root = Path(root); source = root/"analysis/results"; out = root/"data_processed"; plots = root/"results/representations"
    plots.mkdir(parents=True, exist_ok=True)
    cad = load_cad_graph(source/"cad_graph.npz")
    provenance = {"cad_source": "analysis/results/cad_graph.npz", "units": "CAD millimetres", "mesh": "existing CAD graph triangles", "post_scan_leakage": "scans are separate diagnostic representations and training_valid=false"}
    index: dict[str, list[dict[str, Any]]] = {"images": [], "points": [], "voxels": [], "tabular": []}
    for res in resolutions:
        mv = make_multiview(cad, int(res)); dest = out/"images"/f"cad_multiview_{res}.npz"
        index["images"].append(_save_npz(dest, mv, provenance | {"resolution": int(res), "z_buffer": "largest object-to-camera depth"}))
        _plot_images(mv, plots/f"six_view_depth_curvature_{res}.png")
        if int(res) == min(resolutions): _plot_channels(mv, plots/"feature_channel_contact_sheet.png")
    cp = cad_points(cad, point_count); index["points"].append(_save_npz(out/"points"/"cad_surface_points.npz", cp, provenance | {"kind":"CAD-only pre-process"}))
    scans=[]
    for cache in sorted(source.glob("*_all_merge_pointcloud_full.npz")):
        run = cache.name.removesuffix("_all_merge_pointcloud_full.npz")
        sp = scan_points(cache, run, point_count); scans.append(sp)
        index["points"].append(_save_npz(out/"points"/f"scan_{run}_diagnostic.npz", sp, {"source":f"analysis/results/{run}_all_merge_pointcloud_full.npz", "frame":"raw cache native frame; no CAD transform applied", "diagnostic_post_process":True, "not_pre_process_predictor":True}))
    _plot_points(cp, scans, plots/"point_sample_comparison.png")
    voxels=[]
    for res in (32,64,128):
        sv=surface_voxels(cad,res); voxels.append(sv); index["voxels"].append(_save_npz(out/"voxels"/f"cad_surface_sparse_{res}.npz", sv, provenance | {"surface_occupancy":"not filled solid; CAD is nonwatertight", "longest_axis_resolution":res}))
        if res==32: index["voxels"].append(_save_npz(out/"voxels"/"cad_surface_dense_32.npz", dense_surface_tensor(sv), provenance | {"bounded_dense_tensor":True,"surface_occupancy":"not filled solid"}))
    _plot_voxels(voxels, plots/"sparse_voxel_resolutions.png")
    tab=process_tabular(source/"process.json",cad); index["tabular"].append(_save_npz(out/"tabular"/"process_and_geometry.npz",tab,{"source":"analysis/results/process.json", "scope":"pre-process vector + fixed CAD geometry summaries", "labels":"unavailable", "split":"one design, two runs; no generalization split"}))
    for family, entries in index.items(): json_dump(out/family/"index.json", {"family":family,"items":entries})
    json_dump(plots/"index.json", {"plots": [str(x) for x in sorted(plots.glob("*.png"))], "input": provenance, "limitations": ["CAD source mesh is nonwatertight; all voxels are surface-only.", "No thickness, thermal histories, stress, defect labels, or validated scan-to-CAD transform were available.", "Raw scans remain in their own native frames and are diagnostic post-process data only."]})
    return index


def main() -> None:
    ap=argparse.ArgumentParser(description="Prepare CAD-only and diagnostic-scan representations.")
    ap.add_argument("--root", type=Path, default=repo_root()); ap.add_argument("--resolutions", nargs="+", type=int, default=[128,256]); ap.add_argument("--point-count", type=int, default=8192)
    args=ap.parse_args(); result=prepare(args.root,args.resolutions,args.point_count)
    print(json.dumps({k:len(v) for k,v in result.items()}, ensure_ascii=False))


if __name__ == "__main__": main()