"""Abrams-law-shaped phase strengths for paste and ITZ.

Replaces the free paste-compressive (``mortar_pred[:, 0]``) and ITZ-strength
(``itz_pred[:, 2]``) channels with a physically-structured Abrams form

    f_phase = A(z) * exp(-k(z) * wb)          # = A * B^(-w/b),  B = e^k

where ``A`` (the ceiling, i.e. strength at w/b -> 0) and ``k`` (the decay rate)
are *trained per mix*, conditioned on composition + age ``z`` -- crucially NOT on
w/b, so the w/b dependence stays a clean Abrams shape. Conditioning ``A`` on
composition lets silica-fume-rich (UHPC) mixes earn a higher ceiling, which a
single global Abrams constant cannot express.

The Abrams shape (monotone-decreasing, convex in w/b) is enforced *softly* via an
annealed monotonicity + convexity penalty rather than hard-clamping ``k`` >= 0.
This is the "relaxation": the model may bend the shape where the data demands it,
but is pulled back toward Abrams. The penalty is exact and cheap because w/b
enters ``f_phase`` only through the explicit kernel (``z`` excludes w/b).

The ITZ uses a wall-effect-linked local ratio ``rho * wb`` (rho >= 1), tying the
ITZ to the paste w/b with a single extra parameter; ``rho -> 1`` recovers a
densified (UHPC) ITZ. Paste and ITZ ceilings live in separate ranges so the ITZ
stays the weaker phase in normal concrete while still able to rise in UHPC.
"""

from __future__ import annotations

from typing import Dict, List

import torch
from torch import Tensor, nn

from ..missing_features import MissingFeatureEncoder
from ..schema import DEFAULT_SCHEMA, MORTAR_NODE_TYPE_PASTE, SchemaSpec


def _scatter_mean_2d(src: Tensor, index: Tensor, num_graphs: int) -> Tensor:
    """Mean of ``src`` rows grouped by ``index`` into ``num_graphs`` buckets."""

    out = torch.zeros(num_graphs, src.size(-1), device=src.device, dtype=src.dtype)
    cnt = torch.zeros(num_graphs, device=src.device, dtype=src.dtype)
    out.index_add_(0, index, src)
    cnt.index_add_(0, index, torch.ones(src.size(0), device=src.device, dtype=src.dtype))
    return out / cnt.clamp(min=1.0).unsqueeze(-1)


def _scatter_mean_1d(src: Tensor, index: Tensor, num_graphs: int) -> Tensor:
    """Mean of scalar ``src`` grouped by ``index`` into ``num_graphs`` buckets."""

    out = torch.zeros(num_graphs, device=src.device, dtype=src.dtype)
    cnt = torch.zeros(num_graphs, device=src.device, dtype=src.dtype)
    out.index_add_(0, index, src)
    cnt.index_add_(0, index, torch.ones_like(src))
    return out / cnt.clamp(min=1.0)


class AbramsPhaseStrength(nn.Module):
    """Composition-conditioned Abrams kernels for paste and ITZ strength."""

    # ``mix_itz`` layout (see real_data.mix_itz_vector):
    # [cement, scm_type_id, scm_fraction, w/b, admixture, rh, temp, age]
    WB_COL = 3            # ITZ w/b source
    SCM_TYPE_COL = 1
    NUMERIC_COLS = (0, 2, 4, 7)  # cement, scm_fraction, admixture, age  (NOT w/b)

    def __init__(
        self,
        schema: SchemaSpec = DEFAULT_SCHEMA,
        hidden: int = 64,
        num_scm_types: int = 8,
        emb_dim: int = 8,
        use_paste_chem: bool = True,
        a_paste=(20.0, 300.0),
        # ITZ strength channel = ITZ tensile/bond strength -> tensile-scale ceiling.
        a_itz=(0.5, 20.0),
        k_scale: float = 8.0,
        lam_cvx: float = 0.5,
        lam_order: float = 1.0,
        exp_clamp: float = 6.0,
    ):
        super().__init__()
        self.schema = schema
        self.use_paste_chem = use_paste_chem
        self.a_paste = a_paste
        self.a_itz = a_itz
        self.k_scale = k_scale
        self.lam_cvx = lam_cvx
        self.lam_order = lam_order
        self.exp_clamp = exp_clamp
        self.num_scm_types = num_scm_types

        # Paste chemistry used by the coefficient network.
        self.ox_cols: List[int] = []
        if use_paste_chem:
            for name in (
                "cement_CaO_pct",
                "cement_SiO2_pct",
                "scm_CaO_pct",
                "scm_SiO2_pct",
                "cement_type_id",
                "cement_grade_mpa",
            ):
                if name in schema.paste:
                    self.ox_cols.append(schema.paste.index(name))

        # Paste w/b column on the mortar micrograph paste nodes (separate from the
        # ITZ w/b in mix_itz). Falls back to the ITZ w/b if the schema lacks it.
        self.wb_paste_col = (
            schema.paste.index("water_to_binder_ratio")
            if "water_to_binder_ratio" in schema.paste
            else None
        )

        # mix_itz numerics (cement, scm_fraction, admixture, age) are always
        # observed -> BatchNorm. The pooled paste columns (oxides + cement type/
        # grade) go through a MissingFeatureEncoder so cement grade (~72% missing)
        # is MASKED with a learnable embedding, not imputed with a fake value.
        n_numeric = len(self.NUMERIC_COLS)
        self.numeric_bn = nn.BatchNorm1d(n_numeric)
        paste_hidden = 32
        if self.ox_cols:
            self.paste_encoder = MissingFeatureEncoder(len(self.ox_cols), paste_hidden)
            paste_out = paste_hidden
        else:
            self.paste_encoder = None
            paste_out = 0
        self.scm_embed = nn.Embedding(num_scm_types, emb_dim)

        z_dim = n_numeric + paste_out + emb_dim
        self.trunk = nn.Sequential(
            nn.Linear(z_dim, hidden),
            nn.SiLU(),
            nn.Linear(hidden, hidden),
            nn.SiLU(),
        )
        self.head_Ap = nn.Linear(hidden, 1)
        self.head_kp = nn.Linear(hidden, 1)
        self.head_Ai = nn.Linear(hidden, 1)
        self.head_ki = nn.Linear(hidden, 1)
        self.head_rho = nn.Linear(hidden, 1)

    def _build_z(self, batch) -> Tensor:
        mix = batch.mix_itz  # (G, 8) raw, per graph
        parts = [self.numeric_bn(mix[:, list(self.NUMERIC_COLS)])]
        if self.paste_encoder is not None:
            paste = batch.m_node_type == MORTAR_NODE_TYPE_PASTE
            cols = torch.tensor(self.ox_cols, device=mix.device, dtype=torch.long)
            idx = batch.m_batch[paste]
            pooled = _scatter_mean_2d(batch.m_x[paste][:, cols], idx, batch.num_graphs)
            pooled_mask = (
                _scatter_mean_2d(batch.m_mask[paste][:, cols], idx, batch.num_graphs) > 0
            ).to(pooled.dtype)
            parts.append(self.paste_encoder(pooled, pooled_mask))
        scm_id = mix[:, self.SCM_TYPE_COL].long().clamp(0, self.num_scm_types - 1)
        parts.append(self.scm_embed(scm_id))
        return torch.cat(parts, dim=-1)

    def _coeffs(self, z: Tensor):
        h = self.trunk(z)
        a_lo, a_hi = self.a_paste
        ai_lo, ai_hi = self.a_itz
        A_p = a_lo + (a_hi - a_lo) * torch.sigmoid(self.head_Ap(h)).squeeze(-1)
        A_i = ai_lo + (ai_hi - ai_lo) * torch.sigmoid(self.head_Ai(h)).squeeze(-1)
        # k is a free real bounded by tanh; the Abrams shape (k >= 0, convex) is
        # encouraged by the soft penalty, not forced here.
        k_p = self.k_scale * torch.tanh(self.head_kp(h)).squeeze(-1)
        k_i = self.k_scale * torch.tanh(self.head_ki(h)).squeeze(-1)
        rho = 1.0 + torch.nn.functional.softplus(self.head_rho(h)).squeeze(-1)
        return A_p, k_p, A_i, k_i, rho

    def _kernel(self, A: Tensor, k: Tensor, wb: Tensor) -> Tensor:
        # clamp the exponent for numerical safety; downstream clamp bounds the value.
        return A * torch.exp(torch.clamp(-k * wb, max=self.exp_clamp))

    def _paste_wb(self, batch) -> Tensor:
        """Per-graph paste w/b, mean-pooled over the paste nodes of the micrograph."""

        if self.wb_paste_col is None:
            return batch.mix_itz[:, self.WB_COL]
        paste_mask = batch.m_node_type == MORTAR_NODE_TYPE_PASTE
        vals = batch.m_x[paste_mask][:, self.wb_paste_col]
        return _scatter_mean_1d(vals, batch.m_batch[paste_mask], batch.num_graphs)

    def forward(self, batch) -> Dict[str, Tensor]:
        z = self._build_z(batch)
        A_p, k_p, A_i, k_i, rho = self._coeffs(z)
        # Separate w/b per phase: paste from the micrograph, ITZ from the mix vector.
        wb_paste = self._paste_wb(batch)
        wb_itz = batch.mix_itz[:, self.WB_COL]

        f_paste = self._kernel(A_p, k_p, wb_paste).clamp(5.0, 300.0)
        f_itz = self._kernel(A_i, k_i * rho, wb_itz).clamp(0.5, 20.0)  # tensile/bond scale

        if torch.is_grad_enabled():
            penalty = self._shape_penalty(A_p, k_p, A_i, k_i, rho, wb_paste, wb_itz)
            penalty = penalty + self.lam_order * torch.relu(f_itz - f_paste).pow(2).mean()
        else:
            penalty = wb_itz.new_zeros(())

        return {
            "f_paste": f_paste,
            "f_itz": f_itz,
            "physics_penalty": penalty,
        }

    def _pen_one(self, A: Tensor, k: Tensor, wb_p: Tensor) -> Tensor:
        f = self._kernel(A, k, wb_p)
        g1 = torch.autograd.grad(f.sum(), wb_p, create_graph=True)[0]
        g2 = torch.autograd.grad(g1.sum(), wb_p, create_graph=True)[0]
        mono = torch.relu(g1).pow(2).mean()  # penalize df/dwb > 0
        cvx = torch.relu(-g2).pow(2).mean()  # penalize d2f/dwb2 < 0
        return mono + self.lam_cvx * cvx

    def _shape_penalty(self, A_p, k_p, A_i, k_i, rho, wb_paste, wb_itz) -> Tensor:
        # w/b enters f only through the explicit kernel (z excludes w/b), so isolated
        # per-phase leaves give exact, cheap shape gradients. Paste and ITZ are
        # penalized against their own w/b.
        wb_p = wb_paste.detach().clone().requires_grad_(True)
        wb_i = wb_itz.detach().clone().requires_grad_(True)
        return self._pen_one(A_p, k_p, wb_p) + self._pen_one(A_i, k_i * rho, wb_i)
