"""Paper-inspired TAG U-Net conditional diffusion diagnostic.

This is a compact adaptation of Ferguson et al., *Topology-Agnostic Graph
U-Nets for Scalar Field Prediction on Unstructured Meshes* (2024).  It keeps
the paper's key mechanics: balanced kd-tree pooling into 3-D groups of up to
eight, mean pooling, copied unpooling, 12-NN coarse graphs, EdgeConv-style
max aggregation, encoder/decoder concatenation skips, and coordinates at each
convolution.  The paper predicts a scalar field directly; here the U-Net is a
DDPM denoiser conditioned on the same CAD/XLS vectors as ``graph_6000``.

It is deliberately a bounded, CPU weak-label diagnostic: two NG scans are not
enough to validate a generative deformation model.
"""
from __future__ import annotations

import argparse
import json
import math
import re
from pathlib import Path
from typing import Any

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

from gen_models.models.four_conditional_diffusion import (
    ATTR, EXP, INP, MOD, OUT, SCALE_MM, average_precision, chamfer,
    prepare_variant, schedule, tensor,
)
from models.training import seed_everything


NAME = "graph_unet_6000"


def kd_groups(points: np.ndarray, maximum_leaf_size: int = 8) -> np.ndarray:
    """Balanced kd-tree leaves, represented as a fine -> coarse assignment."""
    leaves: list[np.ndarray] = []

    def split(index: np.ndarray) -> None:
        if len(index) <= maximum_leaf_size:
            leaves.append(index)
            return
        local = points[index]
        axis = int(np.argmax(np.ptp(local, axis=0)))
        order = index[np.argsort(local[:, axis], kind="mergesort")]
        midpoint = len(order) // 2
        split(order[:midpoint]); split(order[midpoint:])

    split(np.arange(len(points), dtype=np.int64))
    assignment = np.empty(len(points), dtype=np.int64)
    for group, indices in enumerate(leaves): assignment[indices] = group
    return assignment


def pooled_positions(points: np.ndarray, assignment: np.ndarray) -> np.ndarray:
    count = int(assignment.max()) + 1
    out = np.zeros((count, 3), np.float32)
    np.add.at(out, assignment, points)
    return out / np.bincount(assignment, minlength=count)[:, None]


def bidirectional_knn(points: np.ndarray, k: int = 12) -> np.ndarray:
    """Destination/source edge index for a bidirectional, duplicate-free kNN graph."""
    if len(points) < 2: return np.zeros((2, 0), dtype=np.int64)
    _, neighbours = cKDTree(points).query(points, k=min(k + 1, len(points)), workers=1)
    pairs = {(int(i), int(j)) for i, row in enumerate(np.atleast_2d(neighbours)) for j in row[1:] if i != j}
    pairs |= {(j, i) for i, j in pairs}
    return np.asarray(sorted(pairs), dtype=np.int64).T


def hierarchy(points: np.ndarray, edge0: np.ndarray) -> dict[str, Any]:
    """Three pool stages exactly following the paper's 3-D pooling schedule."""
    positions = [points.astype(np.float32)]
    # Preserve actual CAD topology at level 0.  All coarse levels follow the
    # paper's 12-nearest-neighbour construction after kd-tree pooling.
    edges = [edge0.astype(np.int64)]
    assignments = []
    current = points
    for _ in range(3):
        assign = kd_groups(current, 8)
        assignments.append(assign)
        current = pooled_positions(current, assign)
        positions.append(current)
        edges.append(bidirectional_knn(current, 12))
    return {"positions": positions, "edges": edges, "assignments": assignments}


def as_edge_index(edge: np.ndarray, node_count: int) -> np.ndarray:
    """Make source/destination explicit and add reverse edges if needed."""
    a = edge.astype(np.int64)
    pairs = {(int(s), int(d)) for s, d in a.T if s != d and s < node_count and d < node_count}
    pairs |= {(d, s) for s, d in pairs}
    return np.asarray(sorted(pairs), dtype=np.int64).T


def pool_mean(x: Tensor, assignment: Tensor, output_count: int) -> Tensor:
    out = torch.zeros((output_count, x.shape[1]), device=x.device, dtype=x.dtype)
    out.index_add_(0, assignment, x)
    denom = torch.bincount(assignment, minlength=output_count).to(x.dtype).clamp_min_(1).unsqueeze(1)
    return out / denom


class EdgeConv(nn.Module):
    """EdgeConv with explicit edge delta/distance as a CAD geometry extension."""
    def __init__(self, channels: int, coordinate_scale_mm: float = 1000.) -> None:
        super().__init__()
        self.coordinate_scale_mm = coordinate_scale_mm
        # h(x_i || x_j - x_i || coordinate_i || delta_xyz || unit_direction || length)
        self.message = nn.Sequential(
            nn.Linear(channels * 2 + 10, channels * 2), nn.SiLU(),
            nn.Linear(channels * 2, channels),
        )
        self.norm = nn.BatchNorm1d(channels)

    def forward(self, x: Tensor, position: Tensor, edge: Tensor) -> Tensor:
        # edge is source -> destination; collect incoming source messages at destination.
        src, dst = edge[0], edge[1]
        delta = (position[src] - position[dst]) / self.coordinate_scale_mm
        length = delta.norm(dim=1, keepdim=True)
        direction = delta / length.clamp_min(1e-8)
        message_input = torch.cat((x[dst], x[src] - x[dst], position[dst] / self.coordinate_scale_mm, delta, direction, length), 1)
        message = self.message(message_input)
        out = torch.full_like(x, -torch.inf)
        out.scatter_reduce_(0, dst[:, None].expand_as(message), message, reduce="amax", include_self=True)
        # kNN/coarse and original CAD graphs both contain neighbours; this just
        # makes the operation robust to isolated nodes.
        out = torch.where(torch.isfinite(out), out, x)
        return F.silu(self.norm(out + x))


class TagUNetDenoiser(nn.Module):
    """Three-resolution TAG U-Net serving as a conditional DDPM denoiser."""
    def __init__(self, condition_dim: int, graph: dict[str, Any], width: int = 32) -> None:
        super().__init__()
        self.width = width
        self.positions = []
        self.edges = []
        self.assignments = []
        for i, p in enumerate(graph["positions"]):
            self.register_buffer(f"position_{i}", torch.as_tensor(p, dtype=torch.float32)); self.positions.append(f"position_{i}")
            self.register_buffer(f"edge_{i}", torch.as_tensor(graph["edges"][i], dtype=torch.long)); self.edges.append(f"edge_{i}")
        for i, a in enumerate(graph["assignments"]):
            self.register_buffer(f"assignment_{i}", torch.as_tensor(a, dtype=torch.long)); self.assignments.append(f"assignment_{i}")
        self.cond = nn.Sequential(nn.Linear(condition_dim, width), nn.SiLU(), nn.Linear(width, width))
        self.time = nn.Sequential(nn.Linear(16, width), nn.SiLU(), nn.Linear(width, width))
        self.input = nn.Linear(width * 2 + 1, width)
        self.enc = nn.ModuleList([EdgeConv(width) for _ in range(3)])
        self.bottleneck = EdgeConv(width)
        self.decode_project = nn.ModuleList([nn.Linear(width * 2, width) for _ in range(3)])
        self.dec = nn.ModuleList([EdgeConv(width) for _ in range(3)])
        self.epsilon = nn.Sequential(nn.Linear(width, width), nn.SiLU(), nn.Linear(width, 1))
        self.risk = nn.Sequential(nn.Linear(width, width // 2), nn.SiLU(), nn.Linear(width // 2, 1))

    def _get(self, names: list[str], index: int) -> Tensor: return getattr(self, names[index])
    def _time(self, step: Tensor) -> Tensor:
        f = torch.exp(torch.linspace(0, math.log(1000), 8, device=step.device, dtype=torch.float32))
        p = step[:, None].float() * f[None] / 30.
        return self.time(torch.cat((p.sin(), p.cos()), 1))

    def _single(self, x: Tensor, t: Tensor, condition: Tensor) -> tuple[Tensor, Tensor]:
        h = self.input(torch.cat((x[:, None], self.cond(condition), self._time(t)), 1))
        skips: list[Tensor] = []
        for level in range(3):
            h = self.enc[level](h, self._get(self.positions, level), self._get(self.edges, level)); skips.append(h)
            assignment = self._get(self.assignments, level)
            h = pool_mean(h, assignment, len(self._get(self.positions, level + 1)))
        h = self.bottleneck(h, self._get(self.positions, 3), self._get(self.edges, 3))
        for level in (2, 1, 0):
            h = h[self._get(self.assignments, level)]  # average unpool: copy to constituents
            h = self.decode_project[2 - level](torch.cat((h, skips[level]), 1))
            h = self.dec[2 - level](h, self._get(self.positions, level), self._get(self.edges, level))
        return self.epsilon(h).squeeze(1), self.risk(h).squeeze(1)

    def forward(self, x: Tensor, t: Tensor, condition: Tensor) -> tuple[Tensor, Tensor]:
        # The graph hierarchy is shared CAD geometry, evaluated per run.
        n = len(self._get(self.positions, 0))
        if len(x) == 0 or len(x) % n: raise ValueError("TAG U-Net requires full CAD-node runs")
        out = [self._single(x[s:s+n], t[s:s+n], condition[s:s+n]) for s in range(0,len(x),n)]
        return torch.cat([item[0] for item in out]), torch.cat([item[1] for item in out])


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


def metrics_and_write(model: TagUNetDenoiser, data: dict[str, Any], abar: Tensor, samples: int, sample_steps: int) -> list[dict[str, Any]]:
    condition, target = tensor(data["condition"]), tensor(data["signed_target_scaled"])
    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(); n = int(data["nodes_per_run"]); 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
        truth = target.numpy()[sl] * SCALE_MM; cloud = pos + normal * mean[:, None]; reference = pos + normal * truth[:, None]
        y = data["risk_binary"][sl].astype(bool); p = score[sl] >= .5; inter = int((p & y).sum())
        scan = data[f"scan_points_{run}_cad_mm"]
        result.append({"run": str(run), "MAE_mm": float(np.abs(mean-truth).mean()), "RMSE_mm": float(np.sqrt(np.mean((mean-truth)**2))), "CAD_correspondent_Chamfer_mm": chamfer(cloud, reference), "registered_scan_Chamfer_mm": chamfer(cloud, scan), "risk": {"AP": average_precision(score[sl], y), "IoU_at_0.5": inter/max(1,int((p|y).sum())), "Dice_at_0.5": 2*inter/max(1,int(p.sum()+y.sum()))}, "mean_stochastic_std_mm": float(std.mean()), "max_stochastic_std_mm": float(std.max())})
    runs=len(data["run_id"])
    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("unvalidated_run_level_labels"))
    return result


def attribution(model: TagUNetDenoiser, data: dict[str, Any]) -> None:
    base = tensor(data["condition"]); names = data["process_feature_names"]; start = 10
    c = base.clone().requires_grad_(True); _, logit = model(torch.zeros(len(c)), torch.zeros(len(c), dtype=torch.long), c); logit.mean().backward(); risk = c.grad[:, start:start+len(names)].abs().mean(0).detach().numpy()
    c = base.clone().requires_grad_(True); eps, _ = model(torch.zeros(len(c)), torch.full((len(c),),15,dtype=torch.long),c); abar=schedule(30)[15]; x0=-(1-abar).sqrt()*eps/abar.sqrt(); x0.abs().mean().backward(); deform=c.grad[:,start:start+len(names)].abs().mean(0).detach().numpy()
    payload = {"variant": NAME, "method": "absolute local gradients, fixed DDPM state; not causal attribution", "production_valid": False, "risk_features": [{"rank":i+1,"source_name":str(names[k]),"mean_absolute_gradient":float(risk[k])} for i,k in enumerate(np.argsort(risk)[::-1])], "deformation_features": [{"rank":i+1,"source_name":str(names[k]),"mean_absolute_gradient":float(deform[k])} for i,k in enumerate(np.argsort(deform)[::-1])]}
    (ATTR / f"{NAME}_process_gradient.json").write_text(json.dumps(payload,ensure_ascii=False,indent=2),encoding="utf-8")
    fig, ax = plt.subplots(1,2,figsize=(13,4),constrained_layout=True)
    for axis, values, title, color in zip(ax,(risk,deform),("weak-risk logit sensitivity","denoised deformation sensitivity"),("#495f9d","#a85b35")):
        top=np.argsort(values)[-12:][::-1]; axis.barh([f"process feature {k+1}" for k in top][::-1],values[top][::-1],color=color);axis.set_title(title);axis.set_xlabel("mean absolute local gradient");axis.grid(axis="x",alpha=.2)
    fig.suptitle("graph_unet_6000: fixed-state sensitivity diagnostic (not causal)"); fig.savefig(ATTR/f"{NAME}_process_gradient.png",dpi=160);plt.close(fig)


def figures(data: dict[str, Any], graph: dict[str, Any], history: list[dict[str, float]], metrics: list[dict[str, Any]]) -> None:
    run=str(data["run_id"][-1])
    # Hierarchy itself is a processed-input artifact so the exact graph changes are inspectable.
    np.savez_compressed(INP/f"{NAME}_hierarchy.npz", **{f"positions_{i}":p for i,p in enumerate(graph["positions"])}, **{f"edges_{i}":e for i,e in enumerate(graph["edges"])}, **{f"pool_assignment_{i}_to_{i+1}":a for i,a in enumerate(graph["assignments"])})
    summary={"paper":"Topology-Agnostic Graph U-Nets for Scalar Field Prediction on Unstructured Meshes (Ferguson et al., 2024)","adaptation":"3-level kd-tree TAG U-Net denoiser; source paper is direct scalar regression, not DDPM","pooling":"balanced kd-tree leaves, at most 8 3-D nodes; mean pool / copied unpool","convolution":"EdgeConv max aggregation with CAD coordinates each convolution","level_0":"prepared CAD edge_index with explicit relative delta xyz and Euclidean edge length","coarse_levels":"12-NN bidirectional edges after each pool, with relative delta xyz and Euclidean length recomputed from coarse coordinates","sizes":[len(p) for p in graph["positions"]],"edges":[int(e.shape[1]) for e in graph["edges"]],"condition":"same XLS/global heater values and validity mask as graph_6000","production_valid":False,"limitations":[f"{len(data['run_id'])} NG run(s), no held-out test","registration and defect extents unvalidated","heater footprints unavailable; local heater projection mask false"]}
    (INP/f"{NAME}_hierarchy.json").write_text(json.dumps(summary,ensure_ascii=False,indent=2),encoding="utf-8")
    fig=plt.figure(figsize=(13,4.4),constrained_layout=True); ax=fig.add_subplot(1,2,1,projection="3d")
    colors=["#274c77","#6096ba","#a3cef1","#f6bd60"]
    for level,(p,col) in enumerate(zip(graph["positions"],colors)):
        ids=np.linspace(0,len(p)-1,min(900,len(p)),dtype=int);ax.scatter(p[ids,0],p[ids,1],p[ids,2],s=2+level*5,c=col,label=f"level {level}: {len(p):,} nodes",alpha=.65)
    ax.set_title("CAD kd-tree Graph U-Net hierarchy");ax.set_xlabel("X mm");ax.set_ylabel("Y mm");ax.set_zlabel("Z mm");ax.legend(fontsize=7);ax.view_init(22,-65)
    ax2=fig.add_subplot(1,2,2); ax2.plot([h["epoch"] for h in history],[h["denoise"] for h in history],label="denoise MSE");ax2.plot([h["epoch"] for h in history],[h["risk"] for h in history],label="risk BCE");ax2.set_yscale("log");ax2.set_title("bounded CPU fit loss");ax2.set_xlabel("epoch");ax2.grid(alpha=.2);ax2.legend()
    fig.savefig(INP/f"{NAME}_hierarchy.png",dpi=160);plt.close(fig)
    # Same fixed last-run CAD X-Z view and range used by the graph comparison.
    gen=np.load(OUT/f"{NAME}_generated.npz"); pos=gen["cad_pos_mm"];normal=gen["cad_normal"];disp=gen["generated_signed_normal_mm"][0,-1];risk=gen["risk_probability"][-1];base=np.load(OUT/"graph_6000_generated.npz")
    fig,axes=plt.subplots(2,2,figsize=(10,8),constrained_layout=True)
    for col,(label,a) in enumerate((("graph_6000 baseline",base),("graph_unet_6000 TAG",gen))):
        d=a["generated_signed_normal_mm"][0,-1];cloud=pos+normal*d[:,None];r=a["risk_probability"][-1]
        for row,(value,title,cmap,limit) in enumerate(((abs(d),"generated deformation |mm|","viridis",30),(r,"weak risk probability","magma",1))):
            shown=axes[row,col].scatter(cloud[:,0],cloud[:,2],c=value,s=2,cmap=cmap,vmin=0,vmax=limit,rasterized=True);axes[row,col].set_title(f"{label}: {title}");axes[row,col].set_xlabel("CAD X (mm)");axes[row,col].set_ylabel("CAD Z (mm)");axes[row,col].set_aspect("equal");fig.colorbar(shown,ax=axes[row,col],shrink=.72)
    fig.suptitle(f"{run} same CAD X-Z view / same colour scales: baseline vs TAG U-Net");fig.savefig(OUT/f"{NAME}_{run}_comparison.png",dpi=170);plt.close(fig)
    baseline=json.loads((OUT/"report.json").read_text(encoding="utf-8")); old=next(v for v in baseline["variants"] if v["variant"]=="graph_6000")
    fig,ax=plt.subplots(figsize=(8,3.8));labs=["graph_6000","graph_unet_6000"];mae=[old["metrics"][-1]["MAE_mm"],metrics[-1]["MAE_mm"]];ap=[old["metrics"][-1]["risk"]["AP"] or 0,metrics[-1]["risk"]["AP"] or 0];x=np.arange(2);ax.bar(x-.18,mae,.36,label=f"{run} MAE mm",color="#36648b");a=ax.twinx();a.bar(x+.18,ap,.36,label="weak AP",color="#ce7e3a");ax.set_xticks(x,labs);ax.set_ylabel("MAE (lower)");a.set_ylabel("AP (higher)");ax.set_title("Graph baseline vs paper-inspired TAG U-Net (weak-label fit)");ax.grid(axis="y",alpha=.2);ax.legend(loc="upper left");a.legend(loc="upper right");fig.savefig(OUT/f"{NAME}_metric_comparison.png",dpi=160);plt.close(fig)


def run(args: argparse.Namespace) -> dict[str, Any]:
    seed_everything(); [p.mkdir(parents=True,exist_ok=True) for p in (INP,MOD,OUT,ATTR)]
    data=prepare_variant("graph_6000"); raw=np.load(EXP.parent/"data_processed/graphs/graph_6000.npz",allow_pickle=False); edge0=as_edge_index(raw["edge_index"],len(data["pos_mm"])); graph=hierarchy(data["pos_mm"],edge0)
    run=str(data["run_id"][-1])
    model=TagUNetDenoiser(int(data["condition"].shape[1]),graph,width=args.width);condition=tensor(data["condition"]);target=tensor(data["signed_target_scaled"]);risk=tensor(data["risk_target"]);abar=schedule(args.diffusion_steps);opt=torch.optim.AdamW(model.parameters(),lr=args.learning_rate,weight_decay=1e-5);pos_weight=((len(risk)-risk.sum())/risk.sum().clamp_min(1)).clamp(max=200);history=[];best=(float("inf"),0,None)
    for epoch in range(args.epochs):
        opt.zero_grad();t=torch.randint(0,args.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=pos_weight);loss=denoise+.35*rloss;loss.backward();torch.nn.utils.clip_grad_norm_(model.parameters(),3.);opt.step();score=torch.sigmoid(logit).detach().numpy();history.append({"epoch":epoch+1,"total":float(loss.detach()),"denoise":float(denoise.detach()),"risk":float(rloss.detach()),"AP":average_precision(score,data["risk_binary"].astype(bool)) or 0.});
        if float(loss.detach())<best[0]:best=(float(loss.detach()),epoch+1,{k:v.detach().clone() for k,v in model.state_dict().items()})
    model.load_state_dict(best[2]);model.eval();checkpoint=MOD/f"{NAME}.pt";torch.save({"state_dict":model.state_dict(),"architecture":"paper-inspired 3-level TAG U-Net EdgeConv conditional DDPM denoiser","paper_reproduction":False,"hierarchy_sizes":[len(p) for p in graph["positions"]],"condition_dim":int(condition.shape[1]),"production_valid":False},checkpoint)
    metric=metrics_and_write(model,data,abar,args.samples,args.sample_steps);attribution(model,data);figures(data,graph,history,metric)
    report={"variant":NAME,"architecture":"paper-inspired TAG U-Net with 3 kd-tree pool/unpool levels, EdgeConv max aggregation, concat skips","paper_reproduction":False,"paper_difference":f"the paper uses direct scalar-field regression with larger ReLU EdgeConv networks and much more training data; this compact width-{args.width} SiLU backbone is adapted as a conditional DDPM denoiser","nodes_per_run":int(data["nodes_per_run"]),"condition_dim":int(condition.shape[1]),"hierarchy_sizes":[len(p) for p in graph["positions"]],"history":history,"best_loss":best[0],"best_epoch":best[1],"metrics":metric,"checkpoint":str(checkpoint.relative_to(EXP.parent)),"production_valid":False,"limitation":f"{len(data['run_id'])} weak/unvalidated NG run(s); no held-out generalization result"}
    (OUT/f"{NAME}_report.json").write_text(json.dumps(report,ensure_ascii=False,indent=2),encoding="utf-8")
    # Keep the existing four-representation report intact while exposing this
    # fifth, paper-inspired graph variant in its shared comparison manifest.
    shared_path = OUT / "report.json"
    shared = json.loads(shared_path.read_text(encoding="utf-8"))
    shared["variants"] = [v for v in shared["variants"] if v.get("variant") != NAME] + [report]
    shared["tag_graph_unet_comparison"] = {
        "baseline": "graph_6000",
        "candidate": NAME,
        "fixed_view": f"{run} CAD X-Z; identical deformation/risk scales",
        "artifact": f"gen_models/outputs/{NAME}_{run}_comparison.png",
        "status": "weak-label fit only; not a held-out comparison",
    }
    shared_path.write_text(json.dumps(shared, ensure_ascii=False, indent=2), encoding="utf-8")
    # Do not rebuild the original gallery (it intentionally contains four
    # representations); append one durable, explicit entry for this variant.
    gallery_path = OUT / "index.html"
    gallery = gallery_path.read_text(encoding="utf-8")
    marker = "TAG_GRAPH_UNET_ENTRY"
    base_report = next(v for v in shared["variants"] if v.get("variant") == "graph_6000")
    base_18 = base_report["metrics"][-1]
    tag_18 = metric[-1]
    entry = (f"<section id=tag-graph-unet><h2>TAG Graph U-Net candidate</h2><!-- {marker} -->"
             f"<p>Paper-inspired three-level kd-tree Graph U-Net inside the conditional DDPM. "
             f"It is a weak-label, {len(data['run_id'])}-run fit and is not a validation result.</p>"
             f"<table><tr><th>{run}</th><th>MAE mm ↓</th><th>CAD Chamfer mm ↓</th><th>weak AP ↑</th><th>IoU@0.5 ↑</th><th>uncertainty mm ↓</th></tr>"
             f"<tr><td>graph_6000</td><td>{base_18['MAE_mm']:.3f}</td><td>{base_18['CAD_correspondent_Chamfer_mm']:.3f}</td><td>{base_18['risk']['AP']:.3f}</td><td>{base_18['risk']['IoU_at_0.5']:.4f}</td><td>{base_18['mean_stochastic_std_mm']:.2f}</td></tr>"
             f"<tr><td>graph_unet_6000</td><td>{tag_18['MAE_mm']:.3f}</td><td>{tag_18['CAD_correspondent_Chamfer_mm']:.3f}</td><td>{tag_18['risk']['AP']:.3f}</td><td>{tag_18['risk']['IoU_at_0.5']:.4f}</td><td>{tag_18['mean_stochastic_std_mm']:.2f}</td></tr></table>"
             f"<p><a href={NAME}_{run}_comparison.png>same-view graph baseline vs TAG U-Net</a> · "
             f"<a href={NAME}_metric_comparison.png>metric comparison</a> · "
             f"<a href=../processed_inputs/{NAME}_hierarchy.png>kd-tree hierarchy</a> · "
             f"<a href=../processed_inputs/{NAME}_hierarchy.json>hierarchy JSON</a> · "
             f"<a href=../feature_attributions/{NAME}_process_gradient.png>attribution</a> · "
             f"<a href={NAME}_report.json>full run report</a></p></section>")
    if marker in gallery:
        gallery = re.sub(r"<section id=tag-graph-unet>.*?</section>", entry, gallery, flags=re.S)
    else:
        gallery += entry
    gallery_path.write_text(gallery, encoding="utf-8")
    return report


def main() -> None:
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument("--allow-diagnostic-weak-labels",action="store_true");parser.add_argument("--epochs",type=int,default=10);parser.add_argument("--learning-rate",type=float,default=1e-3);parser.add_argument("--diffusion-steps",type=int,default=30);parser.add_argument("--sample-steps",type=int,default=10);parser.add_argument("--samples",type=int,default=3);parser.add_argument("--width",type=int,default=32);args=parser.parse_args()
    if not args.allow_diagnostic_weak_labels:raise SystemExit("Pass --allow-diagnostic-weak-labels: only weak, unvalidated labels are available.")
    print(json.dumps(run(args),ensure_ascii=False,indent=2))

if __name__=="__main__":main()