#!/usr/bin/env python3
"""Match STEP/CAD, process XLS data, measured PLY, and defect CSV by run.

The exported arrays are diagnostic supervision scaffolds.  The supplied CSVs
contain a point but no defect width/height, and scan-to-CAD registration is not
independently validated, so ``training_valid`` is deliberately false.
"""
from __future__ import annotations

import argparse
import csv
import html
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
import numpy as np
import open3d as o3d
from scipy.sparse import coo_matrix
from scipy.sparse.csgraph import dijkstra
from scipy.spatial import cKDTree

plt.rcParams["font.family"] = ["Apple SD Gothic Neo", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False

RUN_POSE = {"16h": "16h_mirror", "18h01": "18h01"}
WEAK_RADII_MM = np.asarray([5.0, 10.0, 20.0], np.float32)


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


def transform(points: np.ndarray, matrix: np.ndarray) -> np.ndarray:
    return points @ matrix[:3, :3].T + matrix[:3, 3]


def geodesic_distance(node_count: int, edge_index: np.ndarray,
                      edge_length: np.ndarray, seed: int) -> np.ndarray:
    graph = coo_matrix((edge_length, (edge_index[0], edge_index[1])),
                       shape=(node_count, node_count)).tocsr()
    return np.asarray(dijkstra(graph, directed=False, indices=seed), np.float32)


def read_defect_csv(path: Path) -> dict[str, object]:
    with path.open(encoding="utf-8-sig", newline="") as stream:
        rows = list(csv.DictReader(stream))
    if len(rows) != 1:
        raise ValueError(f"Expected one supplied defect row in {path}; got {len(rows)}")
    row = rows[0]
    return {
        "pixel_xy": [float(row["X"]), float(row["Y"])],
        "width_px": float(row["W"]) if row.get("W", "").strip() else None,
        "height_px": float(row["H"]) if row.get("H", "").strip() else None,
        "class": row["Error"],
    }


def closest_cad(mesh_path: Path, points: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    mesh = o3d.io.read_triangle_mesh(str(mesh_path))
    mesh.remove_duplicated_vertices(); mesh.compute_triangle_normals()
    scene = o3d.t.geometry.RaycastingScene()
    scene.add_triangles(o3d.t.geometry.TriangleMesh.from_legacy(mesh))
    hit = scene.compute_closest_points(o3d.core.Tensor(points.astype(np.float32)))
    closest = hit["points"].numpy()
    normals = hit["primitive_normals"].numpy()
    delta = points - closest
    return closest, np.linalg.norm(delta, axis=1).astype(np.float32), np.einsum(
        "ij,ij->i", delta, normals).astype(np.float32)


def heater_values(names: np.ndarray, values: np.ndarray,
                  valid: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    lookup = {str(name).split(" [", 1)[0]: i for i, name in enumerate(names)}
    top = np.full(17, np.nan, np.float32)
    bottom = np.full(17, np.nan, np.float32)
    for zone in range(1, 18):
        for prefix, target in (("상부온도", top), ("하부온도", bottom)):
            key = f"{prefix}{zone:02d}"
            if key in lookup and valid[lookup[key]]:
                target[zone - 1] = values[lookup[key]]
    return top, bottom


def interpolate_between_heaters(values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Fill only NaNs bracketed by valid heater zones; never extrapolate edges."""
    raw = np.asarray(values, np.float32)
    filled = raw.copy()
    observed = np.isfinite(raw)
    imputed = np.zeros(raw.shape, bool)
    valid_index = np.flatnonzero(observed)
    if len(valid_index) < 2:
        return filled, imputed
    candidates = np.flatnonzero(~observed)
    candidates = candidates[(candidates > valid_index[0]) & (candidates < valid_index[-1])]
    if len(candidates):
        filled[candidates] = np.interp(candidates, valid_index, raw[valid_index])
        imputed[candidates] = True
    return filled, imputed


def provisional_heater_projection(cad: np.ndarray, top: np.ndarray,
                                   bottom: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Uniform X-bin projection used only until physical heater footprints exist."""
    scale = float(np.ptp(cad[:, 0]))
    zone = np.floor((cad[:, 0] - cad[:, 0].min()) / max(scale, 1e-9) * 17).astype(np.int32)
    zone = np.clip(zone, 0, 16)
    return zone, top[zone], bottom[zone], (top - bottom)[zone]


def preview(path: Path, run: str, cad: np.ndarray, edges: np.ndarray,
            scan: np.ndarray, scan_abs: np.ndarray, weak: np.ndarray,
            anchor: np.ndarray, top: np.ndarray, bottom: np.ndarray,
            top_imputed: np.ndarray, bottom_imputed: np.ndarray,
            node_top: np.ndarray, node_bottom: np.ndarray, node_delta: np.ndarray,
            note: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fig, axes = plt.subplots(2, 3, figsize=(17, 10), constrained_layout=True)
    axes = axes.ravel()
    take = np.arange(0, edges.shape[1], 2)
    seg = np.stack([cad[edges[0, take]][:, [0, 2]], cad[edges[1, take]][:, [0, 2]]], 1)

    temp = np.stack([top, bottom])
    finite = temp[np.isfinite(temp)]
    image = axes[0].imshow(temp, aspect="auto", cmap="inferno",
                           vmin=float(finite.min()), vmax=float(finite.max()))
    axes[0].set_yticks([0, 1], ["upper / 상부", "lower / 하부"])
    axes[0].set_xticks(range(17), range(1, 18)); axes[0].set_xlabel("heater zone number")
    axes[0].set_title("2×17 sandwich heater map (* = linearly interpolated)")
    fig.colorbar(image, ax=axes[0], label="temperature (°C)")
    midpoint = float(np.nanmean(finite))
    for row in range(2):
        for col in range(17):
            value = temp[row, col]
            if np.isfinite(value):
                was_imputed = (top_imputed if row == 0 else bottom_imputed)[col]
                label = f"{value:.1f}*" if was_imputed else f"{value:.0f}"
                axes[0].text(col, row, label, ha="center", va="center", fontsize=6,
                             color="white" if value < midpoint else "black")
            else:
                axes[0].text(col, row, "NA", ha="center", va="center", fontsize=6, color="cyan")

    aoi = weak[:, 1] >= np.exp(-.5)
    common = dict(s=2.2, cmap="inferno", vmin=float(finite.min()), vmax=float(finite.max()))
    for ax, values, title in ((axes[1], node_top, "upper heater overlay"),
                              (axes[2], node_bottom, "lower heater overlay")):
        draw = ax.scatter(cad[:, 0], cad[:, 2], c=values, **common)
        ax.scatter(cad[aoi, 0], cad[aoi, 2], s=18, facecolors="none",
                   edgecolors="#00e5ff", linewidths=.7)
        ax.scatter(anchor[0], anchor[2], s=70, marker="x", c="#00e5ff")
        ax.set_title(title + "\ncyan = weak 10 mm AOI")
        fig.colorbar(draw, ax=ax, label="temperature (°C)")

    delta_finite = node_delta[np.isfinite(node_delta)]
    delta_lim = max(float(np.max(np.abs(delta_finite))) if len(delta_finite) else 1., 1.)
    draw = axes[3].scatter(cad[:, 0], cad[:, 2], c=node_delta, s=2.2, cmap="coolwarm",
                           vmin=-delta_lim, vmax=delta_lim)
    axes[3].scatter(cad[aoi, 0], cad[aoi, 2], s=18, facecolors="none",
                    edgecolors="#00e5ff", linewidths=.7)
    axes[3].scatter(anchor[0], anchor[2], s=70, marker="x", c="#111")
    axes[3].set_title("upper − lower heater overlay\nAOI shown in cyan")
    fig.colorbar(draw, ax=axes[3], label="temperature difference (°C)")

    show = np.linspace(0, len(scan)-1, min(15000, len(scan)), dtype=int)
    image = axes[4].scatter(scan[show, 0], scan[show, 2], c=scan_abs[show], s=.4,
                            cmap="viridis", vmin=0, vmax=float(np.percentile(scan_abs, 95)))
    fig.colorbar(image, ax=axes[4], label="nearest CAD surface distance (mm)")
    axes[4].set_title("registered measured PLY\ndeformation proxy, not tracked displacement")

    axes[5].add_collection(LineCollection(seg, linewidths=.12, colors="#55768b", alpha=.25))
    axes[5].scatter(cad[:, 0], cad[:, 2], c=weak[:, 1], s=2, cmap="magma", vmin=0, vmax=1)
    axes[5].scatter(anchor[0], anchor[2], s=60, marker="x", c="#00bcd4")
    axes[5].autoscale(); axes[5].set_title("weak 10 mm geodesic area label\nassumed radius, not observed extent")
    for ax in axes[1:]:
        ax.set_aspect("equal"); ax.set_xlabel("CAD X (mm)"); ax.set_ylabel("CAD Z (mm)")
    fig.suptitle(f"{run}: STEP + XLS heaters matched to PLY + defect CSV\n"
                 "CAD overlay assumes heater 01→17 uniformly ordered along CAD X; physical footprints are required\n"
                 + note)
    fig.savefig(path, dpi=150); plt.close(fig)


def prepare(root: Path, out: Path, result_dir: Path) -> dict[str, object]:
    graph_path = root / "data_processed/graphs/graph_12000.npz"
    table_path = root / "data_processed/tabular/process_and_geometry.npz"
    mesh_path = root / "analysis/results/design_fine.stl"
    meta = json.loads((root / "datas_2/viewer/explorer_data/metadata.json").read_text())
    annotations = json.loads((root / "datas_2/viewer/explorer_data/dent_annotations.json").read_text())
    with np.load(graph_path, allow_pickle=False) as raw:
        graph = {key: raw[key] for key in raw.files}
    with np.load(table_path, allow_pickle=False) as raw:
        table = {key: raw[key] for key in raw.files}

    out.mkdir(parents=True, exist_ok=True); visual = out / "visualizations"; visual.mkdir(exist_ok=True)
    result_dir.mkdir(parents=True, exist_ok=True)
    cad, normal, edges = graph["pos_mm"], graph["normal"], graph["edge_index"]
    length = graph["edge_attr"][:, 0]
    entries = []
    csv_dir = next(p for p in (root / "datas_2").glob("3. CSV File*") if p.is_dir())
    for run, pose in RUN_POSE.items():
        row = int(np.where(table["run_id"] == run)[0][0])
        pose_meta = meta["variants"][pose]
        matrix = np.asarray(pose_meta["transform"], np.float64)
        scan_path = root / f"analysis/results/{run}_all_merge_pointcloud_full.npz"
        with np.load(scan_path, allow_pickle=False) as raw:
            scan_raw = raw["points"].astype(np.float32)
        scan_cad = transform(scan_raw, matrix).astype(np.float32)
        scan_closest, scan_abs, scan_signed = closest_cad(mesh_path, scan_cad)

        defect_path = next(csv_dir.glob("*16*.csv" if run == "16h" else "*18*.csv"))
        defect = read_defect_csv(defect_path)
        anchor_raw = np.asarray(annotations[run]["raw_xyz_mm"], np.float32)
        anchor_cad = transform(anchor_raw[None], matrix)[0].astype(np.float32)
        anchor_node = int(cKDTree(cad).query(anchor_cad)[1])
        geo = geodesic_distance(len(cad), edges, length, anchor_node)
        weak_masks = geo[:, None] <= WEAK_RADII_MM[None]
        # A smooth weak target is useful for development, but its scale is an assumption.
        weak_probability = np.exp(-.5 * (geo[:, None] / WEAK_RADII_MM[None]) ** 2).astype(np.float32)
        heater_top_raw, heater_bottom_raw = heater_values(
            table["process_feature_names"], table["process_x"][row], table["process_feature_valid"][row])
        heater_top, heater_top_imputed = interpolate_between_heaters(heater_top_raw)
        heater_bottom, heater_bottom_imputed = interpolate_between_heaters(heater_bottom_raw)
        heater_zone, node_top, node_bottom, node_delta = provisional_heater_projection(
            cad, heater_top, heater_bottom)

        scan_tree = cKDTree(scan_cad)
        cad_scan_abs, cad_scan_index = scan_tree.query(cad, workers=2)
        cad_delta = scan_cad[cad_scan_index] - cad
        cad_scan_signed = np.einsum("ij,ij->i", cad_delta, normal).astype(np.float32)

        registration_valid = bool(pose_meta.get("independently_validated", False))
        extent_observed = defect["width_px"] is not None and defect["height_px"] is not None
        training_valid = np.zeros(len(cad), bool)
        status = (
            f"registration independently validated={registration_valid}; CSV defect extent "
            f"observed={extent_observed}; all targets diagnostic/weak and training_valid=false"
        )
        output = out / f"{run}_cad_ply_targets.npz"
        np.savez_compressed(
            output,
            run_id=np.asarray(run), cad_pos_mm=cad, cad_normal=normal,
            cad_edge_index=edges, process_x=table["process_x"][row],
            process_feature_valid=table["process_feature_valid"][row],
            process_feature_names=table["process_feature_names"],
            scan_pos_raw_mm=scan_raw, scan_pos_cad_mm=scan_cad,
            scan_closest_cad_mm=scan_closest, scan_surface_distance_mm=scan_abs,
            scan_normal_residual_mm=scan_signed,
            cad_to_scan_distance_mm=cad_scan_abs.astype(np.float32),
            cad_to_scan_normal_residual_mm=cad_scan_signed,
            defect_class=np.asarray(defect["class"]), defect_pixel_xy=np.asarray(defect["pixel_xy"], np.float32),
            defect_width_height_px=np.asarray([defect["width_px"] or np.nan, defect["height_px"] or np.nan], np.float32),
            defect_anchor_raw_mm=anchor_raw, defect_anchor_cad_mm=anchor_cad,
            defect_anchor_cad_node=np.asarray(anchor_node, np.int32),
            geodesic_distance_to_defect_seed_mm=geo,
            weak_area_radii_mm=WEAK_RADII_MM, weak_defect_mask=weak_masks,
            weak_defect_probability=weak_probability,
            heater_top_c=heater_top, heater_bottom_c=heater_bottom,
            heater_top_c_raw=heater_top_raw, heater_bottom_c_raw=heater_bottom_raw,
            heater_top_observed=np.isfinite(heater_top_raw), heater_bottom_observed=np.isfinite(heater_bottom_raw),
            heater_top_imputed=heater_top_imputed, heater_bottom_imputed=heater_bottom_imputed,
            heater_zone_index_zero_based=heater_zone,
            provisional_heater_top_c_per_cad_node=node_top,
            provisional_heater_bottom_c_per_cad_node=node_bottom,
            provisional_heater_delta_c_per_cad_node=node_delta,
            heater_projection_valid=np.asarray(False),
            transform_raw_ply_to_cad=matrix, transform_determinant=np.asarray(np.linalg.det(matrix[:3, :3])),
            registration_independently_validated=np.asarray(registration_valid),
            defect_extent_observed=np.asarray(extent_observed), training_valid=training_valid,
        )
        schema = {
            "sample_unit": "one physical production run",
            "model_input": ["CAD node/edge geometry derived from STEP", "XLS process values and validity masks"],
            "post_process_target_sources": ["fused all_merge_pointcloud_full PLY", "defect CSV"],
            "component_scans": "T1/T2/TOP are registration and coverage QA only",
            "arrays": {
                "scan_surface_distance_mm": "PLY point to nearest CAD surface; deformation proxy",
                "scan_normal_residual_mm": "oriented normal component of nearest-surface residual",
                "cad_to_scan_distance_mm": "CAD node to nearest fused scan point",
                "weak_defect_mask": "CAD-node geodesic dilation at weak_area_radii_mm",
                "weak_defect_probability": "Gaussian weak label around point-only CSV seed",
                "heater_top_c/heater_bottom_c": "upper/lower heater channels after internal linear interpolation",
                "heater_*_c_raw": "source values with NaN retained",
                "heater_*_imputed": "true only where an internal NaN was interpolated between valid zone numbers",
                "provisional_heater_*_per_cad_node": "uniform CAD-X projection for visualization only",
                "training_valid": "false until registration and spatial target extent are validated",
            },
        }
        output.with_suffix(".schema.json").write_text(
            json.dumps(schema, ensure_ascii=False, indent=2), encoding="utf-8")
        png = visual / f"{run}_cad_ply_targets.png"
        preview(png, run, cad, edges, scan_cad, scan_abs, weak_probability, anchor_cad,
                heater_top, heater_bottom, heater_top_imputed, heater_bottom_imputed,
                node_top, node_bottom, node_delta, status)
        entries.append({
            "run": run, "pose": pose, "npz": output.name, "preview": f"visualizations/{png.name}",
            "process_source": str(table_path.relative_to(root)), "cad_source": "datas_2/AC28_DT RR UPR STD.stp",
            "ply_source": str(scan_path.relative_to(root)), "defect_source": str(defect_path.relative_to(root)),
            "scan_points": len(scan_raw), "cad_nodes": len(cad), "defect_anchor_node": anchor_node,
            "registration_determinant": float(np.linalg.det(matrix[:3, :3])),
            "registration_independently_validated": registration_valid,
            "defect_extent_observed": extent_observed, "training_valid_nodes": int(training_valid.sum()),
            "heater_projection_valid": False,
            "heater_projection": "provisional uniform zone 01→17 ordering along CAD X",
            "heater_missing_value_method": "linear interpolation only for missing zones bracketed by valid neighboring zone numbers; no edge extrapolation",
            "primary_scan_role": "all_merge_pointcloud_full only: fused post-process surface outcome",
            "component_scan_role": "T1_DOWN/T1_UP/T2_LEFT/T2_RIGHT/TOP are registration, coverage, and sensor-consistency QA only",
            "status": status,
        })

    manifest = {
        "problem": "STEP design + XLS process conditions -> CAD-domain defect risk/shape segmentation and deformation field",
        "role_of_ply": "post-process supervision source during training/validation; not an inference input for pre-process prediction",
        "scan_selection": {
            "primary": "all_merge_pointcloud_full.ply (represented by the existing deterministic 2 mm cache)",
            "component_scans": "T1/T2/TOP are not independent samples or inference inputs; retain for pose/coverage QA",
        },
        "registration_policy": {
            "geometry_features": "match CAD and fused PLY geometry, then verify residuals by region and component sensor",
            "proper_rigid_requirement": "a physical rigid transform has determinant +1",
            "16h_status": "current best-looking mapping has determinant -1; treat only as coordinate-handedness/reflection diagnostic until scanner axes are confirmed",
        },
        "target_semantics": {
            "weak_defect_mask": "geodesic dilation of a CSV seed at assumed 5/10/20 mm radii; not observed defect extent",
            "scan_surface_distance_mm": "registered scan-to-nearest-CAD surface residual; not material-point displacement",
            "scan_normal_residual_mm": "normal component of nearest-surface residual; sign depends on oriented CAD tessellation",
            "provisional_heater_*_per_cad_node": "diagnostic uniform X-bin overlay; not a validated physical heater footprint",
        },
        "entries": entries,
        "blocking_validation": [
            "confirm CSV image origin and pixel-to-PLY calibration",
            "obtain physically valid independently verified PLY-to-STEP registration for every run",
            "collect defect polygons/masks or W/H instead of point-only coordinates",
            "collect many OK and NG runs/design revisions and split by run and CAD revision",
            "provide heater-zone coordinates/footprints and their relation to the STEP coordinate system",
        ],
    }
    (out / "index.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    cards = "".join(
        f'<section><h2>{html.escape(e["run"])}</h2><p>{html.escape(e["status"])}</p>'
        f'<p><a href="{e["npz"]}">raw NPZ</a> · <a href="index.json">metadata</a></p>'
        f'<img src="{e["preview"]}" alt="matched CAD PLY targets for {html.escape(e["run"])}"></section>'
        for e in entries)
    (out / "index.html").write_text(
        '<!doctype html><meta charset="utf-8"><title>CAD–PLY matched targets</title>'
        '<style>body{font:15px system-ui;max-width:1400px;margin:auto;padding:24px;background:#f7f8fa}'
        'section{background:white;padding:18px;margin:18px 0;border:1px solid #ddd}img{max-width:100%}</style>'
        '<h1>CAD–PLY matched supervision</h1><p>PLY is an outcome used to build labels. '
        'At prediction time the intended inputs are STEP geometry and XLS process conditions.</p>' + cards,
        encoding="utf-8")
    (result_dir / "cad_ply_target_summary.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    return manifest


def main() -> None:
    p = argparse.ArgumentParser(description="Build diagnostic CAD–PLY deformation and weak defect-area targets")
    p.add_argument("--root", type=Path, default=root_default())
    p.add_argument("--out", type=Path)
    p.add_argument("--results", type=Path)
    args = p.parse_args(); root = args.root.resolve()
    out = args.out or root / "data_processed/paired"
    results = args.results or root / "results/paired"
    manifest = prepare(root, out, results)
    print(json.dumps({"runs": len(manifest["entries"]), "out": str(out)}, ensure_ascii=False))


if __name__ == "__main__":
    main()
