"""Train the published HiCon-GNN and ANN baseline."""

from __future__ import annotations

import argparse
import csv
import hashlib
import os
import random
import sys
from typing import Dict, List, Optional, Sequence, TypedDict

import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader


class ModelResult(TypedDict):
    metrics: Dict[str, float]
    history: List[Dict[str, float]]
    predictions: Dict[str, np.ndarray]
    checkpoint_path: Optional[str]


from concrete_gnn import (
    DEFAULT_SCHEMA,
    IntegratedMultiscaleModel,
    Standardizer,
    StrengthHead,
    TABULAR_DIM,
    TARGET_NAME,
    TabularANN,
    ConcreteMixDataset,
    collect_strength_predictions,
    collate_multiscale,
    evaluate_strength,
    fit_encoder_standardizers,
    plot_parity,
    plot_r2_bars,
    plot_rve_graph,
    plot_training_curves,
    read_table,
    load_split_indices,
    random_split_indices,
    split_indices_sha256,
    strength_column,
)


OUTPUT_ROOT = "outputs"

PAPER_LOSS_WEIGHTS = {
    "mortar": 0.5,
    "itz": 0.5,
    "itz_prior": 0.05,
    "paste": 0.5,
    "paste_shape": 0.3,
    "paste_hierarchy": 0.15,
    "tensile": 0.3,
    "flexural": 0.3,
    "strength_consistency": 0.1,
    "itz_wb": 0.1,
    "porosity_order": 0.1,
    "paste_mortar_consistency": 0.05,
    "itz_ratio": 0.1,
}


# Literature ranges in ITZ target order.
ITZ_PRIOR_BANDS = (
    (0, 0.0, 70.0),    # itz_thickness_um (matches the [5,70] output transform)
    (1, 0.10, 0.70),   # itz_porosity (floor pulls it above the 0.05 sigmoid min)
    (3, 2.0, 40.0),    # itz_elastic_modulus_GPa (0.3-0.7 x paste must be reachable)
)


def _masked_target_stats(samples, attr: str, dim: int):
    """Return per-target statistics; fewer than two labels remain disabled."""
    vals = torch.stack([getattr(s, attr) for s in samples])  # (N, dim)
    mean = torch.full((dim,), float("nan"))
    std = torch.full((dim,), float("nan"))
    for j in range(dim):
        col = vals[:, j]
        obs = col[~torch.isnan(col)]
        if obs.numel() >= 2:
            mean[j] = obs.mean()
            std[j] = obs.std().clamp_min(1.0e-6)
    return mean, std


def _masked_target_mse(pred, target, mean, std):
    """Masked, per-target-standardized MSE over the supervised label cells."""
    valid = torch.isfinite(mean) & torch.isfinite(std)
    mask = ~torch.isnan(target) & valid
    if not bool(mask.any()):
        return None
    safe_mean = torch.nan_to_num(mean)
    safe_std = torch.nan_to_num(std, nan=1.0)
    pe = (pred - safe_mean) / safe_std
    te = (torch.nan_to_num(target) - safe_mean) / safe_std
    return ((pe - te) ** 2)[mask].mean()


def _itz_range_penalty(itz_pred):
    """Soft hinge keeping per-edge ITZ predictions inside literature bands."""
    if itz_pred is None or itz_pred.numel() == 0:
        return None
    pen = itz_pred.new_zeros(())
    for idx, lo, hi in ITZ_PRIOR_BANDS:
        x = itz_pred[:, idx]
        w = hi - lo
        pen = pen + (torch.relu(lo - x) / w).mean() + (torch.relu(x - hi) / w).mean()
    return pen


# Literature ratios for weakly supervised ITZ properties.
ITZ_MT_RATIO_LO, ITZ_MT_RATIO_HI = 0.35, 0.70     # itz tensile / mortar tensile
ITZ_PE_RATIO_LO, ITZ_PE_RATIO_HI = 0.30, 0.70     # itz modulus / paste modulus


def _itz_strength_ratio_penalty(itz_per_graph, mortar_pred, paste_pred):
    """Apply soft ITZ/mortar and ITZ/paste ratio bands."""
    if itz_per_graph is None or mortar_pred is None or itz_per_graph.numel() == 0:
        return None
    thickness = itz_per_graph[:, 0]
    itz_strength = itz_per_graph[:, 2]   # ITZ tensile/bond strength (MPa)
    itz_modulus = itz_per_graph[:, 3]    # ITZ elastic modulus (GPa)
    has_itz = (thickness > 1.0).to(itz_per_graph.dtype)
    denom = has_itz.sum().clamp(min=1.0)

    def band(ratio, lo, hi):
        return ((torch.relu(lo - ratio) + torch.relu(ratio - hi)) * has_itz).sum() / denom

    mortar_tensile = mortar_pred[:, 1]   # MORTAR_TARGETS slot 1
    pen = band(itz_strength / mortar_tensile.clamp(min=0.1),
               ITZ_MT_RATIO_LO, ITZ_MT_RATIO_HI)
    if paste_pred is not None:
        paste_modulus = paste_pred[:, 2]  # PASTE_TARGETS slot 2 (GPa)
        pen = pen + band(itz_modulus / paste_modulus.clamp(min=0.1),
                         ITZ_PE_RATIO_LO, ITZ_PE_RATIO_HI)
    return pen


def _porosity_ordering_penalty(paste_pred, mortar_pred):
    """Soft prior enforcing paste porosity >= mortar porosity."""
    if paste_pred is None or mortar_pred is None:
        return None
    paste_poro = paste_pred[:, 4]    # PASTE_TARGETS slot 4
    mortar_poro = mortar_pred[:, 4]  # MORTAR_TARGETS slot 4
    return torch.relu(mortar_poro - paste_poro).mean()


# Paste-to-mortar volumetric porosity band.
PMCONS_BAND_LO, PMCONS_BAND_HI = 0.7, 1.4


def _paste_mortar_consistency_penalty(paste_pred, mortar_pred, sand_vf):
    """Penalize paste-to-mortar porosity ratios outside the tolerance band."""
    if paste_pred is None or mortar_pred is None or sand_vf is None:
        return None
    paste_poro = paste_pred[:, 4]
    mortar_poro = mortar_pred[:, 4].clamp(min=1.0e-4)
    implied = paste_poro * (1.0 - sand_vf.reshape(-1).clamp(0.0, 0.95))
    ratio = implied / mortar_poro
    return (torch.relu(PMCONS_BAND_LO - ratio) + torch.relu(ratio - PMCONS_BAND_HI)).mean()


# Tensile and flexural reference bands.
TF_BAND_LO, TF_BAND_HI = 0.6, 1.8                    # default band ([0.6x, 1.8x])
TF_BAND_LO_UHPC, TF_BAND_HI_UHPC = 0.5, 2.5          # released band for fc >= 150 MPa
FT_RATIO_90_150, FT_RATIO_150 = 0.0785, 0.0880
FR_RATIO_50_90, FR_RATIO_90_150, FR_RATIO_150 = 0.124, 0.172, 0.167


def _ft_ref(fc):
    """Piecewise tensile reference (MPa) as a function of predicted compressive fc."""
    ceb = 0.30 * fc.pow(2.0 / 3.0)
    fib = 2.12 * torch.log1p((fc + 8.0) / 10.0)
    high = torch.where(fc >= 150.0, FT_RATIO_150 * fc, FT_RATIO_90_150 * fc)
    ref = torch.where(fc < 50.0, ceb, fib)
    ref = torch.where(fc >= 90.0, high, ref)
    return ref


def _fr_ref(fc):
    """Piecewise flexural reference (MPa) as a function of predicted compressive fc."""
    aci = 0.62 * torch.sqrt(fc)
    high = torch.where(fc >= 150.0, FR_RATIO_150 * fc, FR_RATIO_90_150 * fc)
    ref = torch.where(fc < 50.0, aci, FR_RATIO_50_90 * fc)
    ref = torch.where(fc >= 90.0, high, ref)
    return ref


def _monotonic_tf_penalty(concrete_pred):
    """Penalize tensile and flexural predictions outside broad reference bands."""
    if concrete_pred is None or concrete_pred.size(0) == 0:
        return None
    fc = concrete_pred[:, 0].clamp(min=1.0)
    ft = concrete_pred[:, 1]
    fr = concrete_pred[:, 2]
    ft_ref = _ft_ref(fc)
    fr_ref = _fr_ref(fc)
    # Use the wider band for fc >= 150 MPa.
    uhpc = fc >= 150.0
    lo_m = torch.where(uhpc, fc.new_tensor(TF_BAND_LO_UHPC), fc.new_tensor(TF_BAND_LO))
    hi_m = torch.where(uhpc, fc.new_tensor(TF_BAND_HI_UHPC), fc.new_tensor(TF_BAND_HI))

    def band(x, ref):
        lo, hi = lo_m * ref, hi_m * ref
        return (torch.relu(lo - x) / ref + torch.relu(x - hi) / ref).mean()

    return band(ft, ft_ref) + band(fr, fr_ref)


def _paste_hierarchy_penalty(paste_pred, concrete_mpa, cmask):
    """Penalize paste strength below measured concrete strength."""
    if paste_pred is None or concrete_mpa is None or not bool(cmask.any()):
        return None
    viol = torch.relu(concrete_mpa[cmask] - paste_pred[:, 0][cmask]) / 100.0
    return viol.pow(2).mean()


@torch.no_grad()
def _eval_mortar_rmse(model: nn.Module, loader: DataLoader, device: torch.device):
    """Physical-unit RMSE of the mortar sub-model's compressive output vs. the
    measured mortar strengths, over the masked (supervised) rows only.

    Returns ``(rmse_mpa, n_rows)`` or ``None`` when the model has no mortar
    sub-model (tabular baseline) or no row carries a mortar label.
    """

    model.eval()
    preds, trues = [], []
    for batch in loader:
        batch = batch.to(device)
        model(batch)  # populates model.mortar_pred
        mp = getattr(model, "mortar_pred", None)
        mt = batch.mortar_target
        if mp is None or mt is None:
            return None
        mt0 = mt[:, 0]
        mask = ~torch.isnan(mt0)
        if bool(mask.any()):
            preds.append(mp[mask, 0].detach().cpu())
            trues.append(mt0[mask].detach().cpu())
    if not preds:
        return None
    p = torch.cat(preds)
    t = torch.cat(trues)
    return float((p - t).pow(2).mean().sqrt()), int(p.numel())


@torch.no_grad()
def _eval_concrete_aux(model: nn.Module, loader: DataLoader, device: torch.device):
    """Physical-unit RMSE & R2 of the tensile/flexural capacity heads vs measured
    labels over supervised (non-NaN) rows. Returns ``{name: (rmse, r2, n)}`` or None."""
    model.eval()
    slots = {"tensile": (1, 0), "flexural": (2, 1)}  # name -> (concrete_pred slot, aux col)
    acc = {k: ([], []) for k in slots}
    for batch in loader:
        batch = batch.to(device)
        model(batch)  # populates model.concrete_pred
        cp = getattr(model, "concrete_pred", None)
        cat = getattr(batch, "concrete_aux_target", None)
        if cp is None or cat is None:
            return None
        for k, (slot, col) in slots.items():
            t = cat[:, col]
            mask = ~torch.isnan(t)
            if bool(mask.any()):
                acc[k][0].append(cp[mask, slot].detach().cpu())
                acc[k][1].append(t[mask].detach().cpu())
    out = {}
    for k, (ps, ts) in acc.items():
        if ps:
            p = torch.cat(ps)
            t = torch.cat(ts)
            ss_res = (t - p).pow(2).sum()
            ss_tot = (t - t.mean()).pow(2).sum().clamp_min(1.0e-12)
            out[k] = (
                float((p - t).pow(2).mean().sqrt()),
                float(1.0 - ss_res / ss_tot),
                int(p.numel()),
            )
    return out


def _train_one_model(
    name: str,
    base: nn.Module,
    standardizer: Standardizer,
    train_loader: DataLoader,
    val_loader: DataLoader,
    test_loader: DataLoader,
    device: torch.device,
    epochs: int,
    lr: float,
    checkpoint_dir: Optional[str],
    checkpoint_metadata: Optional[dict] = None,
    mortar_weight: float = 0.0,
    itz_weight: float = 0.0,
    itz_prior_weight: float = 0.0,
    monotonic_weight: float = 0.0,
    paste_weight: float = 0.0,
    paste_hier_weight: float = 0.0,
    paste_shape_weight: float = 0.0,
    tensile_weight: float = 0.0,
    flexural_weight: float = 0.0,
    itz_wb_weight: float = 0.0,
    porosity_order_weight: float = 0.0,
    paste_mortar_consistency_weight: float = 0.0,
    itz_ratio_weight: float = 0.0,
    mortar_stats=None,
    itz_stats=None,
    paste_stats=None,
    concrete_aux_stats=None,
) -> ModelResult:
    model = StrengthHead(base, standardizer).to(device)
    m_mean = mortar_stats[0].to(device) if mortar_stats is not None else None
    m_std = mortar_stats[1].to(device) if mortar_stats is not None else None
    i_mean = itz_stats[0].to(device) if itz_stats is not None else None
    i_std = itz_stats[1].to(device) if itz_stats is not None else None
    p_mean = paste_stats[0].to(device) if paste_stats is not None else None
    p_std = paste_stats[1].to(device) if paste_stats is not None else None
    ca_mean = concrete_aux_stats[0].to(device) if concrete_aux_stats is not None else None
    ca_std = concrete_aux_stats[1].to(device) if concrete_aux_stats is not None else None
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1.0e-4)
    best_val = float("inf")
    best_epoch = -1
    best_state = None
    history: List[Dict[str, float]] = []

    for epoch in range(1, epochs + 1):
        model.train()
        total = 0.0
        mort_total = 0.0
        mort_batches = 0
        itz_total = 0.0
        itz_batches = 0
        prior_total = 0.0
        mono_total = 0.0
        itz_wb_total = 0.0
        poro_order_total = 0.0
        pmcons_total = 0.0
        itz_ratio_total = 0.0
        paste_total = 0.0
        paste_batches = 0
        phier_total = 0.0
        pshape_total = 0.0
        tens_total = 0.0
        tens_batches = 0
        flex_total = 0.0
        flex_batches = 0
        n = 0
        for batch in train_loader:
            batch = batch.to(device)
            pred = model(batch)
            target = batch.concrete_target[:, 0]
            # Ignore rows without a concrete-strength label.
            cmask = ~torch.isnan(target)
            loss = (
                torch.nn.functional.mse_loss(pred[cmask], target[cmask])
                if bool(cmask.any()) else pred.sum() * 0.0
            )
            # Masked mortar supervision.
            mort = getattr(model, "mortar_pred", None)
            mt = batch.mortar_target
            if mortar_weight > 0.0 and mort is not None and mt is not None and m_mean is not None:
                ml = _masked_target_mse(mort, mt, m_mean, m_std)
                if ml is not None:
                    loss = loss + mortar_weight * ml
                    mort_total += float(ml.detach().cpu())
                    mort_batches += 1

            # Masked ITZ supervision.
            itzg = getattr(model, "itz_pred_per_graph", None)
            it = batch.itz_target_graph
            if itz_weight > 0.0 and itzg is not None and it is not None and i_mean is not None:
                il = _masked_target_mse(itzg, it, i_mean, i_std)
                if il is not None:
                    loss = loss + itz_weight * il
                    itz_total += float(il.detach().cpu())
                    itz_batches += 1

            # ITZ range prior.
            if itz_prior_weight > 0.0:
                pen = _itz_range_penalty(getattr(model, "itz_pred", None))
                if pen is not None:
                    loss = loss + itz_prior_weight * pen
                    prior_total += float(pen.detach().cpu())

            # Masked paste supervision.
            pp = getattr(model, "paste_pred", None)
            pt = batch.paste_target
            if paste_weight > 0.0 and pp is not None and pt is not None and p_mean is not None:
                pl = _masked_target_mse(pp, pt, p_mean, p_std)
                if pl is not None:
                    loss = loss + paste_weight * pl
                    paste_total += float(pl.detach().cpu())
                    paste_batches += 1

            # Abrams-shape prior for paste strength.
            if paste_shape_weight > 0.0:
                spen = getattr(model, "paste_abrams_penalty", None)
                if spen is not None:
                    loss = loss + paste_shape_weight * spen
                    pshape_total += float(spen.detach().cpu())

            # Paste strength should not fall below concrete strength.
            if paste_hier_weight > 0.0 and pp is not None:
                conc_mpa = model.standardizer.decode(pred.detach())
                hpen = _paste_hierarchy_penalty(pp, conc_mpa, cmask)
                if hpen is not None:
                    loss = loss + paste_hier_weight * hpen
                    phier_total += float(hpen.detach().cpu())

            # Masked tensile and flexural supervision.
            cp = getattr(model, "concrete_pred", None)
            cat = getattr(batch, "concrete_aux_target", None)
            if tensile_weight > 0.0 and cp is not None and cat is not None and ca_mean is not None:
                tl = _masked_target_mse(cp[:, 1:2], cat[:, 0:1], ca_mean[0:1], ca_std[0:1])
                if tl is not None:
                    loss = loss + tensile_weight * tl
                    tens_total += float(tl.detach().cpu())
                    tens_batches += 1
            if flexural_weight > 0.0 and cp is not None and cat is not None and ca_mean is not None:
                fl = _masked_target_mse(cp[:, 2:3], cat[:, 1:2], ca_mean[1:2], ca_std[1:2])
                if fl is not None:
                    loss = loss + flexural_weight * fl
                    flex_total += float(fl.detach().cpu())
                    flex_batches += 1

            # Strength-consistency prior.
            if monotonic_weight > 0.0 and cp is not None:
                mtf = _monotonic_tf_penalty(cp)
                if mtf is not None:
                    loss = loss + monotonic_weight * mtf
                    mono_total += float(mtf.detach().cpu())

            # ITZ thickness and porosity should not fall as w/b rises.
            if itz_wb_weight > 0.0:
                iwb = getattr(model, "itz_wb_penalty", None)
                if iwb is not None:
                    loss = loss + itz_wb_weight * iwb
                    itz_wb_total += float(iwb.detach().cpu())

            # Paste porosity should exceed mortar porosity; ITZ ordering is structural.
            if porosity_order_weight > 0.0:
                pop = _porosity_ordering_penalty(
                    getattr(model, "paste_pred", None),
                    getattr(model, "mortar_pred", None),
                )
                if pop is not None:
                    loss = loss + porosity_order_weight * pop
                    poro_order_total += float(pop.detach().cpu())

            # Paste-mortar volumetric consistency.
            if paste_mortar_consistency_weight > 0.0:
                pmc = _paste_mortar_consistency_penalty(
                    getattr(model, "paste_pred", None),
                    getattr(model, "mortar_pred", None),
                    getattr(batch, "sand_volume_fraction", None),
                )
                if pmc is not None:
                    loss = loss + paste_mortar_consistency_weight * pmc
                    pmcons_total += float(pmc.detach().cpu())

            # ITZ-to-mortar strength and ITZ-to-paste modulus bands.
            if itz_ratio_weight > 0.0:
                irp = _itz_strength_ratio_penalty(
                    getattr(model, "itz_pred_per_graph", None),
                    getattr(model, "mortar_pred", None),
                    getattr(model, "paste_pred", None),
                )
                if irp is not None:
                    loss = loss + itz_ratio_weight * irp
                    itz_ratio_total += float(irp.detach().cpu())
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            total += float(loss.detach().cpu())
            n += 1

        val = evaluate_strength(model, val_loader, device)
        train_loss = total / max(n, 1)
        history.append({
            "epoch": epoch,
            "train_loss": train_loss,
            "val_rmse": val["rmse"],
            "val_r2": val["r2"],
            "mortar_mse": mort_total / max(mort_batches, 1),
            "itz_mse": itz_total / max(itz_batches, 1),
            "paste_mse": paste_total / max(paste_batches, 1),
            "tensile_mse": tens_total / max(tens_batches, 1),
            "flexural_mse": flex_total / max(flex_batches, 1),
            "monotonic": mono_total / max(n, 1),
            "itz_wb": itz_wb_total / max(n, 1),
            "porosity_order": poro_order_total / max(n, 1),
            "paste_mortar_consistency": pmcons_total / max(n, 1),
            "itz_strength_ratio": itz_ratio_total / max(n, 1),
            "paste_shape": pshape_total / max(n, 1),
        })
        # Rewrite the history CSV every epoch so a crash (e.g. late CUDA OOM)
        # cannot lose the recorded curve.
        if checkpoint_dir is not None:
            os.makedirs(checkpoint_dir, exist_ok=True)
            history_path = os.path.join(checkpoint_dir, f"training_history_{name}.csv")
            with open(history_path, "w", newline="") as fh:
                writer = csv.DictWriter(fh, fieldnames=list(history[0].keys()))
                writer.writeheader()
                writer.writerows(history)
        if val["rmse"] < best_val:
            best_val = val["rmse"]
            best_epoch = epoch
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
        if epoch == 1 or epoch == epochs or epoch % max(1, epochs // 5) == 0:
            mort_msg = (
                f" mortar_mse={mort_total / max(mort_batches, 1):.4f}"
                if mortar_weight > 0.0 and mort_batches > 0
                else ""
            )
            itz_msg = (
                f" itz_mse={itz_total / max(itz_batches, 1):.4f}"
                if itz_weight > 0.0 and itz_batches > 0
                else ""
            )
            prior_msg = (
                f" itz_prior={prior_total / max(n, 1):.4f}"
                if itz_prior_weight > 0.0
                else ""
            )
            paste_msg = (
                f" paste_mse={paste_total / max(paste_batches, 1):.4f}"
                if paste_weight > 0.0 and paste_batches > 0
                else ""
            )
            phier_msg = (
                f" paste>=conc={phier_total / max(n, 1):.4f}"
                if paste_hier_weight > 0.0
                else ""
            )
            pshape_msg = (
                f" paste_shape={pshape_total / max(n, 1):.4f}"
                if paste_shape_weight > 0.0
                else ""
            )
            tens_msg = (
                f" tens_mse={tens_total / max(tens_batches, 1):.4f}"
                if tensile_weight > 0.0 and tens_batches > 0
                else ""
            )
            flex_msg = (
                f" flex_mse={flex_total / max(flex_batches, 1):.4f}"
                if flexural_weight > 0.0 and flex_batches > 0
                else ""
            )
            mono_msg = (
                f" mono={mono_total / max(n, 1):.4f}"
                if monotonic_weight > 0.0
                else ""
            )
            itz_wb_msg = (
                f" itz_wb={itz_wb_total / max(n, 1):.4f}"
                if itz_wb_weight > 0.0
                else ""
            )
            poro_msg = (
                f" poro_order={poro_order_total / max(n, 1):.4f}"
                if porosity_order_weight > 0.0
                else ""
            )
            pmc_msg = (
                f" pm_cons={pmcons_total / max(n, 1):.4f}"
                if paste_mortar_consistency_weight > 0.0
                else ""
            )
            iratio_msg = (
                f" itz_ratio={itz_ratio_total / max(n, 1):.4f}"
                if itz_ratio_weight > 0.0
                else ""
            )
            print(
                f"[{name}] epoch={epoch:03d} train_loss={train_loss:.4f} "
                f"val_rmse={val['rmse']:.3f} val_r2={val['r2']:.3f}"
                f"{mort_msg}{itz_msg}{prior_msg}{paste_msg}{phier_msg}{pshape_msg}"
                f"{tens_msg}{flex_msg}{mono_msg}{itz_wb_msg}{poro_msg}{pmc_msg}{iratio_msg}"
            )

    if best_state is not None:
        model.load_state_dict(best_state)
    test = evaluate_strength(model, test_loader, device)
    print(
        f"[{name}] test_rmse={test['rmse']:.3f} "
        f"test_mae={test['mae']:.3f} test_r2={test['r2']:.3f}"
    )
    reported_metrics = {"compressive": test}
    if mortar_weight > 0.0:
        mort_eval = _eval_mortar_rmse(model, test_loader, device)
        if mort_eval is not None:
            reported_metrics["mortar_compressive"] = {
                "rmse": mort_eval[0],
                "n": mort_eval[1],
            }
            print(
                f"[{name}] test mortar_compressive RMSE={mort_eval[0]:.3f} MPa "
                f"over {mort_eval[1]} supervised rows"
            )
    if tensile_weight > 0.0 or flexural_weight > 0.0:
        aux_eval = _eval_concrete_aux(model, test_loader, device)
        if aux_eval:
            for k, (rmse, r2, nrow) in aux_eval.items():
                reported_metrics[k] = {"rmse": rmse, "r2": r2, "n": nrow}
                print(
                    f"[{name}] test {k} RMSE={rmse:.3f} MPa  R2={r2:.3f}  "
                    f"over {nrow} supervised rows"
                )
    pred, true = collect_strength_predictions(model, test_loader, device)
    predictions = {"pred": pred.numpy()[:, None], "true": true.numpy()[:, None]}
    checkpoint_path = None
    if checkpoint_dir is not None:
        os.makedirs(checkpoint_dir, exist_ok=True)
        checkpoint_path = os.path.join(checkpoint_dir, f"{name}.pt")
        torch.save(
            {
                "format_version": 2,
                "model_name": name,
                "best_epoch": best_epoch,
                "best_val_rmse": best_val,
                "test_metrics": reported_metrics,
                "target_mean": float(standardizer.mean),
                "target_std": float(standardizer.std),
                "base_model_state_dict": model.base.state_dict(),
                "run_metadata": checkpoint_metadata or {},
            },
            checkpoint_path,
        )
        print(f"[{name}] checkpoint saved to {checkpoint_path}")
        if history:
            print(
                f"[{name}] per-epoch history saved to "
                f"{os.path.join(checkpoint_dir, f'training_history_{name}.csv')}"
            )
    return {
        "metrics": test,
        "history": history,
        "predictions": predictions,
        "checkpoint_path": checkpoint_path,
    }


def _save_figures(
    results: Dict[str, ModelResult],
    model_order: Sequence[str],
    sample_graph,
    figure_dir: str,
) -> None:
    os.makedirs(figure_dir, exist_ok=True)
    histories = {name: results[name]["history"] for name in model_order}
    predictions = {name: results[name]["predictions"]["pred"] for name in model_order}
    truths = results[model_order[0]]["predictions"]["true"]
    r2_by_model = {
        name: np.asarray([results[name]["metrics"]["r2"]], dtype=float)
        for name in model_order
    }

    plot_parity(
        predictions=predictions,
        truths=truths,
        model_order=model_order,
        target_names=(TARGET_NAME,),
        save_path=os.path.join(figure_dir, "real_data_parity.png"),
    )
    plot_r2_bars(
        r2_by_model=r2_by_model,
        model_order=model_order,
        target_names=(TARGET_NAME,),
        save_path=os.path.join(figure_dir, "real_data_r2.png"),
    )
    plot_training_curves(
        histories=histories,
        model_order=model_order,
        save_path=os.path.join(figure_dir, "real_data_training_curves.png"),
    )
    plot_rve_graph(
        sample_graph,
        save_path=os.path.join(figure_dir, "real_data_example_rve.png"),
    )

    print(f"\nFigures saved under {figure_dir}:")
    for fname in (
        "real_data_parity.png",
        "real_data_r2.png",
        "real_data_training_curves.png",
        "real_data_example_rve.png",
    ):
        path = os.path.join(figure_dir, fname)
        if os.path.exists(path):
            print(f"  {fname:<32} ({os.path.getsize(path) / 1024.0:.1f} KB)")


def _seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _sha256(path: str) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--data",
        default=os.path.join("data", "HiConGNN_processed_database.csv"),
        help="Processed publication CSV.",
    )
    parser.add_argument(
        "--split",
        default=None,
        help="Optional split manifest. Omit for seeded row-random 70/15/15.",
    )
    parser.add_argument("--epochs", type=int, default=120)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--lr", type=float, default=1.0e-3)
    parser.add_argument("--seed", type=int, default=17)
    parser.add_argument(
        "--model",
        choices=("both", "hicon", "ann"),
        default="both",
        help="Model or models to train.",
    )
    parser.add_argument(
        "--figure-dir",
        default=os.path.join(OUTPUT_ROOT, "figures"),
        help="Output directory for diagnostic figures.",
    )
    parser.add_argument(
        "--checkpoint-dir",
        default=os.path.join(OUTPUT_ROOT, "checkpoints"),
        help="Output directory for best-validation checkpoints.",
    )
    parser.add_argument("--no-checkpoints", action="store_true")
    parser.add_argument("--no-figures", action="store_true")
    parser.add_argument("--cpu", action="store_true")
    args = parser.parse_args()
    schema = DEFAULT_SCHEMA
    _seed_everything(args.seed)

    if args.cpu:
        device = torch.device("cpu")
    else:
        if torch.cuda.is_available():
            device = torch.device("cuda")
        elif getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
            device = torch.device("mps")
        else:
            device = torch.device("cpu")
    df = read_table(args.data)
    if args.split:
        train_idx, val_idx, test_idx = load_split_indices(args.split, len(df))
        split_strategy = "manifest"
    else:
        train_idx, val_idx, test_idx = random_split_indices(len(df), args.seed)
        split_strategy = "row_random"
    train_y = torch.tensor(
        df.iloc[train_idx][strength_column(df)].to_numpy(dtype="float32"),
        dtype=torch.float32,
    )
    standardizer = Standardizer.from_values(train_y)

    print(
        f"device={device} rows={len(df)} split={split_strategy} "
        f"train={len(train_idx)} val={len(val_idx)} test={len(test_idx)}"
    )
    print(
        f"target mean={float(standardizer.mean):.3f} MPa "
        f"std={float(standardizer.std):.3f} MPa"
    )

    train_ds = ConcreteMixDataset(
        df, train_idx, seed=args.seed, target_standardizer=standardizer, schema=schema
    )
    val_ds = ConcreteMixDataset(
        df, val_idx, seed=args.seed + 10, target_standardizer=standardizer, schema=schema
    )
    test_ds = ConcreteMixDataset(
        df, test_idx, seed=args.seed + 20, target_standardizer=standardizer, schema=schema
    )
    loader_rng = torch.Generator().manual_seed(args.seed)
    train_loader = DataLoader(
        train_ds, batch_size=args.batch_size, shuffle=True, collate_fn=collate_multiscale,
        generator=loader_rng,
        # BatchNorm requires at least two samples during training.
        drop_last=True,
    )
    val_loader = DataLoader(
        val_ds, batch_size=args.batch_size, shuffle=False, collate_fn=collate_multiscale
    )
    test_loader = DataLoader(
        test_ds, batch_size=args.batch_size, shuffle=False, collate_fn=collate_multiscale
    )

    models: Dict[str, nn.Module] = {}
    if args.model in ("both", "hicon"):
        models["hierarchical"] = IntegratedMultiscaleModel(schema=schema)
    if args.model in ("both", "ann"):
        models["tabular_ann"] = TabularANN(in_dim=TABULAR_DIM, schema=schema)

    if "hierarchical" in models:
        fit_encoder_standardizers(models["hierarchical"], train_ds.samples)
        print("fitted input standardization: mortar nodes/globals, mesoscale globals, ITZ mix vector")

    # Per-target label scales (masked) for the mortar / ITZ supervision losses.
    mortar_stats = _masked_target_stats(train_ds.samples, "mortar_target", len(schema.mortar_targets))
    itz_stats = _masked_target_stats(train_ds.samples, "itz_target_graph", len(schema.itz_targets))
    paste_stats = _masked_target_stats(train_ds.samples, "paste_target", len(schema.paste_targets))
    concrete_aux_stats = _masked_target_stats(train_ds.samples, "concrete_aux_target", 2)
    n_mortar_lab = int(sum(bool((~torch.isnan(s.mortar_target)).any()) for s in train_ds.samples))
    n_itz_lab = int(sum(bool((~torch.isnan(s.itz_target_graph)).any()) for s in train_ds.samples))
    n_paste_lab = int(sum(bool((~torch.isnan(s.paste_target)).any()) for s in train_ds.samples))
    n_tens_lab = int(sum(bool(~torch.isnan(s.concrete_aux_target[0])) for s in train_ds.samples))
    n_flex_lab = int(sum(bool(~torch.isnan(s.concrete_aux_target[1])) for s in train_ds.samples))
    print(f"supervised train rows -> mortar: {n_mortar_lab}  ITZ: {n_itz_lab}  paste: {n_paste_lab}  "
          f"tensile: {n_tens_lab}  flexural: {n_flex_lab}")

    checkpoint_dir = None if args.no_checkpoints else args.checkpoint_dir
    run_metadata = {
        "seed": args.seed,
        "epochs": args.epochs,
        "batch_size": args.batch_size,
        "learning_rate": args.lr,
        "data_sha256": _sha256(args.data),
        "split_strategy": split_strategy,
        "split_seed": args.seed if split_strategy == "row_random" else None,
        "split_indices_sha256": split_indices_sha256(train_idx, val_idx, test_idx),
        "split_counts": {
            "train": len(train_idx),
            "validation": len(val_idx),
            "test": len(test_idx),
        },
        "loss_weights": PAPER_LOSS_WEIGHTS,
        "python": sys.version.split()[0],
        "torch": str(torch.__version__),
        "numpy": np.__version__,
    }
    if args.split:
        run_metadata["split_sha256"] = _sha256(args.split)
    results: Dict[str, ModelResult] = {}
    for name, model in models.items():
        is_hicon = name == "hierarchical"
        results[name] = _train_one_model(
            name,
            model,
            standardizer,
            train_loader,
            val_loader,
            test_loader,
            device,
            args.epochs,
            args.lr,
            checkpoint_dir,
            checkpoint_metadata=run_metadata,
            mortar_weight=PAPER_LOSS_WEIGHTS["mortar"] if is_hicon else 0.0,
            itz_weight=PAPER_LOSS_WEIGHTS["itz"] if is_hicon else 0.0,
            itz_prior_weight=PAPER_LOSS_WEIGHTS["itz_prior"] if is_hicon else 0.0,
            monotonic_weight=PAPER_LOSS_WEIGHTS["strength_consistency"],
            paste_weight=PAPER_LOSS_WEIGHTS["paste"] if is_hicon else 0.0,
            paste_hier_weight=PAPER_LOSS_WEIGHTS["paste_hierarchy"] if is_hicon else 0.0,
            paste_shape_weight=PAPER_LOSS_WEIGHTS["paste_shape"] if is_hicon else 0.0,
            tensile_weight=PAPER_LOSS_WEIGHTS["tensile"],
            flexural_weight=PAPER_LOSS_WEIGHTS["flexural"],
            itz_wb_weight=PAPER_LOSS_WEIGHTS["itz_wb"] if is_hicon else 0.0,
            porosity_order_weight=PAPER_LOSS_WEIGHTS["porosity_order"] if is_hicon else 0.0,
            paste_mortar_consistency_weight=(
                PAPER_LOSS_WEIGHTS["paste_mortar_consistency"] if is_hicon else 0.0
            ),
            itz_ratio_weight=PAPER_LOSS_WEIGHTS["itz_ratio"] if is_hicon else 0.0,
            mortar_stats=mortar_stats if is_hicon else None,
            itz_stats=itz_stats if is_hicon else None,
            paste_stats=paste_stats if is_hicon else None,
            concrete_aux_stats=concrete_aux_stats,
        )

    print("\nSummary")
    print(f"{'model':<16} {'RMSE':>8} {'MAE':>8} {'R2':>8}")
    for name, result in results.items():
        stats = result["metrics"]
        print(f"{name:<16} {stats['rmse']:8.3f} {stats['mae']:8.3f} {stats['r2']:8.3f}")
    print("\nCheckpoints")
    for name, result in results.items():
        path = result.get("checkpoint_path")
        if path:
            print(f"{name:<16} {path}")

    if not args.no_figures:
        _save_figures(
            results,
            list(models.keys()),
            sample_graph=train_ds[0].graph,
            figure_dir=args.figure_dir,
        )


if __name__ == "__main__":
    main()
