"""Utilities for training HiCon-GNN models on real concrete mix datasets."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Union

import pandas as pd
import torch
from torch import Tensor, nn
from torch.utils.data import DataLoader, Dataset

from .categoricals import (
    CEMENT_TYPES,
    COARSE_AGGREGATE_TYPES,
    CURING_REGIMES,
    FIBRE_TYPES,
    FLY_ASH_TYPES,
    canonical_cement_type,
    canonical_coarse_aggregate_type,
    canonical_curing_regime,
    canonical_fibre_type,
    canonical_fly_ash_type,
    one_hot,
    type_id,
)
from .data import ConcreteSample
from .graph_generator import (
    MixDesign,
    MortarMixDesign,
    binder_chem_features,
    generate_concrete_graph,
    generate_mortar_graph,
    mortar_volume_fraction,
)
from .schema import DEFAULT_SCHEMA, MORTAR_NODE_TYPE_PASTE, SchemaSpec


TARGET_NAME = "compressive_strength_MPa"

NORM_CEMENT = "cement_kg_m3"
NORM_SLAG = "slag_kg_m3"
NORM_FLY_ASH = "fly_ash_kg_m3"
NORM_SILICA_FUME = "silica_fume_kg_m3"
NORM_METAKAOLIN = "metakaolin_kg_m3"
NORM_LIMESTONE = "limestone_powder_kg_m3"
NORM_OTHER_SCM = "other_scm_kg_m3"
# Glass powder and fly-ash class.
NORM_GLASS_POWDER = "glass_powder_kg_m3"
NORM_FLY_ASH_TYPE = "fly_ash_type_norm"
NORM_WATER = "water_kg_m3"
NORM_SUPERPLASTICIZER = "superplasticizer_kg_m3"
NORM_COARSE = "coarse_aggregate_kg_m3"
NORM_FINE = "fine_aggregate_kg_m3"
NORM_FIBRE_CONTENT = "fibre_content_kg_m3"
NORM_FIBRE_LENGTH = "fibre_length_mm"
NORM_FIBRE_DIAMETER = "fibre_diameter_mm"
NORM_FIBRE_TENSILE = "fibre_tensile_strength_mpa"
NORM_FIBRE_MODULUS = "fibre_modulus_gpa"
NORM_MAX_COARSE = "max_coarse_aggregate_size_mm"
NORM_MAX_FINE = "max_fine_aggregate_size_mm"
NORM_CURING_TEMP = "curing_temperature_c"
NORM_AGE = "age_days"
NORM_STRENGTH = "compressive_strength_mpa"
# Rich optional inputs: cement/SCM chemistry and related numeric descriptors.
NORM_CEMENT_CAO = "cement_CaO_pct"
NORM_CEMENT_SIO2 = "cement_SiO2_pct"
NORM_CEMENT_AL2O3 = "cement_Al2O3_pct"
NORM_CEMENT_FE2O3 = "cement_Fe2O3_pct"
NORM_CEMENT_MGO = "cement_MgO_pct"
NORM_CEMENT_SO3 = "cement_SO3_pct"
NORM_CEMENT_ALKALI = "cement_alkali_pct"
NORM_CEMENT_LOI = "cement_LOI_pct"
NORM_SCM_CAO = "scm_CaO_pct"
NORM_SCM_SIO2 = "scm_SiO2_pct"
NORM_SCM_AL2O3 = "scm_Al2O3_pct"
NORM_SCM_FE2O3 = "scm_Fe2O3_pct"
NORM_SCM_MGO = "scm_MgO_pct"
NORM_SCM_LOI = "scm_LOI_pct"
NORM_CEMENT_GRADE = "cement_grade_mpa"
NORM_CURING_HUMIDITY = "curing_humidity_pct"
NORM_SPECIMEN_SIZE = "specimen_size_mm"
# Recycled-coarse-aggregate replacement (% of coarse that is RCA). Converted to a
# 0-1 fraction for the aggregate-node feature; fed raw (%) to the tabular baseline.
NORM_RCA_PCT = "rca_replacement_pct"
# Binder phases, SCM fineness, and sand density.
NORM_CEMENT_C3S = "cement_C3S_pct"
NORM_CEMENT_C2S = "cement_C2S_pct"
NORM_CEMENT_C3A = "cement_C3A_pct"
NORM_CEMENT_C4AF = "cement_C4AF_pct"
NORM_SCM_SSA = "scm_specific_surface_m2_kg"
NORM_FINE_SG = "fine_aggregate_specific_gravity"

# Aggregate type and measured properties.
NORM_COARSE_AGG_TYPE = "coarse_aggregate_type_norm"
NORM_COARSE_AGG_SG = "coarse_aggregate_specific_gravity"
NORM_COARSE_AGG_WA = "coarse_aggregate_water_absorption_pct"
NORM_COARSE_AGG_CS = "coarse_aggregate_crushing_strength_mpa"

# Canonical free-text categoricals (string columns; one-hot for the TabularANN).
NORM_CEMENT_TYPE = "cement_type_norm"
NORM_FIBRE_TYPE = "fibre_type_norm"
NORM_CURING_REGIME = "curing_regime_norm"

# Maskable rich columns fed (value + observed-mask channel) to the TabularANN.
RICH_TABULAR_COLUMNS = (
    NORM_CEMENT_CAO,
    NORM_CEMENT_SIO2,
    NORM_CEMENT_AL2O3,
    NORM_CEMENT_FE2O3,
    NORM_CEMENT_MGO,
    NORM_CEMENT_SO3,
    NORM_CEMENT_ALKALI,
    NORM_CEMENT_LOI,
    NORM_SCM_CAO,
    NORM_SCM_SIO2,
    NORM_SCM_AL2O3,
    NORM_SCM_FE2O3,
    NORM_SCM_MGO,
    NORM_SCM_LOI,
    NORM_CEMENT_GRADE,
    NORM_CURING_HUMIDITY,
    NORM_SPECIMEN_SIZE,
    NORM_RCA_PCT,
    NORM_CEMENT_C3S,
    NORM_CEMENT_C2S,
    NORM_CEMENT_C3A,
    NORM_CEMENT_C4AF,
    NORM_SCM_SSA,
    NORM_FINE_SG,
)

# One-hot categorical blocks appended to the tabular vector (after the rich block).
CATEGORICAL_TABULAR = (
    (NORM_CEMENT_TYPE, canonical_cement_type, CEMENT_TYPES, "other"),
    (NORM_FIBRE_TYPE, canonical_fibre_type, FIBRE_TYPES, "none"),
    (NORM_CURING_REGIME, canonical_curing_regime, CURING_REGIMES, "other"),
    (NORM_COARSE_AGG_TYPE, canonical_coarse_aggregate_type, COARSE_AGGREGATE_TYPES, "natural"),
)

# OPC-typical fallbacks used for the GNN paste chemistry when a row does not
# carry measured oxides.
_DEFAULT_CEMENT_CHEM = (64.0, 21.0, 5.0, 3.0)       # CaO, SiO2, Al2O3, Fe2O3
_DEFAULT_CEMENT_CHEM_EXT = (2.0, 2.5, 0.6)          # MgO, SO3, Na2O-eq alkali
_DEFAULT_SCM_CHEM = (4.0, 55.0)                     # CaO, SiO2
_DEFAULT_SCM_CHEM_EXT = (15.0, 5.0, 2.0, 3.0)       # Al2O3, Fe2O3, MgO, LOI

TableInput = Union[str, Path, pd.DataFrame]


def read_table(data: TableInput) -> pd.DataFrame:
    if isinstance(data, pd.DataFrame):
        return data.copy()
    path = Path(data)
    if path.suffix.lower() != ".csv":
        raise ValueError(f"Expected a CSV file, got {path}")
    return pd.read_csv(path)


def row_value(row, name: str, default: float = 0.0) -> float:
    value = row.get(name, default)
    if pd.isna(value):
        return default
    return float(value)


def row_value_optional(row, name: str) -> Optional[float]:
    """Return a numeric value or ``None`` when it is unavailable."""

    if name not in row.index:
        return None
    value = row[name]
    if pd.isna(value):
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _chem_or_default(row, name: str, default: float) -> float:
    """Measured oxide value when the column is present & non-NaN, else ``default``."""
    value = row_value_optional(row, name)
    return default if value is None else value


def _category(row, name: str, canon_fn, vocab, default: str) -> str:
    """Return a canonical categorical value."""
    if name not in row.index:
        return default
    value = row[name]
    if value is None or (isinstance(value, float) and pd.isna(value)):
        return default
    text = str(value).strip()
    if text == "" or text.lower() == "nan":
        return default
    return text if text in vocab else canon_fn(text)


@dataclass
class Standardizer:
    mean: Tensor
    std: Tensor

    @classmethod
    def from_values(cls, values: Tensor) -> "Standardizer":
        # Ignore NaN targets (e.g. pure-paste rows with no concrete strength).
        values = values[~torch.isnan(values)]
        return cls(values.mean(), values.std().clamp_min(1.0e-6))

    def encode(self, values: Tensor) -> Tensor:
        mean = self.mean.to(device=values.device, dtype=values.dtype)
        std = self.std.to(device=values.device, dtype=values.dtype)
        return (values - mean) / std

    def decode(self, values: Tensor) -> Tensor:
        mean = self.mean.to(device=values.device, dtype=values.dtype)
        std = self.std.to(device=values.device, dtype=values.dtype)
        return values * std + mean


def strength_column(df: pd.DataFrame) -> str:
    if NORM_STRENGTH not in df.columns:
        raise ValueError(f"Missing required column: {NORM_STRENGTH}")
    return NORM_STRENGTH


# Target name to CSV column mappings. Unlisted targets remain NaN.
_MORTAR_TARGET_COLS = {
    "mortar_compressive_strength_MPa": "mortar_compressive_strength_mpa",
    "mortar_tensile_strength_MPa": "mortar_tensile_strength_mpa",
    "mortar_flexural_strength_MPa": "mortar_flexural_strength_mpa",
    "mortar_elastic_modulus_GPa": "mortar_elastic_modulus_gpa",
    "mortar_porosity": "mortar_porosity",
}
_ITZ_TARGET_COLS = {
    "itz_thickness_um": "itz_thickness_um",
    "itz_porosity": "itz_porosity",
    "itz_strength_MPa": "itz_strength_mpa",
    "itz_elastic_modulus_GPa": "itz_elastic_modulus_gpa",
}
# Paste target columns.
_PASTE_TARGET_COLS = {
    "paste_compressive_strength_MPa": "paste_compressive_strength_mpa",
    "paste_tensile_strength_MPa": "paste_tensile_strength_mpa",
    "paste_elastic_modulus_GPa": "paste_elastic_modulus_gpa",
    "paste_fracture_energy_J_m2": "paste_fracture_energy_j_m2",
    "paste_porosity": "paste_porosity",
}


def _target_vector(row, names, col_map) -> Tensor:
    """Physical-unit label vector; ``NaN`` marks every unmeasured slot/row."""
    t = torch.full((len(names),), float("nan"), dtype=torch.float32)
    for i, name in enumerate(names):
        col = col_map.get(name)
        if col is not None:
            v = row_value_optional(row, col)
            if v is not None:
                t[i] = v
    return t


def _normalize_mixed_porosity(t: Tensor, names) -> Tensor:
    """Convert mixed fraction and percent porosity labels to fractions."""
    for i, name in enumerate(names):
        if "porosity" in name and not torch.isnan(t[i]) and t[i] >= 1.0:
            t[i] = t[i] / 100.0
    return t


def _validate_fraction_porosity(t: Tensor, names) -> Tensor:
    """Require porosity labels to be stored as fractions."""
    for i, name in enumerate(names):
        if "porosity" not in name or torch.isnan(t[i]):
            continue
        if not 0.0 <= float(t[i]) <= 1.0:
            raise ValueError(f"{name} must be stored as a fraction in [0, 1]")
    return t


def mortar_target_from_row(row, schema: SchemaSpec) -> Tensor:
    """Build the mortar target vector."""
    t = _target_vector(row, schema.mortar_targets, _MORTAR_TARGET_COLS)
    return _validate_fraction_porosity(t, schema.mortar_targets)


def itz_target_from_row(row, schema: SchemaSpec) -> Tensor:
    """Build the ITZ target vector."""
    t = _target_vector(row, schema.itz_targets, _ITZ_TARGET_COLS)
    return _normalize_mixed_porosity(t, schema.itz_targets)


# Concrete tensile/flexural label columns for the masked capacity-head losses
# (slots align with CONCRETE_TARGETS[1:3] = tensile, flexural).
CONCRETE_AUX_TARGET_NAMES = ("tensile_strength_MPa", "flexural_strength_MPa")
_CONCRETE_AUX_TARGET_COLS = {
    "tensile_strength_MPa": "tensile_strength_mpa",
    "flexural_strength_MPa": "flexural_strength_mpa",
}


def concrete_aux_target_from_row(row) -> Tensor:
    """Per-graph concrete tensile/flexural label vector (physical MPa, NaN where
    unmeasured) for the masked tensile/flexural capacity-head losses."""
    return _target_vector(row, CONCRETE_AUX_TARGET_NAMES, _CONCRETE_AUX_TARGET_COLS)


def paste_target_from_row(row, schema: SchemaSpec) -> Tensor:
    """Build the paste target vector."""
    t = _target_vector(row, schema.paste_targets, _PASTE_TARGET_COLS)
    return _normalize_mixed_porosity(t, schema.paste_targets)


def _porosity_slot(names) -> Optional[int]:
    return next((i for i, n in enumerate(names) if "porosity" in n), None)


def derive_paste_porosity(paste_target: Tensor, mortar_target: Tensor,
                          sand_vf: float, schema: SchemaSpec) -> Tensor:
    """Derive missing paste porosity from mortar porosity and sand fraction."""
    p_idx = _porosity_slot(schema.paste_targets)
    m_idx = _porosity_slot(schema.mortar_targets)
    if p_idx is None or m_idx is None or not torch.isnan(paste_target[p_idx]):
        return paste_target
    m_poro = mortar_target[m_idx]
    if bool(torch.isnan(m_poro)) or sand_vf >= 0.95:
        return paste_target
    derived = float(m_poro) / (1.0 - sand_vf)
    if derived > 0.55:
        return paste_target
    out = paste_target.clone()
    out[p_idx] = min(max(derived, 0.02), 0.50)
    return out


def row_to_mix(row) -> MixDesign:
    cement = row_value(row, NORM_CEMENT)
    slag = row_value(row, NORM_SLAG)
    fly_ash = row_value(row, NORM_FLY_ASH)
    silica_fume = row_value(row, NORM_SILICA_FUME)
    metakaolin = row_value(row, NORM_METAKAOLIN)
    limestone = row_value(row, NORM_LIMESTONE)
    other_scm = row_value(row, NORM_OTHER_SCM)
    glass_powder = row_value(row, NORM_GLASS_POWDER)
    water = row_value(row, NORM_WATER)
    coarse = row_value(row, NORM_COARSE)
    superplasticizer = row_value(row, NORM_SUPERPLASTICIZER)
    age = row_value(row, NORM_AGE)

    # Glass powder is an explicit supplementary binder.
    extra_scm = metakaolin + limestone + other_scm + glass_powder
    binder = max(cement + slag + fly_ash + silica_fume + extra_scm, 1.0)
    scm = slag + fly_ash + silica_fume + extra_scm
    primary_count = sum(v > 0.0 for v in (slag, fly_ash, silica_fume))
    if primary_count > 1 or (primary_count >= 1 and extra_scm > 0.0):
        scm_type = 3
    elif slag > 0.0:
        scm_type = 1
    elif fly_ash > 0.0:
        scm_type = 2
    elif silica_fume > 0.0 or extra_scm > 0.0:
        scm_type = 4
    else:
        scm_type = 0

    aggregate_volume_fraction = max(0.0, min(0.70, coarse / 2650.0))
    water_to_binder = water / binder

    # Sand fraction from constituent absolute volumes.
    fine_agg = row_value(row, NORM_FINE)
    rho_sand = _chem_or_default(row, NORM_FINE_SG, 2.65)
    v_paste = (cement / 3.15 + slag / 2.90 + fly_ash / 2.30 + silica_fume / 2.20
               + metakaolin / 2.50 + limestone / 2.70 + other_scm / 2.50
               + glass_powder / 2.50 + water / 1.00)
    v_sand = fine_agg / max(rho_sand, 1.0)
    sand_vf = float(v_sand / (v_sand + v_paste)) if (v_sand + v_paste) > 0.0 else 0.0

    fibre_content = row_value(row, NORM_FIBRE_CONTENT)
    fibre_length = row_value(row, NORM_FIBRE_LENGTH)
    fibre_diameter = row_value(row, NORM_FIBRE_DIAMETER)
    fibre_tensile = row_value(row, NORM_FIBRE_TENSILE)
    fibre_modulus = row_value(row, NORM_FIBRE_MODULUS)
    fibre_type = _category(row, NORM_FIBRE_TYPE, canonical_fibre_type, FIBRE_TYPES, "none")

    # Missing descriptors remain masked.
    max_coarse = row_value_optional(row, NORM_MAX_COARSE)
    max_fine = row_value_optional(row, NORM_MAX_FINE)
    curing_temp = row_value_optional(row, NORM_CURING_TEMP)
    humidity_pct = row_value_optional(row, NORM_CURING_HUMIDITY)
    relative_humidity = humidity_pct / 100.0 if humidity_pct is not None else 0.95
    rca_pct = row_value_optional(row, NORM_RCA_PCT)
    rca_fraction = max(0.0, min(1.0, rca_pct / 100.0)) if rca_pct is not None else 0.0
    coarse_agg_type = _category(
        row, NORM_COARSE_AGG_TYPE, canonical_coarse_aggregate_type,
        COARSE_AGGREGATE_TYPES, "natural",
    )
    coarse_agg_type_id = float(type_id(coarse_agg_type, COARSE_AGGREGATE_TYPES))
    coarse_agg_sg = row_value_optional(row, NORM_COARSE_AGG_SG)
    coarse_agg_wa = row_value_optional(row, NORM_COARSE_AGG_WA)
    coarse_agg_cs = row_value_optional(row, NORM_COARSE_AGG_CS)
    cement_grade = row_value_optional(row, NORM_CEMENT_GRADE)

    return MixDesign(
        aggregate_volume_fraction=aggregate_volume_fraction,
        sand_volume_fraction=sand_vf,
        rca_replacement_fraction=rca_fraction,
        coarse_aggregate_type_id=coarse_agg_type_id,
        coarse_aggregate_specific_gravity=coarse_agg_sg,
        coarse_aggregate_water_absorption_pct=coarse_agg_wa,
        coarse_aggregate_crushing_strength_mpa=coarse_agg_cs,
        water_to_binder_ratio=max(0.10, min(1.20, water_to_binder)),
        cement_content_kg_m3=cement,
        cement_grade_mpa=(cement_grade if cement_grade is not None else float("nan")),
        scm_type_id=scm_type,
        scm_fraction=max(0.0, min(0.70, scm / binder)),
        admixture_dosage=superplasticizer / binder,
        relative_humidity=relative_humidity,
        temperature_C=curing_temp if curing_temp is not None else 20.0,
        temperature_observed=curing_temp is not None,
        curing_age_days=age,
        max_coarse_aggregate_size_mm=max_coarse,
        max_fine_aggregate_size_mm=max_fine,
        fibre_content_kg_m3=fibre_content,
        fibre_length_mm=fibre_length,
        fibre_diameter_mm=fibre_diameter,
        fibre_tensile_strength_MPa=fibre_tensile,
        fibre_modulus_GPa=fibre_modulus,
        fibre_type_id=float(type_id(fibre_type, FIBRE_TYPES)),
    )


def mortar_mix_from_row(row, mix: MixDesign) -> MortarMixDesign:
    fine = row_value(row, NORM_FINE)
    coarse = row_value(row, NORM_COARSE)
    # Convert concrete-basis quantities to the mortar volume.
    mortar_vf = mortar_volume_fraction(coarse / 2650.0)
    # Use typical chemistry where measurements are missing.
    cement_chem = (
        _chem_or_default(row, NORM_CEMENT_CAO, _DEFAULT_CEMENT_CHEM[0]),
        _chem_or_default(row, NORM_CEMENT_SIO2, _DEFAULT_CEMENT_CHEM[1]),
        _chem_or_default(row, NORM_CEMENT_AL2O3, _DEFAULT_CEMENT_CHEM[2]),
        _chem_or_default(row, NORM_CEMENT_FE2O3, _DEFAULT_CEMENT_CHEM[3]),
    )
    cement_chem_ext = (
        _chem_or_default(row, NORM_CEMENT_MGO, _DEFAULT_CEMENT_CHEM_EXT[0]),
        _chem_or_default(row, NORM_CEMENT_SO3, _DEFAULT_CEMENT_CHEM_EXT[1]),
        _chem_or_default(row, NORM_CEMENT_ALKALI, _DEFAULT_CEMENT_CHEM_EXT[2]),
    )
    scm_chem = (
        _chem_or_default(row, NORM_SCM_CAO, _DEFAULT_SCM_CHEM[0]),
        _chem_or_default(row, NORM_SCM_SIO2, _DEFAULT_SCM_CHEM[1]),
    )
    scm_chem_ext = (
        _chem_or_default(row, NORM_SCM_AL2O3, _DEFAULT_SCM_CHEM_EXT[0]),
        _chem_or_default(row, NORM_SCM_FE2O3, _DEFAULT_SCM_CHEM_EXT[1]),
        _chem_or_default(row, NORM_SCM_MGO, _DEFAULT_SCM_CHEM_EXT[2]),
        _chem_or_default(row, NORM_SCM_LOI, _DEFAULT_SCM_CHEM_EXT[3]),
    )
    cement_type = _category(row, NORM_CEMENT_TYPE, canonical_cement_type, CEMENT_TYPES, "other")
    cement_bogue = (
        _chem_or_default(row, NORM_CEMENT_C3S, 55.0),
        _chem_or_default(row, NORM_CEMENT_C2S, 18.0),
        _chem_or_default(row, NORM_CEMENT_C3A, 8.0),
        _chem_or_default(row, NORM_CEMENT_C4AF, 10.0),
    )
    scm_ssa = _chem_or_default(row, NORM_SCM_SSA, 400.0)
    sand_sg = _chem_or_default(row, NORM_FINE_SG, 2.65)
    # Glass powder on a mortar-volume basis.
    glass_powder = row_value(row, NORM_GLASS_POWDER) / mortar_vf
    fly_ash_type = _category(row, NORM_FLY_ASH_TYPE, canonical_fly_ash_type, FLY_ASH_TYPES, "none")
    return MortarMixDesign(
        sand_volume_fraction=max(0.20, min(0.60, fine / 2650.0)),
        water_to_binder_ratio=mix.water_to_binder_ratio,
        cement_content_kg_m3=mix.cement_content_kg_m3 / mortar_vf,
        cement_grade_mpa=mix.cement_grade_mpa,
        scm_type_id=mix.scm_type_id,
        scm_fraction=mix.scm_fraction,
        admixture_dosage=mix.admixture_dosage,
        cement_chem=cement_chem,
        cement_chem_ext=cement_chem_ext,
        scm_chem=scm_chem,
        scm_chem_ext=scm_chem_ext,
        cement_bogue=cement_bogue,
        scm_specific_surface_m2_kg=scm_ssa,
        sand_specific_gravity=sand_sg,
        cement_type_id=float(type_id(cement_type, CEMENT_TYPES)),
        curing_relative_humidity=mix.relative_humidity,
        curing_temperature_C=mix.temperature_C,
        curing_temperature_observed=mix.temperature_observed,
        max_sand_size_mm=mix.max_fine_aggregate_size_mm,
        curing_age_days=mix.curing_age_days,
        fibre_content_kg_m3=mix.fibre_content_kg_m3 / mortar_vf,
        fibre_length_mm=mix.fibre_length_mm,
        fibre_diameter_mm=mix.fibre_diameter_mm,
        fibre_tensile_strength_MPa=mix.fibre_tensile_strength_MPa,
        fibre_modulus_GPa=mix.fibre_modulus_GPa,
        fibre_type_id=mix.fibre_type_id,
        glass_powder_kg_m3=glass_powder,
        fly_ash_type_id=float(type_id(fly_ash_type, FLY_ASH_TYPES)),
    )


def mix_itz_vector(mix: MixDesign, mortar_mix: MortarMixDesign) -> Tensor:
    """Build the per-graph binder, curing and chemistry vector for ITZ edges."""
    return torch.tensor(
        [
            mix.cement_content_kg_m3,
            float(mix.scm_type_id),
            mix.scm_fraction,
            mix.water_to_binder_ratio,
            mix.admixture_dosage,
            mix.relative_humidity,
            mix.temperature_C,
            mix.curing_age_days,
            *binder_chem_features(mortar_mix),
            mortar_mix.glass_powder_kg_m3,
            mortar_mix.fly_ash_type_id,
            float(mortar_mix.cement_type_id),
            float(mortar_mix.cement_grade_mpa),
        ],
        dtype=torch.float32,
    )


RAW_MIX_COLUMNS = (
    NORM_CEMENT,
    NORM_SLAG,
    NORM_FLY_ASH,
    NORM_SILICA_FUME,
    NORM_METAKAOLIN,
    NORM_LIMESTONE,
    NORM_OTHER_SCM,
    NORM_WATER,
    NORM_SUPERPLASTICIZER,
    NORM_COARSE,
    NORM_FINE,
    NORM_FIBRE_CONTENT,
    NORM_FIBRE_LENGTH,
    NORM_FIBRE_DIAMETER,
    NORM_FIBRE_TENSILE,
    NORM_FIBRE_MODULUS,
    NORM_AGE,
)
# Tabular input = raw mix columns, then the rich maskable columns, then one
# observed-mask channel per rich column (1 = measured, 0 = missing), then one-hot
# blocks for the categorical columns. The mask channels let the ANN distinguish
# "not reported" from a genuine zero.
_CATEGORICAL_WIDTH = sum(len(vocab) for _, _, vocab, _ in CATEGORICAL_TABULAR)
TABULAR_DIM = len(RAW_MIX_COLUMNS) + 2 * len(RICH_TABULAR_COLUMNS) + _CATEGORICAL_WIDTH


def raw_mix_vector(row) -> Tensor:
    """Build the numeric, mask, and categorical input for ``TabularANN``."""

    base = [row_value(row, name) for name in RAW_MIX_COLUMNS]

    rich_values: List[float] = []
    rich_mask: List[float] = []
    for name in RICH_TABULAR_COLUMNS:
        v = row_value_optional(row, name)
        rich_values.append(0.0 if v is None else v)
        rich_mask.append(0.0 if v is None else 1.0)

    categorical: List[float] = []
    for name, canon_fn, vocab, default in CATEGORICAL_TABULAR:
        categorical.extend(one_hot(_category(row, name, canon_fn, vocab, default), vocab))

    return torch.tensor(base + rich_values + rich_mask + categorical, dtype=torch.float32)


class ConcreteMixDataset(Dataset):
    def __init__(
        self,
        data: TableInput,
        indices: Optional[Sequence[int]] = None,
        seed: int = 17,
        target_standardizer: Optional[Standardizer] = None,
        schema: SchemaSpec = DEFAULT_SCHEMA,
    ):
        df = read_table(data)
        if indices is not None:
            df = df.iloc[list(indices)].reset_index(drop=True)
        self.df = df.reset_index(drop=True)
        self.schema = schema
        y = torch.tensor(
            self.df[strength_column(self.df)].to_numpy(dtype="float32"),
            dtype=torch.float32,
        )
        self.y_scaled = (
            target_standardizer.encode(y) if target_standardizer is not None else y
        )
        self.samples: List[ConcreteSample] = []
        self._build_samples(seed)

    def _build_samples(self, seed: int) -> None:
        for i, row in self.df.iterrows():
            mix = row_to_mix(row)
            mortar_mix = mortar_mix_from_row(row, mix)
            graph = generate_concrete_graph(
                mix,
                schema=self.schema,
                seed=seed + 1000 + i,
            )
            micro = generate_mortar_graph(
                mortar_mix, schema=self.schema, seed=seed + 2000 + i
            )
            # Targets are fully data-driven: a blank compressive_strength_mpa (e.g.
            # pure-paste rows, whose strength lives in paste_compressive_strength_mpa)
            # yields a NaN concrete target -> masked out of the concrete loss/eval.
            # Likewise a blank mortar_compressive_strength_mpa -> masked mortar target.
            y = torch.zeros(len(self.schema.concrete_targets), dtype=torch.float32)
            y[0] = self.y_scaled[i]
            mortar_t = mortar_target_from_row(row, self.schema)
            paste_t = derive_paste_porosity(
                paste_target_from_row(row, self.schema), mortar_t,
                mix.sand_volume_fraction, self.schema,
            )
            self.samples.append(
                ConcreteSample(
                    graph=graph,
                    micrograph=micro,
                    mix_itz=mix_itz_vector(mix, mortar_mix),
                    concrete_target=y,
                    concrete_aux_target=concrete_aux_target_from_row(row),
                    mortar_target=mortar_t,
                    itz_target_graph=itz_target_from_row(row, self.schema),
                    paste_target=paste_t,
                    sand_volume_fraction=torch.tensor(mix.sand_volume_fraction, dtype=torch.float32),
                    tabular_mix=raw_mix_vector(row),
                )
            )

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> ConcreteSample:
        return self.samples[idx]


class StrengthHead(nn.Module):
    """Expose standardized compression while retaining physical model outputs."""

    def __init__(self, base: nn.Module, standardizer: Standardizer):
        super().__init__()
        self.base = base
        self.standardizer = standardizer
        self.mortar_pred: Optional[Tensor] = None
        self.itz_pred: Optional[Tensor] = None
        self.itz_pred_per_graph: Optional[Tensor] = None
        self.paste_pred: Optional[Tensor] = None
        self.paste_abrams_penalty: Optional[Tensor] = None
        self.itz_wb_penalty: Optional[Tensor] = None
        self.concrete_pred: Optional[Tensor] = None

    def forward(self, batch) -> Tensor:
        out = self.base(batch)
        self.mortar_pred = out.get("mortar_pred") if isinstance(out, dict) else None
        self.itz_pred = out.get("itz_pred") if isinstance(out, dict) else None
        self.itz_pred_per_graph = out.get("itz_pred_per_graph") if isinstance(out, dict) else None
        self.paste_pred = out.get("paste_pred") if isinstance(out, dict) else None
        self.paste_abrams_penalty = out.get("paste_abrams_penalty") if isinstance(out, dict) else None
        self.itz_wb_penalty = out.get("itz_wb_penalty") if isinstance(out, dict) else None
        self.concrete_pred = out.get("concrete_pred") if isinstance(out, dict) else None
        return self.standardizer.encode(out["concrete_pred"][:, 0])


def metrics(pred_mpa: Tensor, true_mpa: Tensor) -> Dict[str, float]:
    # Drop rows without a concrete target (NaN), e.g. pure-paste specimens that
    # only supervise the paste head -- they must not enter concrete R2/RMSE.
    keep = ~torch.isnan(true_mpa)
    pred_mpa, true_mpa = pred_mpa[keep], true_mpa[keep]
    resid = true_mpa - pred_mpa
    ss_res = resid.pow(2).sum()
    ss_tot = (true_mpa - true_mpa.mean()).pow(2).sum().clamp_min(1.0e-12)
    return {
        "rmse": float(resid.pow(2).mean().sqrt()),
        "mae": float(resid.abs().mean()),
        "r2": float(1.0 - ss_res / ss_tot),
    }


@torch.no_grad()
def collect_strength_predictions(
    model: StrengthHead, loader: DataLoader, device: torch.device
) -> tuple:
    model.eval()
    pred_scaled, true_scaled = [], []
    for batch in loader:
        batch = batch.to(device)
        pred_scaled.append(model(batch).cpu())
        true_scaled.append(batch.concrete_target[:, 0].cpu())
    pred = model.standardizer.decode(torch.cat(pred_scaled))
    true = model.standardizer.decode(torch.cat(true_scaled))
    # Drop rows with no concrete target (NaN), e.g. pure-paste specimens, so both
    # metrics and parity figures see only concrete-supervised rows.
    keep = ~torch.isnan(true)
    return pred[keep], true[keep]


def evaluate_strength(
    model: StrengthHead, loader: DataLoader, device: torch.device
) -> Dict[str, float]:
    pred, true = collect_strength_predictions(model, loader, device)
    return metrics(pred, true)


def fit_encoder_standardizers(model: nn.Module, samples: Sequence) -> None:
    """Fit masked feature statistics on the training samples."""

    from .missing_features import masked_feature_stats, masked_feature_stats_by_type

    if not samples or not hasattr(model, "mortar"):
        return

    schema = model.schema

    # Mesoscale globals.
    gf = torch.stack([s.graph.global_features for s in samples])
    gm = torch.stack([s.graph.global_mask for s in samples])
    model.mesoscale.global_encoder.set_feature_stats(*masked_feature_stats(gf, gm))

    # Mesoscale nodes by type.
    nx = torch.cat([s.graph.x for s in samples], dim=0)
    nxm = torch.cat([s.graph.x_mask for s in samples], dim=0)
    nnt = torch.cat([s.graph.node_type for s in samples], dim=0)
    model.mesoscale.node_encoder.set_feature_stats(
        *masked_feature_stats_by_type(nx, nxm, nnt, schema.num_node_types)
    )

    # Mesoscale edges by type.
    ex = torch.cat([s.graph.edge_attr for s in samples], dim=0)
    exm = torch.cat([s.graph.edge_attr_mask for s in samples], dim=0)
    eet = torch.cat([s.graph.edge_type for s in samples], dim=0)
    model.mesoscale.edge_encoder.set_feature_stats(
        *masked_feature_stats_by_type(ex, exm, eet, schema.num_edge_types)
    )

    # Mortar nodes by type.
    mx = torch.cat([s.micrograph.x for s in samples], dim=0)
    mxm = torch.cat([s.micrograph.x_mask for s in samples], dim=0)
    mnt = torch.cat([s.micrograph.node_type for s in samples], dim=0)
    mortar_node_mean, mortar_node_std = masked_feature_stats_by_type(
        mx, mxm, mnt, schema.num_mortar_node_types
    )
    model.mortar.node_encoder.set_feature_stats(
        mortar_node_mean, mortar_node_std
    )

    # Paste composition.
    if hasattr(model, "paste_sub"):
        d = model.paste_sub.input_dim
        model.paste_sub.encoder.set_feature_stats(
            mortar_node_mean[MORTAR_NODE_TYPE_PASTE][:d],
            mortar_node_std[MORTAR_NODE_TYPE_PASTE][:d],
        )

    # Mortar globals.
    mg = torch.stack([s.micrograph.global_features for s in samples])
    mgm = torch.stack([s.micrograph.global_mask for s in samples])
    model.mortar.global_encoder.set_feature_stats(*masked_feature_stats(mg, mgm))

    # ITZ mix features. Geometry columns retain identity statistics.
    mix_itz = torch.stack([s.mix_itz for s in samples])
    mean, std = masked_feature_stats(mix_itz, (~torch.isnan(mix_itz)).float())
    enc = model.itz.encoder
    full_mean = enc.feat_mean[0].clone()
    full_std = enc.feat_std[0].clone()
    k = mix_itz.size(1)
    full_mean[:k] = mean
    full_std[:k] = std
    enc.set_feature_stats(full_mean, full_std)

    # Abrams paste features.
    abrams = model.paste_abrams
    if abrams.paste_encoder is not None:
        cols = abrams.ox_cols
        abrams.paste_encoder.set_feature_stats(
            mortar_node_mean[MORTAR_NODE_TYPE_PASTE][cols],
            mortar_node_std[MORTAR_NODE_TYPE_PASTE][cols],
        )
