"""Diagnostic weak-label GNN training for CAD risk and deformation fields.

This command deliberately requires an opt-in flag. The current two runs have
unvalidated registration, point-dilated defect areas, and one defect class.
The result is a pipeline/training diagnostic, never a production model.
"""
from __future__ import annotations

import argparse
import html
import json
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
import torch.nn.functional as F

from .architectures import PhysicsConstrainedRiskGNN
from .data import ROOT, fixed_physical_normalize, graph_inputs, split_embedded_units
from .training import seed_everything

OUT = ROOT / "results/models/weak_gnn"
DEFORMATION_SCALE_MM = 30.0


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


def _process_vector(data: dict[str, np.ndarray]) -> np.ndarray:
    """Use interpolated heater values while retaining observed/imputed flags."""
    raw = np.asarray(data["process_x"], np.float32).copy()
    observed = np.asarray(data["process_feature_valid"], bool)
    names = np.asarray(data["process_feature_names"]).astype(str)
    base_names, units = split_embedded_units(names)
    imputed = np.zeros(len(raw), bool)
    base_names = base_names.tolist()
    for prefix, filled, filled_mask in (
        ("상부온도", data["heater_top_c"], data["heater_top_imputed"]),
        ("하부온도", data["heater_bottom_c"], data["heater_bottom_imputed"]),
    ):
        for zone in range(1, 18):
            key = f"{prefix}{zone:02d}"
            if key in base_names:
                index = base_names.index(key)
                if filled_mask[zone - 1]:
                    raw[index] = filled[zone - 1]
                    imputed[index] = True
    usable = observed | imputed
    values = np.where(usable, raw, 0).astype(np.float32)
    normalized = fixed_physical_normalize(values[None], np.asarray(base_names), units)[0]
    return np.concatenate([normalized, observed.astype(np.float32), imputed.astype(np.float32)])


def load_training_data() -> dict[str, Any]:
    graph = max(graph_inputs(), key=lambda item: item["nodes"])
    paths = sorted((ROOT / "data_processed/paired").glob("*_cad_ply_targets.npz"))
    if not paths:
        raise RuntimeError("At least one paired target file is required for the diagnostic experiment")
    runs = []
    for path in paths:
        with np.load(path, allow_pickle=False) as raw:
            item = {key: raw[key] for key in raw.files}
        if item["cad_pos_mm"].shape[0] != graph["nodes"]:
            raise RuntimeError(f"{path.name}: target nodes do not match the finest CAD graph")
        runs.append((path, item))

    node_one = np.concatenate([graph["x"], graph["mask"].astype(np.float32)], axis=1)
    node = np.concatenate([node_one for _ in runs], axis=0)
    count = graph["nodes"]
    edge = np.concatenate([graph["edge_index"] + index * count for index in range(len(runs))], axis=1)
    edge_attr = np.concatenate([graph["edge_attr"] for _ in runs], axis=0).astype(np.float32)
    edge_attr[:, 0] /= 1000.0
    batch = np.repeat(np.arange(len(runs)), count)
    process = np.stack([_process_vector(item) for _, item in runs])
    risk = np.concatenate([item["weak_defect_probability"][:, 1] for _, item in runs]).astype(np.float32)
    risk_binary = np.concatenate([item["weak_defect_mask"][:, 1] for _, item in runs]).astype(np.float32)
    signed = np.concatenate([item["cad_to_scan_normal_residual_mm"] for _, item in runs]).astype(np.float32)
    magnitude = np.concatenate([item["cad_to_scan_distance_mm"] for _, item in runs]).astype(np.float32)
    deformation_clip_by_run = []
    for (_, item), run in zip(runs, [str(item["run_id"].item()) for _, item in runs]):
        distance = np.asarray(item["cad_to_scan_distance_mm"], np.float32)
        normal = np.asarray(item["cad_to_scan_normal_residual_mm"], np.float32)
        deformation_clip_by_run.append({
            "run": run, "nodes": int(len(distance)),
            "distance_nodes_at_or_above_cap": int((distance >= DEFORMATION_SCALE_MM).sum()),
            "signed_nodes_outside_cap": int((np.abs(normal) >= DEFORMATION_SCALE_MM).sum()),
            "raw_max_distance_mm": float(distance.max()),
            "raw_max_abs_signed_mm": float(np.abs(normal).max()),
        })
    signed = np.clip(signed, -DEFORMATION_SCALE_MM, DEFORMATION_SCALE_MM) / DEFORMATION_SCALE_MM
    magnitude = np.clip(magnitude, 0, DEFORMATION_SCALE_MM) / DEFORMATION_SCALE_MM
    classes = [str(item["defect_class"].item()) for _, item in runs]
    class_names = sorted(set(classes))
    labels = np.asarray([class_names.index(value) for value in classes], np.int64)
    return {
        "node": node, "edge": edge, "edge_attr": edge_attr, "batch": batch,
        "process": process, "risk": risk, "risk_binary": risk_binary,
        "signed": signed, "magnitude": magnitude, "class_labels": labels,
        "class_names": class_names, "paths": paths, "runs": [str(item["run_id"].item()) for _, item in runs],
        "cad_pos_mm": runs[0][1]["cad_pos_mm"], "nodes_per_run": count,
        "deformation_clip_by_run": deformation_clip_by_run,
    }


def soft_dice_loss(logit: Tensor, target: Tensor) -> Tensor:
    probability = torch.sigmoid(logit)
    return 1 - (2 * (probability * target).sum() + 1) / (probability.sum() + target.sum() + 1)


def train(epochs: int, learning_rate: float) -> dict[str, Any]:
    seed_everything()
    data = load_training_data()
    node = _tensor(data["node"]); edge = _tensor(data["edge"], torch.long)
    edge_attr = _tensor(data["edge_attr"]); batch = _tensor(data["batch"], torch.long)
    process = _tensor(data["process"]); risk = _tensor(data["risk"])
    risk_binary = _tensor(data["risk_binary"]); signed_target = _tensor(data["signed"])
    magnitude_target = _tensor(data["magnitude"]); class_target = _tensor(data["class_labels"], torch.long)
    model = PhysicsConstrainedRiskGNN(
        node.shape[1], edge_attr.shape[1], process.shape[1], max(1, len(data["class_names"])))
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=1e-5)
    positive = risk_binary.sum().clamp_min(1)
    positive_weight = ((len(risk_binary) - positive) / positive).clamp(max=200)
    src, dst = edge
    # Prepared topology is bidirectional. Keep one canonical orientation per pair.
    undirected = src < dst
    history = []
    selected_loss = float("inf")
    selected_epoch = 0
    selected_state: dict[str, Tensor] | None = None
    type_status = "trained_multiclass" if len(data["class_names"]) > 1 else "unavailable_single_observed_class"
    for epoch in range(epochs):
        model.train(); optimizer.zero_grad(set_to_none=True)
        output = model(node, edge, edge_attr, process, batch)
        bce = F.binary_cross_entropy_with_logits(output["risk_logit"], risk, pos_weight=positive_weight)
        dice = soft_dice_loss(output["risk_logit"], risk)
        deformation = F.huber_loss(output["signed_normal_deformation"], signed_target, delta=.15)
        deformation = deformation + F.huber_loss(output["deformation_magnitude"], magnitude_target, delta=.15)
        probability = torch.sigmoid(output["risk_logit"])
        risk_smooth = (probability[src[undirected]] - probability[dst[undirected]]).square().mean()
        deformation_smooth = (
            output["signed_normal_deformation"][src[undirected]] -
            output["signed_normal_deformation"][dst[undirected]]).square().mean()
        if len(data["class_names"]) > 1:
            type_loss = F.cross_entropy(output["defect_type_logits"], class_target)
        else:
            type_loss = output["defect_type_logits"].sum() * 0
        loss = bce + dice + .20 * deformation + .01 * risk_smooth + .005 * deformation_smooth + type_loss
        loss.backward(); optimizer.step()
        loss_value = float(loss.detach())
        if loss_value < selected_loss:
            selected_loss = loss_value
            selected_epoch = epoch + 1
            selected_state = {key: value.detach().clone() for key, value in model.state_dict().items()}
        history.append({
            "epoch": epoch + 1, "total": loss_value, "bce": float(bce.detach()),
            "dice_loss": float(dice.detach()), "deformation": float(deformation.detach()),
            "risk_graph_smoothness": float(risk_smooth.detach()),
            "deformation_graph_smoothness": float(deformation_smooth.detach()),
            "type_loss": float(type_loss.detach()),
        })

    if selected_state is None:
        raise RuntimeError("No training epoch completed")
    model.load_state_dict(selected_state)
    model.eval()
    with torch.no_grad():
        output = model(node, edge, edge_attr, process, batch)
        probability = torch.sigmoid(output["risk_logit"])
        prediction = probability >= .5
        intersection = (prediction & risk_binary.bool()).sum().item()
        union = (prediction | risk_binary.bool()).sum().item()
        dice_value = (2 * intersection) / max(1, prediction.sum().item() + risk_binary.sum().item())
        iou_value = intersection / max(1, union)
        constraint_violation = torch.relu(
            output["signed_normal_deformation"].abs() - output["deformation_magnitude"]).max().item()
        per_run = []
        count = data["nodes_per_run"]
        for index, run in enumerate(data["runs"]):
            sl = slice(index * count, (index + 1) * count)
            run_prediction = prediction[sl]
            run_target = risk_binary[sl].bool()
            run_intersection = (run_prediction & run_target).sum().item()
            run_union = (run_prediction | run_target).sum().item()
            per_run.append({
                "run": run,
                "target_positive_nodes": int(run_target.sum()),
                "predicted_positive_nodes": int(run_prediction.sum()),
                "training_iou": run_intersection / max(1, run_union),
                "training_dice": 2 * run_intersection / max(1, run_prediction.sum().item() + run_target.sum().item()),
                "max_fitted_risk_probability": float(probability[sl].max()),
            })

    OUT.mkdir(parents=True, exist_ok=True)
    checkpoint = OUT / "diagnostic_weak_gnn.pt"
    torch.save({
        "state_dict": model.state_dict(), "production_valid": False,
        "class_names": data["class_names"], "output_channels": ["risk", "signed_deformation", "magnitude"],
        "warning": f"{len(data['runs'])} unvalidated NG run(s); weak point-dilated labels; training diagnostic only",
    }, checkpoint)
    plot_results(data, history, output)
    result = {
        "status": "diagnostic_weak_label_training_completed",
        "production_valid": False, "generalization_measured": False,
        "input": "STEP-derived CAD graph + XLS process values/masks; interpolated heaters include imputation flags",
        "target_source": "registered fused all_merge_pointcloud_full PLY + defect CSV weak geodesic area",
        "runs": data["runs"], "independent_samples": len(data["runs"]),
        "defect_classes": data["class_names"], "defect_type_training": type_status,
        "defect_type_reason": "The supplied run-level target has only one class and no reviewed OK run.",
        "epochs": epochs, "learning_rate": learning_rate,
        "loss_initial": history[0]["total"], "loss_final": selected_loss,
        "loss_last_epoch": history[-1]["total"], "selected_epoch": selected_epoch,
        "training_threshold_iou": iou_value, "training_threshold_dice": dice_value,
        "training_metrics_by_run": per_run,
        "deformation_target_preprocessing": {
            "diagnostic_clip_mm": DEFORMATION_SCALE_MM,
            "reason": "bound the influence of unvalidated nearest-scan registration residuals",
            "counts_by_run": data["deformation_clip_by_run"],
        },
        "hard_constraint": "predicted magnitude >= abs(predicted signed normal deformation)",
        "hard_constraint_max_violation": constraint_violation,
        "soft_constraints": ["risk graph smoothness", "signed-deformation graph smoothness"],
        "physics_constraints_unavailable": [
            "heat-equation residual: missing time histories, material conductivity/capacity and physical heater footprints",
            "thermo-mechanical residual: missing temperature-dependent material law, thickness, contact and boundary conditions",
        ],
        "checkpoint": str(checkpoint.relative_to(ROOT)),
        "figures": ["results/models/weak_gnn/training.png", "results/models/weak_gnn/predictions.png"],
        "history": history,
    }
    (OUT / "report.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    metrics_rows = "".join(
        "<tr>"
        f"<td>{html.escape(row['run'])}</td>"
        f"<td>{row['target_positive_nodes']}</td>"
        f"<td>{row['predicted_positive_nodes']}</td>"
        f"<td>{row['training_iou']:.4f}</td>"
        f"<td>{row['training_dice']:.4f}</td>"
        "</tr>"
        for row in per_run
    )
    (OUT / "index.html").write_text(f"""<!doctype html>
<html lang="ko"><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>Diagnostic weak-label GNN</title>
<style>
body{{font:16px/1.5 system-ui,sans-serif;max-width:1180px;margin:2rem auto;padding:0 1rem;color:#172033;background:#f6f8fb}}
.warning{{background:#fff2cc;border-left:5px solid #d99b00;padding:1rem}} img{{max-width:100%;background:white;border:1px solid #ccd3df;border-radius:8px}}
table{{border-collapse:collapse;background:white}} th,td{{padding:.55rem .8rem;border:1px solid #ccd3df;text-align:right}} th:first-child,td:first-child{{text-align:left}}
code{{background:#e9edf3;padding:.15rem .3rem;border-radius:3px}}
</style>
<h1>CAD risk-area + deformation GNN</h1>
<p class="warning"><strong>Diagnostic training-set fit only.</strong> {len(data['runs'])} unvalidated NG run(s), weak point-dilated masks, one CAD and one run-level class. These values do not measure generalization or production accuracy.</p>
<p>Input: STEP-derived graph + XLS process values/masks. Outputs: CAD-node risk probability, signed normal deformation, deformation magnitude, and a graph-level defect-type head.</p>
<p>Final loss: <strong>{result['loss_final']:.4f}</strong>; pooled training IoU: <strong>{iou_value:.4f}</strong>; pooled training Dice: <strong>{dice_value:.4f}</strong>. Defect-type status: <code>{html.escape(type_status)}</code>.</p>
<table><thead><tr><th>Run</th><th>Target nodes</th><th>Predicted nodes</th><th>Training IoU</th><th>Training Dice</th></tr></thead><tbody>{metrics_rows}</tbody></table>
<h2>Training trace</h2><img src="training.png" alt="Diagnostic GNN training curve">
<h2>Fitted fields</h2><img src="predictions.png" alt="Weak target, fitted risk and fitted deformation fields">
<p><a href="report.json">Full JSON report</a> · <a href="diagnostic_weak_gnn.pt">Diagnostic checkpoint</a></p>
</html>""", encoding="utf-8")
    return result


def plot_results(data: dict[str, Any], history: list[dict[str, float]], output: dict[str, Tensor]) -> None:
    epochs = [row["epoch"] for row in history]
    fig, ax = plt.subplots(figsize=(8, 4), constrained_layout=True)
    for key in ("total", "bce", "dice_loss", "deformation"):
        ax.plot(epochs, [row[key] for row in history], label=key)
    ax.set_xlabel("epoch"); ax.set_ylabel("diagnostic training loss"); ax.set_yscale("log")
    ax.legend(); ax.grid(alpha=.25); ax.set_title("Weak-label GNN training curve")
    fig.savefig(OUT / "training.png", dpi=150); plt.close(fig)

    probability = torch.sigmoid(output["risk_logit"]).cpu().numpy()
    magnitude = output["deformation_magnitude"].cpu().numpy() * DEFORMATION_SCALE_MM
    target = data["risk"]
    count = data["nodes_per_run"]; pos = data["cad_pos_mm"]
    fig, axes = plt.subplots(len(data["runs"]), 3, figsize=(15, 4.8 * len(data["runs"])), constrained_layout=True)
    if len(data["runs"]) == 1: axes = axes[None]
    for index, run in enumerate(data["runs"]):
        sl = slice(index * count, (index + 1) * count)
        panels = ((target[sl], "weak 10 mm risk target", "magma", 0, 1),
                  (probability[sl], "fitted risk probability", "magma", 0, 1),
                  (magnitude[sl], "fitted deformation magnitude (30 mm target cap)", "viridis", 0, DEFORMATION_SCALE_MM))
        for ax, (value, title, cmap, low, high) in zip(axes[index], panels):
            if "risk" in title:
                ax.scatter(pos[:, 0], pos[:, 2], c="#d9dde3", s=1, rasterized=True)
                shown = value > .02
                draw = ax.scatter(pos[shown, 0], pos[shown, 2], c=value[shown], s=10,
                                  cmap=cmap, vmin=low, vmax=high, rasterized=True)
                peak = int(np.argmax(value))
                ax.scatter(pos[peak, 0], pos[peak, 2], marker="x", c="#00b7ff", s=45,
                           linewidths=1.4, label="maximum")
            else:
                draw = ax.scatter(pos[:, 0], pos[:, 2], c=value, s=2, cmap=cmap,
                                  vmin=low, vmax=high, rasterized=True)
            ax.set_aspect("equal"); ax.set_xlabel("CAD X (mm)"); ax.set_ylabel("CAD Z (mm)")
            ax.set_title(f"{run}: {title}"); fig.colorbar(draw, ax=ax)
    fig.suptitle(f"Training-set visualization only; {len(data['runs'])} weak/unvalidated NG run(s), no held-out performance")
    fig.savefig(OUT / "predictions.png", dpi=150); plt.close(fig)


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=60)
    parser.add_argument("--learning-rate", type=float, default=2e-3)
    args = parser.parse_args()
    if not args.allow_diagnostic_weak_labels:
        raise SystemExit("Refusing weak/unvalidated targets. Pass --allow-diagnostic-weak-labels for the requested diagnostic overfit experiment.")
    result = train(args.epochs, args.learning_rate)
    print(json.dumps({key: result[key] for key in (
        "status", "loss_initial", "loss_final", "training_threshold_iou",
        "training_threshold_dice", "defect_type_training")}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()