"""Generate concrete and mortar graphs from a mix design."""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np
import torch
from torch import Tensor

from .schema import (
    DEFAULT_SCHEMA,
    EDGE_TYPE_AGGREGATE_AGGREGATE,
    EDGE_TYPE_ITZ,
    EDGE_TYPE_MORTAR_MORTAR,
    MORTAR_NODE_TYPE_PASTE,
    MORTAR_NODE_TYPE_SAND,
    NODE_TYPE_AGGREGATE,
    NODE_TYPE_MORTAR,
    SchemaSpec,
)


# Mix design


def mortar_volume_fraction(coarse_aggregate_volume_fraction: float) -> float:
    """Return the mortar fraction after bounding coarse aggregate at 0.70."""

    phi_coarse = min(0.70, max(0.0, coarse_aggregate_volume_fraction))
    return max(0.30, 1.0 - phi_coarse)



@dataclass
class MixDesign:
    """Concrete-scale mix design and curing metadata."""

    rve_size_mm: float = 150.0
    aggregate_size_bins_mm: Sequence[float] = (8.0, 12.0, 16.0, 20.0, 24.0, 32.0)
    aggregate_size_fractions: Sequence[float] = (0.10, 0.20, 0.25, 0.20, 0.15, 0.10)
    aggregate_volume_fraction: float = 0.40
    # Sand (fine-aggregate) volume fraction of the MORTAR phase, assuming sand is
    # non-porous. Used to derive paste porosity from mortar porosity and for the
    # volumetric consistency prior: phi_mortar = phi_paste * (1 - sand_volume_fraction).
    sand_volume_fraction: float = 0.0
    # Recycled aggregate fraction (0 = natural, 1 = all recycled).
    rca_replacement_fraction: float = 0.0
    # Aggregate type and optional measured properties.
    coarse_aggregate_type_id: float = 0.0
    coarse_aggregate_specific_gravity: Optional[float] = None
    coarse_aggregate_water_absorption_pct: Optional[float] = None
    coarse_aggregate_crushing_strength_mpa: Optional[float] = None
    water_to_binder_ratio: float = 0.42
    cement_content_kg_m3: float = 360.0
    cement_grade_mpa: float = 42.5  # cement strength class (MPa); default typical OPC
    scm_type_id: int = 1
    scm_fraction: float = 0.12
    admixture_dosage: float = 0.5
    relative_humidity: float = 0.95
    temperature_C: float = 20.0
    curing_age_days: float = 28.0
    fibre_content_kg_m3: float = 0.0
    fibre_length_mm: float = 0.0
    fibre_diameter_mm: float = 0.0
    fibre_tensile_strength_MPa: float = 0.0
    fibre_modulus_GPa: float = 0.0
    # Fibre material-class id (index into categoricals.FIBRE_TYPES; 0 = none).
    fibre_type_id: float = 0.0
    # Measured nominal max grain sizes (mm); ``None`` means not reported, in
    # which case the default bin distribution is used and the corresponding
    # global feature is masked out. ``temperature_observed`` drives the same
    # masking for curing temperature (which always carries a numeric fallback).
    max_coarse_aggregate_size_mm: Optional[float] = None
    max_fine_aggregate_size_mm: Optional[float] = None
    temperature_observed: bool = True

    def effective_aggregate_bins_mm(self) -> np.ndarray:
        """Default bins rescaled so their top equals the measured max size.

        Rescaling (rather than clipping) preserves the relative gradation shape
        while respecting the reported nominal maximum coarse-aggregate size.
        """

        bins = np.asarray(self.aggregate_size_bins_mm, dtype=float)
        if self.max_coarse_aggregate_size_mm and self.max_coarse_aggregate_size_mm > 0:
            top = bins.max()
            if top > 0:
                bins = bins * (float(self.max_coarse_aggregate_size_mm) / top)
        return bins

    def sample_radius(self, rng: np.random.Generator) -> float:
        bins = self.effective_aggregate_bins_mm()
        weights = np.asarray(self.aggregate_size_fractions, dtype=float)
        weights = weights / weights.sum()
        idx = int(rng.choice(len(bins), p=weights))
        return float(bins[idx] * 0.5 * rng.uniform(0.85, 1.0))


@dataclass
class MortarMixDesign:
    """Mortar-scale mix description, distinct from :class:`MixDesign`.

    The mortar scale operates on a small RVE (default 6 mm) so each micrograph
    holds on the order of 10-30 nodes - enough for relational message passing
    but cheap enough to train end-to-end alongside the mesoscale GNN.
    """

    rve_size_mm: float = 6.0
    sand_size_bins_mm: Sequence[float] = (0.6, 1.0, 1.5, 2.0)
    sand_size_fractions: Sequence[float] = (0.2, 0.3, 0.3, 0.2)
    sand_volume_fraction: float = 0.35
    water_to_binder_ratio: float = 0.45
    cement_content_kg_m3: float = 360.0
    cement_grade_mpa: float = 42.5  # cement strength class (MPa); default typical OPC
    scm_type_id: int = 1
    scm_fraction: float = 0.20
    admixture_dosage: float = 0.5
    cement_chem: Tuple[float, float, float, float] = (64.0, 21.0, 5.0, 3.0)
    # Extra cement oxides (MgO, SO3, Na2O-equivalent alkali).
    cement_chem_ext: Tuple[float, float, float] = (2.0, 2.5, 0.6)
    scm_chem: Tuple[float, float] = (4.0, 55.0)
    # Extended SCM oxides (Al2O3, Fe2O3, MgO, LOI); generic-pozzolan defaults.
    scm_chem_ext: Tuple[float, float, float, float] = (15.0, 5.0, 2.0, 3.0)
    # Bogue clinker phases (C3S, C2S, C3A, C4AF); OPC-typical defaults when not measured.
    cement_bogue: Tuple[float, float, float, float] = (55.0, 18.0, 8.0, 10.0)
    # SCM Blaine-equivalent specific surface (m2/kg); generic default when not measured.
    scm_specific_surface_m2_kg: float = 400.0
    # Fine-aggregate (sand) specific gravity; typical siliceous default.
    sand_specific_gravity: float = 2.65
    # Cement-type categorical id (index into categoricals.CEMENT_TYPES);
    # default 1.0 = "type_i_ii" (generic Portland cement).
    cement_type_id: float = 1.0
    curing_relative_humidity: float = 0.95
    curing_temperature_C: float = 23.0
    curing_age_days: float = 28.0
    fibre_content_kg_m3: float = 0.0
    fibre_length_mm: float = 0.0
    fibre_diameter_mm: float = 0.0
    fibre_tensile_strength_MPa: float = 0.0
    fibre_modulus_GPa: float = 0.0
    # Fibre material-class id (index into categoricals.FIBRE_TYPES; 0 = none),
    # carried so the mortar sub-model conditions on fibre type, not just dosage.
    fibre_type_id: float = 0.0
    # Measured nominal max fine/sand grain size (mm); ``None`` means not
    # reported. ``curing_temperature_observed`` masks the mortar curing-temp
    # global when the source did not report it.
    max_sand_size_mm: Optional[float] = None
    curing_temperature_observed: bool = True
    # Glass powder (pozzolanic SCM, kg/m^3 mortar) and fly-ash class id (index
    # into categoricals.FLY_ASH_TYPES; 0 = none). Default 0/absent. (Quartz powder
    # is treated as fine aggregate, not a binder channel.)
    glass_powder_kg_m3: float = 0.0
    fly_ash_type_id: float = 0.0

    def effective_sand_bins_mm(self) -> np.ndarray:
        """Default sand bins rescaled to the measured max fine grain size."""

        bins = np.asarray(self.sand_size_bins_mm, dtype=float)
        if self.max_sand_size_mm and self.max_sand_size_mm > 0:
            top = bins.max()
            if top > 0:
                bins = bins * (float(self.max_sand_size_mm) / top)
        return bins


# Graph containers


@dataclass
class ConcreteGraph:
    x: Tensor
    x_mask: Tensor
    edge_index: Tensor
    edge_attr: Tensor
    edge_attr_mask: Tensor
    global_features: Tensor
    global_mask: Tensor
    pos: Tensor
    node_type: Tensor
    edge_type: Tensor


@dataclass
class MortarMicrograph:
    x: Tensor
    x_mask: Tensor
    edge_index: Tensor
    edge_attr: Tensor
    node_type: Tensor
    global_features: Tensor
    global_mask: Tensor


# Geometry


# Ignore trace coarse aggregate that does not form a continuous phase.
MIN_AGG_VOLUME_FRACTION = 0.01


def _sample_aggregates(
    mix: MixDesign,
    rng: np.random.Generator,
    max_attempts: int = 6000,
) -> Tuple[np.ndarray, np.ndarray]:
    rve_area = mix.rve_size_mm ** 2

    if mix.aggregate_volume_fraction < MIN_AGG_VOLUME_FRACTION:
        return np.zeros((0, 2), dtype=float), np.zeros((0,), dtype=float)

    target_area = mix.aggregate_volume_fraction * rve_area

    centroids: List[np.ndarray] = []
    radii: List[float] = []
    placed_area = 0.0
    attempts = 0
    while placed_area < target_area and attempts < max_attempts:
        attempts += 1
        r = mix.sample_radius(rng)
        c = rng.uniform(r, mix.rve_size_mm - r, size=2)
        ok = True
        for cj, rj in zip(centroids, radii):
            if np.linalg.norm(c - cj) < (r + rj) * 1.05:
                ok = False
                break
        if not ok:
            continue
        centroids.append(c)
        radii.append(r)
        placed_area += math.pi * r * r

    if not centroids:
        # Keep the graph nonempty if placement fails.
        centroids = [np.array([mix.rve_size_mm / 2.0, mix.rve_size_mm / 2.0])]
        radii = [mix.rve_size_mm * 0.05]
    return np.asarray(centroids), np.asarray(radii)


def _mortar_seed_grid(
    mix: MixDesign,
    agg_centroids: np.ndarray,
    agg_radii: np.ndarray,
    n_side: int = 8,
) -> np.ndarray:
    xs = np.linspace(0, mix.rve_size_mm, n_side + 2)[1:-1]
    grid = np.array([(x, y) for x in xs for y in xs])
    if len(agg_centroids) == 0:
        return grid
    diff = grid[:, None, :] - agg_centroids[None, :, :]
    d2 = np.einsum("nki,nki->nk", diff, diff)
    inside = np.any(d2 < (agg_radii * 1.05) ** 2, axis=1)
    kept = grid[~inside]
    if len(kept) == 0:
        kept = np.array([[mix.rve_size_mm / 2.0, mix.rve_size_mm / 2.0]])
    return kept


def _voronoi_adjacency(
    points: np.ndarray,
    rve_size: float,
    rng: np.random.Generator,
    n_samples: int = 4000,
) -> Dict[Tuple[int, int], int]:
    """Approximate Voronoi adjacency by sampling.

    For each random sample we look up its two nearest seeds; that pair shares
    a Voronoi boundary.  The number of samples falling on each pair is a
    monotone proxy for the shared boundary length.
    """

    if points.shape[0] < 2:
        return {}
    samples = rng.uniform(0.0, rve_size, size=(n_samples, 2))
    diff = samples[:, None, :] - points[None, :, :]
    d2 = np.einsum("nki,nki->nk", diff, diff)
    nearest_two = np.argpartition(d2, kth=1, axis=1)[:, :2]
    adjacency: Dict[Tuple[int, int], int] = {}
    a = nearest_two[:, 0]
    b = nearest_two[:, 1]
    lo = np.minimum(a, b)
    hi = np.maximum(a, b)
    for i, j in zip(lo, hi):
        if i == j:
            continue
        key = (int(i), int(j))
        adjacency[key] = adjacency.get(key, 0) + 1
    return adjacency


def _voronoi_cell_fraction(
    seeds: np.ndarray,
    rve_size: float,
    rng: np.random.Generator,
    n_samples: int = 4000,
) -> np.ndarray:
    """Fraction of the RVE area nearest to each seed (Monte-Carlo Voronoi cells).

    Returns an array summing to 1.0; used to apportion the non-particle area
    across mortar / paste seed nodes so each carries a heterogeneous area share.
    """

    n = len(seeds)
    if n == 0:
        return np.zeros(0, dtype=np.float64)
    if n == 1:
        return np.ones(1, dtype=np.float64)
    samples = rng.uniform(0.0, rve_size, size=(n_samples, 2))
    diff = samples[:, None, :] - seeds[None, :, :]
    d2 = np.einsum("nki,nki->nk", diff, diff)
    nearest = np.argmin(d2, axis=1)
    counts = np.bincount(nearest, minlength=n).astype(np.float64)
    total = counts.sum()
    if total <= 0:
        return np.full(n, 1.0 / n, dtype=np.float64)
    return counts / total


# Concrete features


def _aggregate_feature_vector(
    radius: float,
    rng: np.random.Generator,
    schema: SchemaSpec,
    area_fraction: float = 0.0,
    rca_replacement_fraction: float = 0.0,
    coarse_aggregate_type_id: float = 0.0,
    specific_gravity: float = None,
    water_absorption_pct: float = None,
    crushing_strength_mpa: float = None,
) -> np.ndarray:
    feats = np.zeros(len(schema.aggregate), dtype=np.float32)
    feats[0] = 2.0 * radius
    feats[1] = float(rng.uniform(0.6, 1.0))
    feats[2] = float(rng.normal(120.0, 15.0))
    feats[3] = float(rng.normal(20.0, 4.0))
    feats[4] = float(rng.uniform(0.7, 1.0))
    feats[5] = float(rng.uniform(0.0, 0.05))
    feats[6] = float(rng.uniform(0.2, 2.0))
    feats[7] = float(rng.uniform(0.4, 1.0))
    feats[8] = float(rng.normal(2.65, 0.05))
    feats[9] = float(rng.normal(1600.0, 50.0))
    # Real reported aggregate quality overrides the synthetic prior in its slot.
    if crushing_strength_mpa is not None:
        feats[2] = float(crushing_strength_mpa)
    if water_absorption_pct is not None:
        feats[6] = float(water_absorption_pct)
    if specific_gravity is not None:
        feats[8] = float(specific_gravity)
    feats[schema.aggregate.index("area_fraction")] = float(area_fraction)
    feats[schema.aggregate.index("rca_replacement_fraction")] = float(rca_replacement_fraction)
    feats[schema.aggregate.index("coarse_aggregate_type_id")] = float(coarse_aggregate_type_id)
    return feats


def _mortar_placeholder_vector(
    mix: MixDesign,
    rng: np.random.Generator,
    schema: SchemaSpec,
    area_fraction: float = 0.0,
) -> np.ndarray:
    """A coarse prior for mortar features; overwritten by the sub-model."""

    feats = np.zeros(len(schema.mortar), dtype=np.float32)
    w = mix.water_to_binder_ratio
    base = 85.0 - 65.0 * w
    feats[0] = float(rng.normal(base, 5.0))
    feats[1] = float(rng.normal(0.1 * base, 0.5))
    feats[2] = float(rng.normal(0.16 * base, 0.5))
    feats[3] = float(rng.normal(15.0 + 0.35 * base, 2.0))
    feats[4] = float(rng.uniform(1e-12, 1e-10))
    feats[5] = float(rng.uniform(1e-12, 1e-10))
    feats[6] = float(rng.uniform(0.08, 0.20))
    feats[7] = float(rng.uniform(1.0, 3.0))
    # Fibre conditioning (slots 8-11): broadcast mix-level fibre attributes
    # to every mortar node. Aspect ratio = length / diameter when both > 0.
    feats[8] = float(mix.fibre_content_kg_m3)
    if mix.fibre_diameter_mm > 0.0:
        feats[9] = float(mix.fibre_length_mm / mix.fibre_diameter_mm)
    feats[10] = float(mix.fibre_tensile_strength_MPa)
    feats[11] = float(mix.fibre_modulus_GPa)
    feats[schema.mortar.index("mortar_area_fraction")] = float(area_fraction)
    feats[schema.mortar.index("fibre_type_id")] = float(mix.fibre_type_id)
    return feats


def _itz_placeholder_vector(
    mix: MixDesign, rng: np.random.Generator, schema: SchemaSpec
) -> np.ndarray:
    feats = np.zeros(len(schema.itz), dtype=np.float32)
    w = mix.water_to_binder_ratio
    feats[0] = float(rng.uniform(5.0, 70.0))  # thickness prior matches [5,70] transform
    # Porosity rises with w/b (kept physical) but spans the wider [0.05, 0.70].
    feats[1] = float(np.clip(0.05 + 0.45 * w + rng.uniform(-0.03, 0.03), 0.05, 0.70))
    # Strength channel = ITZ tensile/bond strength -> tensile-scale prior (~half
    # the mortar tensile); modulus prior sits in the 0.3-0.7 x paste band. These
    # only calibrate encoder standardization; the sub-model overwrites them.
    feats[2] = float(rng.normal(4.0, 1.5))
    feats[3] = float(rng.normal(7.0, 2.0))
    feats[4] = float(rng.uniform(1e-11, 1e-9))
    feats[5] = float(rng.uniform(1e-11, 1e-9))
    return feats


def _global_feature_vector(
    mix: MixDesign, schema: SchemaSpec
) -> Tuple[np.ndarray, np.ndarray]:
    """Return ``(values, observed_mask)`` for the concrete global vector.

    Maskable descriptors (curing temperature, max coarse / fine grain size)
    are flagged 0 in the mask when the mix did not report them, so the
    MissingFeatureEncoder treats them as missing rather than as a real value.
    """

    values = {
        "relative_humidity": mix.relative_humidity,
        "temperature_C": mix.temperature_C,
        "curing_age_days": mix.curing_age_days,
        "water_to_binder_ratio": mix.water_to_binder_ratio,
        "cement_content_kg_m3": mix.cement_content_kg_m3,
        "scm_fraction": mix.scm_fraction,
        "aggregate_volume_fraction": mix.aggregate_volume_fraction,
        "fibre_content_kg_m3": mix.fibre_content_kg_m3,
        "max_coarse_aggregate_size_mm": mix.max_coarse_aggregate_size_mm or 0.0,
        "max_fine_aggregate_size_mm": mix.max_fine_aggregate_size_mm or 0.0,
    }
    observed = {
        "temperature_C": mix.temperature_observed,
        "max_coarse_aggregate_size_mm": mix.max_coarse_aggregate_size_mm is not None,
        "max_fine_aggregate_size_mm": mix.max_fine_aggregate_size_mm is not None,
    }
    g = np.zeros(len(schema.glob), dtype=np.float32)
    mask = np.zeros(len(schema.glob), dtype=np.float32)
    for i, name in enumerate(schema.glob):
        if name in values:
            g[i] = values[name]
            mask[i] = 1.0 if observed.get(name, True) else 0.0
    return g, mask


# Concrete graph


def generate_concrete_graph(
    mix: MixDesign,
    schema: SchemaSpec = DEFAULT_SCHEMA,
    seed: Optional[int] = None,
) -> ConcreteGraph:
    rng = np.random.default_rng(seed)
    agg_xy, agg_r = _sample_aggregates(mix, rng)
    mortar_xy = _mortar_seed_grid(mix, agg_xy, agg_r)

    n_agg = len(agg_xy)
    n_mortar = len(mortar_xy)
    all_points = np.concatenate([agg_xy, mortar_xy], axis=0)
    adjacency = _voronoi_adjacency(all_points, rve_size=mix.rve_size_mm, rng=rng)

    # Augment mortar-mortar edges via k-nearest neighbours (Voronoi cells of
    # mortar seeds are frequently interrupted by aggregate cells).
    if n_mortar > 1:
        k = min(4, n_mortar - 1)
        d2 = np.sum((mortar_xy[:, None, :] - mortar_xy[None, :, :]) ** 2, axis=-1)
        np.fill_diagonal(d2, np.inf)
        nbr = np.argpartition(d2, kth=k - 1, axis=1)[:, :k]
        for i in range(n_mortar):
            for j in nbr[i]:
                a = n_agg + int(i)
                b = n_agg + int(j)
                if a == b:
                    continue
                key = (min(a, b), max(a, b))
                adjacency.setdefault(key, 1)

    pad_node = schema.node_pad_dim
    pad_edge = schema.edge_pad_dim
    n_total = n_agg + n_mortar

    x = np.zeros((n_total, pad_node), dtype=np.float32)
    x_mask = np.zeros((n_total, pad_node), dtype=np.float32)

    # Per-node area fractions. Aggregates use exact disk area; mortar seeds split
    # the remaining (non-aggregate) area by their Monte-Carlo Voronoi cell share.
    rve_area = float(mix.rve_size_mm ** 2)
    agg_area_frac = (np.pi * np.asarray(agg_r, dtype=np.float64) ** 2) / rve_area
    mortar_total_frac = max(0.0, 1.0 - float(agg_area_frac.sum()))
    mortar_cell = _voronoi_cell_fraction(mortar_xy, mix.rve_size_mm, rng)
    mortar_area_frac = mortar_cell * mortar_total_frac

    for i in range(n_agg):
        feats = _aggregate_feature_vector(
            agg_r[i], rng, schema, area_fraction=float(agg_area_frac[i]),
            rca_replacement_fraction=mix.rca_replacement_fraction,
            coarse_aggregate_type_id=mix.coarse_aggregate_type_id,
            specific_gravity=mix.coarse_aggregate_specific_gravity,
            water_absorption_pct=mix.coarse_aggregate_water_absorption_pct,
            crushing_strength_mpa=mix.coarse_aggregate_crushing_strength_mpa,
        )
        agg_dim = len(schema.aggregate)
        x[i, :agg_dim] = feats
        x_mask[i, :agg_dim] = 1.0

    for j in range(n_mortar):
        feats = _mortar_placeholder_vector(
            mix, rng, schema, area_fraction=float(mortar_area_frac[j])
        )
        mortar_dim = len(schema.mortar)
        x[n_agg + j, :mortar_dim] = feats
        x_mask[n_agg + j, :mortar_dim] = 1.0

    node_type = np.concatenate(
        [
            np.full(n_agg, NODE_TYPE_AGGREGATE, dtype=np.int64),
            np.full(n_mortar, NODE_TYPE_MORTAR, dtype=np.int64),
        ]
    )

    # Edges
    edge_index: List[List[int]] = []
    edge_attr: List[np.ndarray] = []
    edge_mask: List[np.ndarray] = []
    edge_type: List[int] = []

    for (a, b), shared in adjacency.items():
        type_a = NODE_TYPE_AGGREGATE if a < n_agg else NODE_TYPE_MORTAR
        type_b = NODE_TYPE_AGGREGATE if b < n_agg else NODE_TYPE_MORTAR
        pa = all_points[a]
        pb = all_points[b]
        dist = float(np.linalg.norm(pa - pb))
        boundary = float(shared) * (mix.rve_size_mm / 60.0)

        feats = np.zeros(pad_edge, dtype=np.float32)
        mask = np.zeros(pad_edge, dtype=np.float32)

        if type_a != type_b:
            etype = EDGE_TYPE_ITZ
            itz_vals = _itz_placeholder_vector(mix, rng, schema)
            feats[: len(itz_vals)] = itz_vals
            mask[: len(itz_vals)] = 1.0
        elif type_a == NODE_TYPE_MORTAR:
            etype = EDGE_TYPE_MORTAR_MORTAR
            feats[0] = dist
            feats[1] = boundary
            feats[2] = float(pb[0] - pa[0])
            feats[3] = float(pb[1] - pa[1])
            mask[:4] = 1.0
        else:
            ai = a
            bi = b
            surface_gap = max(0.0, dist - float(agg_r[ai]) - float(agg_r[bi]))
            if surface_gap > 1.5 * (mix.rve_size_mm / 60.0):
                continue
            etype = EDGE_TYPE_AGGREGATE_AGGREGATE
            feats[0] = surface_gap
            feats[1] = dist
            mask[:2] = 1.0

        for src, dst in ((a, b), (b, a)):
            edge_index.append([src, dst])
            edge_attr.append(feats.copy())
            edge_mask.append(mask.copy())
            edge_type.append(etype)

    if not edge_index:
        edge_index = [[0, 0]]
        edge_attr = [np.zeros(pad_edge, dtype=np.float32)]
        edge_mask = [np.zeros(pad_edge, dtype=np.float32)]
        edge_type = [EDGE_TYPE_ITZ]

    edge_index_t = torch.tensor(np.asarray(edge_index).T, dtype=torch.long)
    edge_attr_t = torch.tensor(np.stack(edge_attr), dtype=torch.float32)
    edge_mask_t = torch.tensor(np.stack(edge_mask), dtype=torch.float32)
    edge_type_t = torch.tensor(edge_type, dtype=torch.long)
    pos_t = torch.tensor(all_points, dtype=torch.float32)
    x_t = torch.tensor(x, dtype=torch.float32)
    x_mask_t = torch.tensor(x_mask, dtype=torch.float32)
    node_type_t = torch.tensor(node_type, dtype=torch.long)

    g, g_observed = _global_feature_vector(mix, schema)
    g_t = torch.tensor(g, dtype=torch.float32)
    g_mask = torch.tensor(g_observed, dtype=torch.float32)
    g_t = g_t * g_mask  # zero out unreported descriptors

    return ConcreteGraph(
        x=x_t,
        x_mask=x_mask_t,
        edge_index=edge_index_t,
        edge_attr=edge_attr_t,
        edge_attr_mask=edge_mask_t,
        global_features=g_t,
        global_mask=g_mask,
        pos=pos_t,
        node_type=node_type_t,
        edge_type=edge_type_t,
    )
# Mortar graph


def _sample_sand(
    mix: MortarMixDesign,
    rng: np.random.Generator,
    max_attempts: int = 800,
) -> Tuple[np.ndarray, np.ndarray]:
    bins = mix.effective_sand_bins_mm()
    weights = np.asarray(mix.sand_size_fractions, dtype=float)
    weights = weights / weights.sum()

    rve_area = mix.rve_size_mm ** 2
    target_area = mix.sand_volume_fraction * rve_area

    centroids: List[np.ndarray] = []
    radii: List[float] = []
    placed = 0.0
    attempts = 0
    while placed < target_area and attempts < max_attempts:
        attempts += 1
        idx = int(rng.choice(len(bins), p=weights))
        r = float(bins[idx] * 0.5 * rng.uniform(0.85, 1.0))
        # Keep grains inside the mortar RVE.
        r = min(r, 0.45 * mix.rve_size_mm)
        c = rng.uniform(r, mix.rve_size_mm - r, size=2)
        if any(np.linalg.norm(c - cj) < (r + rj) * 1.05 for cj, rj in zip(centroids, radii)):
            continue
        centroids.append(c)
        radii.append(r)
        placed += math.pi * r * r

    if not centroids:
        centroids = [np.array([mix.rve_size_mm / 2.0, mix.rve_size_mm / 2.0])]
        radii = [0.5]
    return np.asarray(centroids), np.asarray(radii)


def _paste_grid(mix: MortarMixDesign, sand_xy: np.ndarray, sand_r: np.ndarray) -> np.ndarray:
    n_side = 5
    xs = np.linspace(0, mix.rve_size_mm, n_side + 2)[1:-1]
    grid = np.array([(x, y) for x in xs for y in xs])
    if len(sand_xy) == 0:
        return grid
    diff = grid[:, None, :] - sand_xy[None, :, :]
    d2 = np.einsum("nki,nki->nk", diff, diff)
    inside = np.any(d2 < (sand_r * 1.05) ** 2, axis=1)
    kept = grid[~inside]
    if len(kept) == 0:
        kept = np.array([[mix.rve_size_mm / 2.0, mix.rve_size_mm / 2.0]])
    return kept


def _sand_feature_vector(
    radius: float,
    rng: np.random.Generator,
    schema: SchemaSpec,
    area_fraction: float = 0.0,
    sand_specific_gravity: float = 2.65,
) -> np.ndarray:
    v = np.zeros(len(schema.sand), dtype=np.float32)
    v[0] = 2.0 * radius
    v[1] = float(rng.uniform(2.0, 3.5))
    v[2] = float(rng.uniform(2.4, 3.2))
    v[3] = float(rng.uniform(0.5, 2.0))
    v[4] = float(rng.normal(80.0, 10.0))
    v[5] = float(rng.normal(70.0, 8.0))
    v[schema.sand.index("sand_area_fraction")] = float(area_fraction)
    v[schema.sand.index("sand_specific_gravity")] = float(sand_specific_gravity)
    return v


def _paste_feature_vector(
    mix: MortarMixDesign, schema: SchemaSpec, area_fraction: float = 0.0
) -> np.ndarray:
    v = np.zeros(len(schema.paste), dtype=np.float32)
    v[0] = float(mix.cement_type_id)
    v[1:5] = np.asarray(mix.cement_chem, dtype=np.float32)
    v[5] = mix.cement_content_kg_m3
    v[6] = float(mix.scm_type_id)
    v[7:9] = np.asarray(mix.scm_chem, dtype=np.float32)
    v[9] = mix.scm_fraction * mix.cement_content_kg_m3
    v[10] = mix.water_to_binder_ratio
    v[11] = mix.admixture_dosage
    v[schema.paste.index("paste_area_fraction")] = float(area_fraction)
    ext = np.asarray(mix.cement_chem_ext, dtype=np.float32)
    v[schema.paste.index("cement_MgO_pct")] = ext[0]
    v[schema.paste.index("cement_SO3_pct")] = ext[1]
    v[schema.paste.index("cement_alkali_pct")] = ext[2]
    sext = np.asarray(mix.scm_chem_ext, dtype=np.float32)
    v[schema.paste.index("scm_Al2O3_pct")] = sext[0]
    v[schema.paste.index("scm_Fe2O3_pct")] = sext[1]
    v[schema.paste.index("scm_MgO_pct")] = sext[2]
    v[schema.paste.index("scm_LOI_pct")] = sext[3]
    bg = np.asarray(mix.cement_bogue, dtype=np.float32)
    v[schema.paste.index("cement_C3S_pct")] = bg[0]
    v[schema.paste.index("cement_C2S_pct")] = bg[1]
    v[schema.paste.index("cement_C3A_pct")] = bg[2]
    v[schema.paste.index("cement_C4AF_pct")] = bg[3]
    v[schema.paste.index("scm_specific_surface_m2_kg")] = float(mix.scm_specific_surface_m2_kg)
    v[schema.paste.index("glass_powder_kg_m3")] = float(mix.glass_powder_kg_m3)
    v[schema.paste.index("fly_ash_type_id")] = float(mix.fly_ash_type_id)
    if "cement_grade_mpa" in schema.paste:
        v[schema.paste.index("cement_grade_mpa")] = float(mix.cement_grade_mpa)
    # Initialize paste-property slots before the sub-model replaces them.
    w = mix.water_to_binder_ratio
    paste_property_priors = {
        "paste_compressive_strength_MPa": max(5.0, 95.0 - 70.0 * w),
        "paste_tensile_strength_MPa": max(0.5, 0.09 * (95.0 - 70.0 * w)),
        "paste_elastic_modulus_GPa": float(np.clip(10.0 + 14.0 * (1.0 - w), 5.0, 60.0)),
        "paste_fracture_energy_J_m2": 20.0,
        "paste_porosity": float(np.clip(0.05 + 0.45 * w, 0.02, 0.6)),
    }
    for name, prior in paste_property_priors.items():
        if name in schema.paste:
            v[schema.paste.index(name)] = float(prior)
    return v


def binder_chem_features(mix: MortarMixDesign) -> list:
    """Flat cement + SCM chemistry/phase vector (18 values) shared by the paste
    nodes and the ITZ edge inputs. The ITZ is a modified cement paste, so it sees
    the same binder chemistry the paste does:

    cement oxides (CaO, SiO2, Al2O3, Fe2O3), cement extended (MgO, SO3, alkali),
    Bogue phases (C3S, C2S, C3A, C4AF), SCM oxides (CaO, SiO2), SCM extended
    (Al2O3, Fe2O3, MgO, LOI), SCM fineness (specific surface).
    """

    return [
        *mix.cement_chem,        # 4: CaO, SiO2, Al2O3, Fe2O3
        *mix.cement_chem_ext,    # 3: MgO, SO3, alkali
        *mix.cement_bogue,       # 4: C3S, C2S, C3A, C4AF
        *mix.scm_chem,           # 2: CaO, SiO2
        *mix.scm_chem_ext,       # 4: Al2O3, Fe2O3, MgO, LOI
        float(mix.scm_specific_surface_m2_kg),  # 1
    ]


def _mortar_global_vector(
    mix: MortarMixDesign, schema: SchemaSpec
) -> Tuple[np.ndarray, np.ndarray]:
    """Return ``(values, observed_mask)`` for the mortar global vector."""

    values = {
        "mortar_water_to_binder_ratio": mix.water_to_binder_ratio,
        "mortar_cement_content_kg_m3": mix.cement_content_kg_m3,
        "mortar_scm_fraction": mix.scm_fraction,
        "mortar_admixture_dosage": mix.admixture_dosage,
        "mortar_curing_relative_humidity": mix.curing_relative_humidity,
        "mortar_curing_temperature_C": mix.curing_temperature_C,
        "mortar_curing_age_days": mix.curing_age_days,
        # Full fibre descriptor (dosage is mortar-basis; geometry/strength are
        # intensive so passed as-is). Zero for fibre-free mixes.
        "mortar_fibre_content_kg_m3": mix.fibre_content_kg_m3,
        "mortar_fibre_length_mm": mix.fibre_length_mm,
        "mortar_fibre_diameter_mm": mix.fibre_diameter_mm,
        "mortar_fibre_tensile_strength_MPa": mix.fibre_tensile_strength_MPa,
        "mortar_fibre_modulus_GPa": mix.fibre_modulus_GPa,
        "mortar_fibre_type_id": mix.fibre_type_id,
    }
    observed = {"mortar_curing_temperature_C": mix.curing_temperature_observed}
    g = np.zeros(len(schema.mortar_global), dtype=np.float32)
    mask = np.zeros(len(schema.mortar_global), dtype=np.float32)
    for i, name in enumerate(schema.mortar_global):
        if name in values:
            g[i] = values[name]
            mask[i] = 1.0 if observed.get(name, True) else 0.0
    return g, mask


def generate_mortar_graph(
    mix: MortarMixDesign,
    schema: SchemaSpec = DEFAULT_SCHEMA,
    seed: Optional[int] = None,
) -> MortarMicrograph:
    rng = np.random.default_rng(seed)
    sand_xy, sand_r = _sample_sand(mix, rng)
    paste_xy = _paste_grid(mix, sand_xy, sand_r)

    n_sand = len(sand_xy)
    n_paste = len(paste_xy)
    pts = np.concatenate([sand_xy, paste_xy], axis=0)
    adjacency = _voronoi_adjacency(pts, rve_size=mix.rve_size_mm, rng=rng, n_samples=3000)

    pad = schema.mortar_node_pad_dim
    n_total = n_sand + n_paste
    x = np.zeros((n_total, pad), dtype=np.float32)
    mask = np.zeros((n_total, pad), dtype=np.float32)
    node_type = np.zeros(n_total, dtype=np.int64)

    # Sand uses exact disk area; paste seeds split the remaining area by their
    # Monte-Carlo Voronoi cell share.
    mortar_rve_area = float(mix.rve_size_mm ** 2)
    sand_area_frac = (np.pi * np.asarray(sand_r, dtype=np.float64) ** 2) / mortar_rve_area
    paste_total_frac = max(0.0, 1.0 - float(sand_area_frac.sum()))
    paste_cell = _voronoi_cell_fraction(paste_xy, mix.rve_size_mm, rng, n_samples=3000)
    paste_area_frac = paste_cell * paste_total_frac

    for i in range(n_sand):
        v = _sand_feature_vector(
            sand_r[i], rng, schema, area_fraction=float(sand_area_frac[i]),
            sand_specific_gravity=mix.sand_specific_gravity,
        )
        x[i, : len(v)] = v
        mask[i, : len(v)] = 1.0
        node_type[i] = MORTAR_NODE_TYPE_SAND

    for j in range(n_paste):
        v = _paste_feature_vector(mix, schema, area_fraction=float(paste_area_frac[j]))
        x[n_sand + j, : len(v)] = v
        mask[n_sand + j, : len(v)] = 1.0
        node_type[n_sand + j] = MORTAR_NODE_TYPE_PASTE

    # Any unreported feature carried as NaN (e.g. cement grade) is MASKED (mask=0)
    # and zero-filled, so the encoders treat it as missing rather than a fake value.
    nan = np.isnan(x)
    mask[nan] = 0.0
    x = np.nan_to_num(x)

    edge_index: List[List[int]] = []
    edge_attr: List[List[float]] = []
    for a, b in adjacency:
        d = float(np.linalg.norm(pts[a] - pts[b]))
        same_phase = 1.0 if node_type[a] == node_type[b] else 0.0
        for src, dst in ((a, b), (b, a)):
            edge_index.append([src, dst])
            edge_attr.append([d, same_phase])

    if not edge_index:
        edge_index = [[0, 0]]
        edge_attr = [[0.0, 0.0]]

    g, g_observed = _mortar_global_vector(mix, schema)
    g_t = torch.tensor(g, dtype=torch.float32)
    g_mask = torch.tensor(g_observed, dtype=torch.float32)
    g_t = g_t * g_mask  # zero out unreported descriptors

    return MortarMicrograph(
        x=torch.tensor(x, dtype=torch.float32),
        x_mask=torch.tensor(mask, dtype=torch.float32),
        edge_index=torch.tensor(np.asarray(edge_index).T, dtype=torch.long),
        edge_attr=torch.tensor(edge_attr, dtype=torch.float32),
        node_type=torch.tensor(node_type, dtype=torch.long),
        global_features=g_t,
        global_mask=g_mask,
    )
