"""Four compact conditional-diffusion diagnostics for prepared CAD/run data.

Each representation predicts CAD-correspondent signed normal displacement and a
weak defect-risk field.  Generated points are always reconstructed as
``p_generated = p_CAD + n_CAD * displacement``.  This is intentionally a
weak-label experiment, not a production training pipeline.
"""
from __future__ import annotations

import argparse, html, json, math
from pathlib import Path
from typing import Any

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial import cKDTree
import torch
from torch import Tensor, nn
import torch.nn.functional as F

from models.data import ROOT, fixed_physical_normalize, split_embedded_units
from models.training import seed_everything

# Keep a bundled, verified ASCII font for every plot. Korean content in the
# raster figures is transliterated (the HTML report preserves 찍힘) because the
# available TTC/OTF fonts caused missing tick and legend glyphs on this runtime.
matplotlib.rcParams.update({"font.family": "DejaVu Sans", "axes.unicode_minus": False})

EXP = ROOT / "gen_models"
INP, MOD, OUT, ATTR = (EXP / name for name in ("processed_inputs", "models", "outputs", "feature_attributions"))
SCALE_MM = 30.0
VARIANTS = ("image_2d", "voxel_3d", "point_6000", "graph_6000")


def tensor(x: np.ndarray, dtype: torch.dtype = torch.float32) -> Tensor:
    return torch.as_tensor(x, dtype=dtype)


def process_vector(item: dict[str, np.ndarray]) -> tuple[np.ndarray, list[str]]:
    raw = item["process_x"].astype(np.float32).copy(); observed = item["process_feature_valid"].astype(bool)
    names = item["process_feature_names"].astype(str); base, units = split_embedded_units(names); base = base.tolist()
    imputed = np.zeros(len(raw), bool)
    for prefix, values, flags in (("상부온도", item["heater_top_c"], item["heater_top_imputed"]), ("하부온도", item["heater_bottom_c"], item["heater_bottom_imputed"])):
        for zone in range(1, 18):
            name = f"{prefix}{zone:02d}"
            if name in base and flags[zone - 1]:
                raw[base.index(name)] = values[zone - 1]; imputed[base.index(name)] = True
    valid = observed | imputed
    value = fixed_physical_normalize(np.where(valid, raw, 0)[None], np.asarray(base), units)[0]
    return np.r_[value, observed.astype(np.float32), imputed.astype(np.float32)], base + [f"{n}:observed" for n in base] + [f"{n}:imputed" for n in base]


def standardize(x: np.ndarray, mask: np.ndarray) -> np.ndarray:
    """Representation-only fixed summary scaling; no outcome information is fit."""
    value = np.where(mask, np.nan_to_num(x, nan=0., posinf=0., neginf=0.), 0.).astype(np.float32)
    scale = np.maximum(np.median(np.abs(value), axis=0), 1e-3)
    return np.clip(value / scale, -10, 10).astype(np.float32)


def knn_neighbors(coordinates: np.ndarray, count: int) -> np.ndarray:
    """Deterministic local message neighborhoods; first query is self."""
    _, index = cKDTree(coordinates).query(coordinates, k=min(count + 1, len(coordinates)), workers=1)
    index = np.atleast_2d(index)
    result = index[:, 1:]
    if result.shape[1] < count:
        result = np.c_[result, np.repeat(np.arange(len(coordinates))[:, None], count - result.shape[1], axis=1)]
    return result.astype(np.int64)


def load_representation(name: str) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, dict[str, Any]]:
    if name == "image_2d":
        with np.load(ROOT / "data_processed/images/cad_multiview_128.npz", allow_pickle=False) as a:
            keep = (a["silhouette"] & a["visibility"]).ravel(); ids = np.flatnonzero(keep)
            ids = ids[np.linspace(0, len(ids) - 1, min(6000, len(ids)), dtype=np.int64)]
            pos = a["xyz_mm"].reshape(-1, 3)[ids]; normal = a["normal"].reshape(-1, 3)[ids]
            x = a["x"].transpose(0, 2, 3, 1).reshape(-1, a["x"].shape[1])[ids]
            mask = a["feature_valid"].transpose(0, 2, 3, 1).reshape(-1, a["x"].shape[1])[ids]
        # Explicit 2-D neighbourhoods within each view (view,y,x); the large
        # view offset prevents a message crossing between different cameras.
        image_coordinates = np.c_[ids // (128 * 128) * 1000, (ids % (128 * 128)) // 128, ids % 128]
        return pos, normal, standardize(x, mask), mask, knn_neighbors(image_coordinates, 4), {"source": "data_processed/images/cad_multiview_128.npz", "representation": "six CAD orthographic views with explicit 2-D pixel-neighbor message aggregation", "nodes": len(pos), "denoiser_adapter": "4 nearest selected image-grid pixels in the same view; not a CNN"}
    if name == "voxel_3d":
        with np.load(ROOT / "data_processed/voxels/cad_surface_sparse_64.npz", allow_pickle=False) as a:
            coords = a["voxel_coords"].astype(np.float32); pos = a["grid_origin_mm"] + coords * float(a["voxel_pitch_mm"])
            x, mask = a["x"], a["feature_valid"]
            # Voxel records do not store normals; nearest CAD normals are filled below.
        return pos, np.empty_like(pos), standardize(x, mask), mask, knn_neighbors(coords, 6), {"source": "data_processed/voxels/cad_surface_sparse_64.npz", "representation": "64-grid sparse CAD surface voxels with 3-D voxel-neighbor message aggregation", "nodes": len(pos), "denoiser_adapter": "6 nearest sparse voxel-grid neighbours; not a dense 3-D CNN"}
    if name == "point_6000":
        with np.load(ROOT / "data_processed/points/cad_surface_points.npz", allow_pickle=False) as a:
            ids = np.linspace(0, len(a["pos_mm"]) - 1, 6000, dtype=np.int64)
            pos, normal, x, mask = a["pos_mm"][ids], a["normal"][ids], a["x"][ids], a["feature_valid"][ids]
        return pos, normal, standardize(x, mask), mask, knn_neighbors(pos, 8), {"source": "data_processed/points/cad_surface_points.npz", "representation": "deterministic 6,000 CAD surface points with PointNet-style local/global pooling", "nodes": len(pos), "denoiser_adapter": "8-nearest local pooling plus per-run global max/mean context"}
    if name == "graph_6000":
        with np.load(ROOT / "data_processed/graphs/graph_6000.npz", allow_pickle=False) as a:
            pos, normal, x, mask, edge = a["pos_mm"], a["normal"], a["x"], a["feature_valid"], a["edge_index"]
            # One topology-aware adapter pass: append each node's mean adjacent
            # CAD feature. The DDPM is still deliberately small, but its graph
            # conditioning changes if the prepared topology changes.
            summed = np.zeros_like(x); degree = np.zeros((len(x), 1), np.float32)
            np.add.at(summed, edge[0], x[edge[1]]); np.add.at(degree[:, 0], edge[0], 1)
            x = np.c_[x, summed / np.maximum(degree, 1)]; mask = np.c_[mask, mask]
            adjacency = [[] for _ in range(len(x))]
            for source, destination in edge.T: adjacency[int(source)].append(int(destination))
            neighbours = np.asarray([(row[:8] + [index] * 8)[:8] for index, row in enumerate(adjacency)], np.int64)
        return pos, normal, standardize(x, mask), mask, neighbours, {"source": "data_processed/graphs/graph_6000.npz", "representation": f"CAD graph nodes ({len(pos):,} nodes) with edge-index message passing at each denoise step", "nodes": len(pos), "denoiser_adapter": "prepared graph edge-index, up to 8 adjacent CAD nodes per message"}
    raise ValueError(name)


def paired_records() -> list[tuple[str, dict[str, np.ndarray]]]:
    records = []
    for path in sorted((ROOT / "data_processed/paired").glob("*_cad_ply_targets.npz")):
        with np.load(path, allow_pickle=False) as a: item = {key: a[key] for key in a.files}
        records.append((str(item["run_id"].item()), item))
    if not records: raise RuntimeError("Four-representation diagnostic requires at least one paired run")
    return records


def prepare_variant(name: str) -> dict[str, Any]:
    """Map each prepared representation onto the paired CAD target only by geometry."""
    INP.mkdir(parents=True, exist_ok=True)
    pos, normal, features, mask, neighbours, info = load_representation(name); records = paired_records()
    cad_pos = records[0][1]["cad_pos_mm"]; cad_normal = records[0][1]["cad_normal"]
    _, nearest = cKDTree(cad_pos).query(pos, workers=1)
    if not np.isfinite(normal).all() or not np.any(normal): normal = cad_normal[nearest]
    normal /= np.maximum(np.linalg.norm(normal, axis=1, keepdims=True), 1e-6)
    all_condition=[]; all_target=[]; all_risk=[]; all_binary=[]; all_local=[]; process_names: list[str] = []
    for run, item in records:
        proc, process_names = process_vector(item)
        valid = bool(item["heater_projection_valid"].item())
        # Values are deliberately zeroed when the projection validity is false;
        # the mask is present explicitly for a future physically mapped zone field.
        local = np.c_[item["provisional_heater_top_c_per_cad_node"][nearest], item["provisional_heater_bottom_c_per_cad_node"][nearest], item["provisional_heater_delta_c_per_cad_node"][nearest]] / 500.
        local_mask = np.full((len(pos), 1), float(valid), np.float32); local *= local_mask
        condition = np.c_[pos / 1000., normal, local, local_mask, np.repeat(proc[None], len(pos), axis=0), features, mask.astype(np.float32)]
        all_condition.append(condition.astype(np.float32)); all_target.append(np.clip(item["cad_to_scan_normal_residual_mm"][nearest], -SCALE_MM, SCALE_MM) / SCALE_MM)
        all_risk.append(item["weak_defect_probability"][nearest, 1]); all_binary.append(item["weak_defect_mask"][nearest, 1]); all_local.append(valid)
    output = {"run_id": np.asarray([r for r,_ in records]), "pos_mm": pos, "normal": normal, "nearest_paired_cad_node": nearest.astype(np.int32), "condition": np.concatenate(all_condition), "signed_target_scaled": np.concatenate(all_target).astype(np.float32), "risk_target": np.concatenate(all_risk).astype(np.float32), "risk_binary": np.concatenate(all_binary).astype(bool), "nodes_per_run": np.int32(len(pos)), "neighbor_index_one_run": neighbours, "process_feature_names": np.asarray(process_names), "heater_projection_valid": np.asarray(all_local, bool), "representation_features": features, "representation_feature_valid": mask, "run_defect_class": np.asarray([str(item["defect_class"].item()) for _,item in records])}
    output.update({f"scan_points_{run}_cad_mm": item["scan_pos_cad_mm"] for run,item in records})
    np.savez_compressed(INP / f"{name}.npz", **output)
    (INP / f"{name}.json").write_text(json.dumps({**info, "condition_dim": int(output["condition"].shape[1]), "nearest_paired_cad_mapping": "cKDTree CAD position only", "heater_projection_valid": all_local, "label_status": "weak/unvalidated diagnostic"}, ensure_ascii=False, indent=2), encoding="utf-8")
    return {key: value for key, value in output.items()} | {"info": info}


class TopologyConditionalDDPM(nn.Module):
    """Denoiser with a representation-specific local message every DDPM step.

    ``point_6000`` additionally follows the PointNet pattern: local kNN pooling
    plus global mean/max pooling for each process run. The other three variants
    use their own 2-D, 3-D or CAD-graph neighbourhood index.
    """
    def __init__(self, condition_dim: int, nodes_per_run: int, neighbours: np.ndarray, variant: str, width: int = 96, run_count: int = 2) -> None:
        super().__init__(); self.nodes_per_run=nodes_per_run; self.variant=variant
        full = np.concatenate([neighbours + i*nodes_per_run for i in range(run_count)], axis=0)
        self.register_buffer("neighbours", torch.as_tensor(full, dtype=torch.long))
        self.c=nn.Sequential(nn.Linear(condition_dim,width),nn.SiLU(),nn.Linear(width,width),nn.SiLU()); self.t=nn.Sequential(nn.Linear(16,width),nn.SiLU(),nn.Linear(width,width))
        message_width = width * (3 if variant == "point_6000" else 1)
        self.e=nn.Sequential(nn.Linear(1 + width + message_width + width,width),nn.SiLU(),nn.Linear(width,1)); self.r=nn.Sequential(nn.Linear(width + message_width,width//2),nn.SiLU(),nn.Linear(width//2,1))
    def forward(self, x: Tensor, t: Tensor, c: Tensor) -> tuple[Tensor, Tensor]:
        z=self.c(c); local=z[self.neighbours].mean(1)
        if self.variant == "point_6000":
            groups=z.reshape(-1,self.nodes_per_run,z.shape[-1]); mean=groups.mean(1,keepdim=True).expand_as(groups).reshape_as(z); maximum=groups.max(1,keepdim=True).values.expand_as(groups).reshape_as(z); message=torch.cat([local,mean,maximum],1)
        else: message=local
        f=torch.exp(torch.linspace(0, math.log(1000), 8, dtype=x.dtype)); p=t[:,None].to(x.dtype)*f[None]/30.; h=self.t(torch.cat([p.sin(),p.cos()],1)); return self.e(torch.cat([x[:,None],z,message,h],1)).squeeze(1), self.r(torch.cat([z,message],1)).squeeze(1)


def schedule(steps: int) -> Tensor: return torch.cumprod(1-torch.linspace(1e-4,.02,steps),0)
def average_precision(score: np.ndarray, y: np.ndarray) -> float | None:
    if not y.any(): return None
    order=np.argsort(-score,kind="stable"); y=y[order].astype(float); return float((np.cumsum(y)/np.arange(1,len(y)+1)*y).sum()/y.sum())


def ascii_process_label(index: int) -> str:
    """Portable raster label; exact Korean source names remain in NPZ/JSON."""
    return f"process feature {index + 1}"
def chamfer(a: np.ndarray,b: np.ndarray) -> float:
    ix=np.linspace(0,len(a)-1,min(768,len(a)),dtype=int); iy=np.linspace(0,len(b)-1,min(768,len(b)),dtype=int); d=torch.cdist(tensor(a[ix])[None],tensor(b[iy])[None]).squeeze(); return float((d.min(1).values.mean()+d.min(0).values.mean()).item()/2)


@torch.no_grad()
def sample(model: TopologyConditionalDDPM, condition: Tensor, abar: Tensor, sample_steps: int, seed: int) -> Tensor:
    x=torch.randn(len(condition),generator=torch.Generator().manual_seed(seed)); seq=np.unique(np.linspace(0,len(abar)-1,sample_steps,dtype=int))[::-1]
    for i,k in enumerate(seq):
        eps,_=model(x,torch.full((len(x),),int(k),dtype=torch.long),condition); x0=(x-(1-abar[k]).sqrt()*eps)/abar[k].sqrt(); previous=-1 if i==len(seq)-1 else int(seq[i+1]); x=x0 if previous<0 else abar[previous].sqrt()*x0+(1-abar[previous]).sqrt()*eps
    return x.clamp(-1,1)


def train_variant(name: str, data: dict[str, Any], epochs: int, lr: float, diffusion_steps: int, samples: int, sample_steps: int) -> dict[str, Any]:
    condition=tensor(data["condition"]); target=tensor(data["signed_target_scaled"]); risk=tensor(data["risk_target"]); binary=data["risk_binary"].astype(bool); n=int(data["nodes_per_run"]); runs=len(data["run_id"]); abar=schedule(diffusion_steps); model=TopologyConditionalDDPM(condition.shape[1],n,data["neighbor_index_one_run"],name,run_count=runs); opt=torch.optim.AdamW(model.parameters(),lr=lr,weight_decay=1e-5); pw=((len(risk)-risk.sum())/risk.sum().clamp_min(1)).clamp(max=200); hist=[]; best=(float("inf"),0,None)
    for epoch in range(epochs):
        opt.zero_grad(); t=torch.randint(0,diffusion_steps,(len(target),)); noise=torch.randn_like(target); xt=abar[t].sqrt()*target+(1-abar[t]).sqrt()*noise; eps,logit=model(xt,t,condition); denoise=F.mse_loss(eps,noise); rloss=F.binary_cross_entropy_with_logits(logit,risk,pos_weight=pw); loss=denoise+.35*rloss; loss.backward();opt.step(); score=torch.sigmoid(logit).detach().numpy(); hist.append({"epoch":epoch+1,"total":float(loss.detach()),"denoise":float(denoise.detach()),"risk":float(rloss.detach()),"AP":average_precision(score,binary) or 0.});
        loss_value = float(loss.detach())
        if loss_value < best[0]: best=(loss_value,epoch+1,{k:v.detach().clone() for k,v in model.state_dict().items()})
    model.load_state_dict(best[2]); model.eval(); generated=torch.stack([sample(model,condition,abar,sample_steps,20260915+i) for i in range(samples)]).numpy()
    with torch.no_grad(): _, logit=model(torch.zeros_like(target),torch.zeros(len(target),dtype=torch.long),condition)
    score=torch.sigmoid(logit).numpy(); pos,normal=data["pos_mm"],data["normal"]; result=[]
    for j,run in enumerate(data["run_id"]):
        sl=slice(j*n,(j+1)*n); mean=generated[:,sl].mean(0)*SCALE_MM; std=generated[:,sl].std(0)*SCALE_MM; tgt=target.numpy()[sl]*SCALE_MM; pred=pos+normal*mean[:,None]; ref=pos+normal*tgt[:,None]; y=binary[sl]; p=score[sl]>=.5; both=int((p&y).sum()); risk_metrics={"AP":average_precision(score[sl],y),"IoU_at_0.5":both/max(1,int((p|y).sum())),"Dice_at_0.5":2*both/max(1,int(p.sum()+y.sum()))}; scan=data[f"scan_points_{run}_cad_mm"]; result.append({"run":str(run),"MAE_mm":float(abs(mean-tgt).mean()),"RMSE_mm":float(np.sqrt(np.mean((mean-tgt)**2))),"CAD_correspondent_Chamfer_mm":chamfer(pred,ref),"registered_scan_Chamfer_mm":chamfer(pred,scan),"risk":risk_metrics,"mean_stochastic_std_mm":float(std.mean()),"max_stochastic_std_mm":float(std.max())})
    MOD.mkdir(parents=True,exist_ok=True); torch.save({"state_dict":model.state_dict(),"condition_dim":condition.shape[1],"nodes_per_run":n,"production_valid":False,"variant":name,"denoiser_adapter":data["info"]["denoiser_adapter"]},MOD/f"{name}.pt")
    np.savez_compressed(OUT/f"{name}_generated.npz",run_id=data["run_id"],cad_pos_mm=pos,cad_normal=normal,generated_signed_normal_mm=generated.reshape(samples,runs,n)*SCALE_MM,generated_pointcloud_mm=(pos[None,None]+normal[None,None]*generated.reshape(samples,runs,n,1)*SCALE_MM),risk_probability=score.reshape(runs,n),defect_type=data["run_defect_class"],defect_type_status=np.asarray("constant_diagnostic_unlearnable_single_class"))
    return {"variant":name,"representation":data["info"],"nodes_per_run":n,"condition_dim":int(condition.shape[1]),"history":hist,"best_loss":best[0],"best_epoch":best[1],"metrics":result,"heater_projection_valid":data["heater_projection_valid"].tolist(),"checkpoint":str((MOD/f"{name}.pt").relative_to(ROOT))}


def attribution(name: str, data: dict[str,Any], checkpoint: Path) -> None:
    """Fixed-state risk and denoised-deformation gradient diagnostics."""
    base_condition=tensor(data["condition"]); model=TopologyConditionalDDPM(base_condition.shape[1],int(data["nodes_per_run"]),data["neighbor_index_one_run"],name,run_count=len(data["run_id"])); model.load_state_dict(torch.load(checkpoint,weights_only=True)["state_dict"]); model.eval()
    names=data["process_feature_names"]; start=10
    condition=base_condition.clone().requires_grad_(True); zero=torch.zeros(len(condition)); t0=torch.zeros(len(condition),dtype=torch.long); _,logit=model(zero,t0,condition);logit.mean().backward();risk_value=condition.grad[:,start:start+len(names)].abs().mean(0).detach().numpy()
    condition=base_condition.clone().requires_grad_(True); step=15; timestep=torch.full((len(condition),),step,dtype=torch.long);epsilon,_=model(zero,timestep,condition);abar=schedule(30)[step];x0=-(1-abar).sqrt()*epsilon/abar.sqrt();x0.abs().mean().backward();deformation_value=condition.grad[:,start:start+len(names)].abs().mean(0).detach().numpy()
    fig,axes=plt.subplots(1,2,figsize=(15,4.5),constrained_layout=True)
    for ax,value,title,color in ((axes[0],risk_value,"weak-risk logit sensitivity","#495f9d"),(axes[1],deformation_value,"denoised deformation sensitivity","#a85b35")):
        top=np.argsort(value)[-12:][::-1];labels=[ascii_process_label(int(index)) for index in top][::-1]
        ax.barh(labels,value[top][::-1],color=color,label="mean absolute gradient");ax.set_xlabel("mean absolute local gradient");ax.set_ylabel("XLS process feature index");ax.set_title(title);ax.legend(loc="lower right",fontsize=8);ax.grid(axis="x",alpha=.22)
    fig.suptitle(f"{name}: fixed-state local sensitivity diagnostic (not causal attribution)");fig.savefig(ATTR/f"{name}_process_gradient.png",dpi=150);plt.close(fig)
    risk_order = np.argsort(risk_value)[::-1]
    deformation_order = np.argsort(deformation_value)[::-1]
    payload = {
        "variant": name,
        "method": "absolute local gradients with fixed diffusion state; risk uses mean logit at t=0, deformation uses mean |x0 estimate| at t=15 of 30 with x_t=0",
        "production_valid": False,
        "causal": False,
        "checkpoint": str(checkpoint.relative_to(ROOT)),
        "risk_features": [{"rank": rank + 1, "feature_index_one_based": int(index) + 1,"source_name": str(names[index]), "mean_absolute_gradient": float(risk_value[index])} for rank,index in enumerate(risk_order)],
        "deformation_features": [{"rank": rank + 1, "feature_index_one_based": int(index) + 1,"source_name": str(names[index]), "mean_absolute_gradient": float(deformation_value[index])} for rank,index in enumerate(deformation_order)],
    }
    (ATTR / f"{name}_process_gradient.json").write_text(
        json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")


def processed_input_figure(name: str, data: dict[str, Any]) -> None:
    """Show the real prepared representation and the two condition vectors."""
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2), constrained_layout=True)
    ax = axes[0]
    if name == "image_2d":
        with np.load(ROOT / "data_processed/images/cad_multiview_128.npz", allow_pickle=False) as source:
            image = np.where(source["silhouette"][0], source["depth_mm"][0], np.nan)
        shown = ax.imshow(image, cmap="viridis"); fig.colorbar(shown, ax=ax, label="CAD depth (mm)")
        ax.set_title("real 2-D input: CAD view 0 depth pixels"); ax.set_xlabel("pixel x"); ax.set_ylabel("pixel y")
    elif name == "voxel_3d":
        fig.delaxes(ax); ax=fig.add_subplot(1,2,1,projection="3d")
        pos=data["pos_mm"];feature=np.linalg.norm(data["representation_features"],axis=1)
        draw=ax.scatter(pos[:,0],pos[:,1],pos[:,2],c=feature,s=2,cmap="viridis",rasterized=True)
        fig.colorbar(draw,ax=ax,label="norm of prepared representation features",shrink=.72)
        ax.set_title("real 3-D input: sparse CAD surface voxels");ax.set_xlabel("CAD X (mm)");ax.set_ylabel("CAD Y (mm)");ax.set_zlabel("CAD Z (mm)");ax.view_init(24,-62)
    else:
        pos = data["pos_mm"]; feature = np.linalg.norm(data["representation_features"], axis=1)
        draw = ax.scatter(pos[:, 0], pos[:, 2], c=feature, s=2, cmap="viridis", rasterized=True)
        if name == "graph_6000":
            neighbours = data["neighbor_index_one_run"]
            for node in np.linspace(0, len(pos) - 1, min(550, len(pos)), dtype=int):
                for neighbour in neighbours[node, :2]:
                    ax.plot(pos[[node, neighbour], 0], pos[[node, neighbour], 2], color="#607d8b", lw=.18, alpha=.30)
        fig.colorbar(draw, ax=ax, label="norm of prepared representation features")
        ax.set_title({"voxel_3d":"real 3-D input: sparse CAD surface voxels", "point_6000":"real point input: 6,000 CAD surface points", "graph_6000":f"real graph input: {len(pos):,} CAD nodes + edges"}[name])
        ax.set_xlabel("CAD X (mm)"); ax.set_ylabel("CAD Z (mm)"); ax.set_aspect("equal")
    process_width = len(data["process_feature_names"]) // 3
    for index,run in enumerate(data["run_id"]):
        values=data["condition"][index*int(data["nodes_per_run"]),10:10+process_width]
        axes[1].plot(np.arange(1,process_width+1),values,marker="o",ms=2.7,label=f"{run} XLS process")
    axes[1].set_title("condition actually supplied to every CAD element")
    axes[1].set_xlabel("process feature index (source names retained in NPZ)"); axes[1].set_ylabel("fixed physical normalized value")
    axes[1].legend(); axes[1].grid(alpha=.25)
    axes[1].text(.01,.02,"local provisional heater values: zero-imputed\nlocal heater validity mask: false for these runs",transform=axes[1].transAxes,fontsize=8,bbox={"facecolor":"#fff2cc","edgecolor":"#d99b00","alpha":.9})
    fig.suptitle(f"{name}: prepared input and XLS/heater condition summary")
    fig.savefig(INP / f"{name}.png", dpi=170); plt.close(fig)


def figures(reports: list[dict[str,Any]], run: str) -> None:
    OUT.mkdir(parents=True,exist_ok=True); fig,axes=plt.subplots(2,2,figsize=(12,7),constrained_layout=True)
    for ax,rep in zip(axes.ravel(),reports):
        h=rep["history"];ax.plot([x["epoch"] for x in h],[x["denoise"] for x in h],label="denoise MSE");ax.plot([x["epoch"] for x in h],[x["risk"] for x in h],label="risk BCE");ax.set_yscale("log");ax.set_title(rep["variant"]);ax.grid(alpha=.2);ax.legend(fontsize=8)
    fig.suptitle("Four conditional diffusion denoising losses (weak-label diagnostic)");fig.savefig(OUT/"loss_comparison.png",dpi=160);plt.close(fig)
    # Fixed last-run CAD X-Z camera and 0..30 mm colour scale for every variant.
    fig,axes=plt.subplots(2,4,figsize=(18,8),constrained_layout=True)
    for col,rep in enumerate(reports):
        a=np.load(OUT/f"{rep['variant']}_generated.npz");pos=a["cad_pos_mm"];normal=a["cad_normal"];disp=a["generated_signed_normal_mm"][0,-1];cloud=pos+normal*disp[:,None]; risk=a["risk_probability"][-1]
        for row,(values,title,cmap,hi) in enumerate(((abs(disp),"generated deformation |mm|","viridis",30),(risk,"weak risk / unvalidated run-level type","magma",1))):
            ax=axes[row,col];d=ax.scatter(cloud[:,0],cloud[:,2],c=values,s=2,cmap=cmap,vmin=0,vmax=hi,rasterized=True);ax.set_aspect("equal");ax.set_title(f"{rep['variant']}\n{title}");ax.set_xlabel("CAD X (mm)");ax.set_ylabel("CAD Z (mm)");fig.colorbar(d,ax=ax,shrink=.7)
    fig.suptitle(f"{run} fixed CAD projection and colour scales; generated sample 0") ;fig.savefig(OUT/f"{run}_four_variant_comparison.png",dpi=170);plt.close(fig)
    fig,ax=plt.subplots(figsize=(12,3.6));ax.axis("off");labels=["2D CAD views\n4 pixel messages", "3D sparse voxels\n6 voxel messages", "6k CAD points\nkNN + global pooling", "~6k CAD graph\nedge messages"];x=np.linspace(.08,.72,4)
    for xx,label in zip(x,labels): ax.add_patch(plt.Rectangle((xx,.57),.18,.2,fc="#e8f0fe",ec="#2855a6"));ax.text(xx+.09,.67,label,ha="center",va="center");ax.annotate("",(.49,.46),(xx+.09,.57),arrowprops={"arrowstyle":"->"})
    ax.add_patch(plt.Rectangle((.40,.25),.2,.2,fc="#eaf6ec",ec="#278343"));ax.text(.5,.35,"process_x + masks\nprovisional heaters + mask",ha="center",va="center");ax.add_patch(plt.Rectangle((.72,.25),.2,.2,fc="#fff1e6",ec="#ae5f15"));ax.text(.82,.35,"DDPM signed d\n→ CAD + normal × d\nrisk overlay",ha="center",va="center");ax.annotate("",(.72,.35),(.60,.35),arrowprops={"arrowstyle":"->"});ax.text(.5,.04,"All local heater projections are mask=false (heater_projection_valid=false); values are zero-imputed.",ha="center",color="#8a4b00")
    fig.savefig(OUT/"architecture.png",dpi=170,bbox_inches="tight");plt.close(fig)


def gallery(report: dict[str,Any]) -> None:
    run=report["runs"][-1]
    rows="".join(f"<tr><td>{r['variant']}</td><td>{r['nodes_per_run']}</td><td>{r['best_loss']:.4f}</td><td>{r['metrics'][-1]['MAE_mm']:.3f}</td><td>{r['metrics'][-1]['RMSE_mm']:.3f}</td><td>{r['metrics'][-1]['CAD_correspondent_Chamfer_mm']:.3f}</td><td>{r['metrics'][-1]['registered_scan_Chamfer_mm']:.3f}</td><td>{r['metrics'][-1]['risk']['AP']}</td></tr>" for r in report["variants"])
    inputs=" · ".join(f"<a href=../processed_inputs/{r['variant']}.png>{r['variant']} prepared input</a>" for r in report["variants"])
    attributions=" · ".join(f"<a href=../feature_attributions/{r['variant']}_process_gradient.png>{r['variant']} attribution</a> (<a href=../feature_attributions/{r['variant']}_process_gradient.json>JSON</a>)" for r in report["variants"])
    (OUT/"index.html").write_text(f"""<!doctype html><meta charset=utf-8><title>CAD conditional diffusion diagnostics</title><style>body{{font:16px/1.5 system-ui;max-width:1280px;margin:2rem auto;padding:0 1rem;background:#f6f8fb}}.w{{background:#fff2cc;border-left:5px solid #d99b00;padding:1rem}}img{{max-width:100%;background:white;border:1px solid #ccd3df}}table{{border-collapse:collapse;background:white}}td,th{{padding:.5rem;border:1px solid #ccd3df;text-align:right}}td:first-child,th:first-child{{text-align:left}}</style><h1>CAD + process conditional-diffusion representations</h1><p class=w><b>Weak-label diagnostic only.</b> {len(report['runs'])} NG run(s), registration and defect extents are unvalidated, and local heater projection is invalid. Type is an observed diagnostic label, never a validated classifier.</p><h2>Prepared representation and condition inputs</h2><p>{inputs}</p><img src=architecture.png><h2>Denoising loss</h2><img src=loss_comparison.png><h2>Same {run} CAD view and colour scale</h2><img src={run}_four_variant_comparison.png><table><tr><th>representation</th><th>nodes/run</th><th>best loss</th><th>{run} MAE mm</th><th>{run} RMSE mm</th><th>CAD-correspondent Chamfer mm</th><th>registered-scan Chamfer mm</th><th>{run} weak AP</th></tr>{rows}</table><h2>Feature attribution diagnostics</h2><p>{attributions}</p><p><a href=report.json>full report</a></p>""",encoding="utf-8")


def run(args: argparse.Namespace) -> dict[str,Any]:
    seed_everything(); [p.mkdir(parents=True,exist_ok=True) for p in (INP,MOD,OUT,ATTR)]; reports=[]
    for name in VARIANTS:
        data=prepare_variant(name); processed_input_figure(name,data); rep=train_variant(name,data,args.epochs,args.learning_rate,args.diffusion_steps,args.samples,args.sample_steps); attribution(name,data,MOD/f"{name}.pt");reports.append(rep)
    runs=[str(x) for x in data["run_id"]]
    figures(reports,runs[-1]); report={"status":"four_conditional_diffusion_weak_label_diagnostic_completed","production_valid":False,"runs":runs,"variants":reports,"common_contract":"CAD/process conditioning -> signed-normal deformation diffusion -> p=CAD+n*d plus weak risk overlay","defect_type":{"value":data["run_defect_class"].tolist(),"status":"unvalidated_run_level_labels"},"critical_limitations":["registration_independently_validated=false","defect extent/identity not fully reviewed","heater_projection_valid=false; local values zero-imputed with explicit mask",f"only {len(runs)} NG run(s) and one CAD; no held-out result"],"output_directories":{"processed_inputs":"gen_models/processed_inputs","models":"gen_models/models","outputs":"gen_models/outputs","feature_attributions":"gen_models/feature_attributions"}}
    (OUT/"report.json").write_text(json.dumps(report,ensure_ascii=False,indent=2),encoding="utf-8");gallery(report);(ATTR/"README.md").write_text("Gradient plots are local post-fit sensitivity diagnostics of the weak-label risk head. They are not causal process attribution or feature importance.\n",encoding="utf-8");return report


def main() -> None:
    p=argparse.ArgumentParser(description=__doc__);p.add_argument("--allow-diagnostic-weak-labels",action="store_true");p.add_argument("--epochs",type=int,default=28);p.add_argument("--learning-rate",type=float,default=2e-3);p.add_argument("--diffusion-steps",type=int,default=30);p.add_argument("--sample-steps",type=int,default=12);p.add_argument("--samples",type=int,default=4);args=p.parse_args()
    if not args.allow_diagnostic_weak_labels: raise SystemExit("Pass --allow-diagnostic-weak-labels: paired registration, defect extent and heater projection are invalid.")
    report=run(args);print(json.dumps({"status":report["status"],"variants":[{"name":x["variant"],"loss":x["best_loss"]} for x in report["variants"]]},ensure_ascii=False,indent=2))
if __name__=="__main__": main()