"""HiCon-GNN and its structure-blind ANN baseline."""

from __future__ import annotations

from typing import Dict, Tuple

import torch
from torch import Tensor, nn

from ..schema import (
    DEFAULT_SCHEMA,
    EDGE_TYPE_ITZ,
    MORTAR_NODE_TYPE_PASTE,
    NODE_TYPE_AGGREGATE,
    NODE_TYPE_MORTAR,
    SchemaSpec,
)
from .abrams import AbramsPhaseStrength
from .concrete_gnn import ConcreteMesoGNN
from .itz_model import ITZSubModel
from .mortar_gnn import MortarSubModel
from .paste_model import PasteSubModel
from .outputs import transform_concrete_outputs

# ITZ porosity equals paste porosity times a bounded wall-effect ratio.
ITZ_PORO_R_MIN, ITZ_PORO_R_MAX = 1.1, 1.5
ITZ_PORO_CAP = 0.70
PASTE_POROSITY_SLOT = 4


def _slot_indices(target_names, feature_names):
    return [feature_names.index(n) for n in target_names if n in feature_names]


def _mortar_slot_indices(schema: SchemaSpec):
    return _slot_indices(schema.mortar_targets, schema.mortar)


def _itz_slot_indices(schema: SchemaSpec):
    return _slot_indices(schema.itz_targets, schema.itz)


def _paste_slot_indices(schema: SchemaSpec):
    return _slot_indices(schema.paste_targets, schema.paste)


def _paste_composition_dim(schema: SchemaSpec) -> int:
    """Number of paste-composition columns before the property slots."""

    n_prop = sum(1 for n in schema.paste_targets if n in schema.paste)
    return len(schema.paste) - n_prop


def pool_paste_composition(
    m_x: Tensor,
    m_mask: Tensor,
    m_node_type: Tensor,
    m_batch: Tensor,
    num_graphs: int,
    comp_dim: int,
) -> Tuple[Tensor, Tensor]:
    """Mean paste composition and observed mask for each graph."""

    paste = m_node_type == MORTAR_NODE_TYPE_PASTE
    xc = m_x[paste][:, :comp_dim]
    mc = m_mask[paste][:, :comp_dim]
    idx = m_batch[paste]
    out = torch.zeros(num_graphs, comp_dim, device=m_x.device, dtype=m_x.dtype)
    cnt = torch.zeros(num_graphs, device=m_x.device, dtype=m_x.dtype)
    out.index_add_(0, idx, xc)
    cnt.index_add_(0, idx, torch.ones(xc.size(0), device=m_x.device, dtype=m_x.dtype))
    pooled = out / cnt.clamp(min=1.0).unsqueeze(-1)
    mout = torch.zeros(num_graphs, comp_dim, device=m_x.device, dtype=m_x.dtype)
    mout.index_add_(0, idx, mc)
    pooled_mask = (mout > 0).to(m_x.dtype)
    return pooled, pooled_mask


def inject_paste(
    x: Tensor,
    x_mask: Tensor,
    paste_pred: Tensor,
    node_type: Tensor,
    batch: Tensor,
    slot_indices: list,
) -> Tuple[Tensor, Tensor]:
    """Write paste predictions into paste-node property slots."""

    paste_nodes = (node_type == MORTAR_NODE_TYPE_PASTE).to(x.dtype)
    pred_per_node = paste_pred[batch]
    injected = torch.zeros_like(x)
    inj_mask = torch.zeros_like(x)
    for k, idx in enumerate(slot_indices):
        injected[:, idx] = pred_per_node[:, k] * paste_nodes
        inj_mask[:, idx] = paste_nodes
    x_new = x * (1.0 - inj_mask) + injected
    x_mask_new = torch.maximum(x_mask, inj_mask)
    return x_new, x_mask_new


def inject_mortar(
    x: Tensor,
    x_mask: Tensor,
    mortar_pred: Tensor,
    node_type: Tensor,
    batch: Tensor,
    slot_indices: list,
) -> Tuple[Tensor, Tensor]:
    """Write mortar predictions into mortar-node feature slots."""

    mortar_nodes = (node_type == NODE_TYPE_MORTAR).to(x.dtype)
    pred_per_node = mortar_pred[batch]
    injected = torch.zeros_like(x)
    inj_mask = torch.zeros_like(x)
    for k, idx in enumerate(slot_indices):
        injected[:, idx] = pred_per_node[:, k] * mortar_nodes
        inj_mask[:, idx] = mortar_nodes
    x_new = x * (1.0 - inj_mask) + injected
    x_mask_new = torch.maximum(x_mask, inj_mask)
    return x_new, x_mask_new


def inject_itz(
    edge_attr: Tensor,
    edge_attr_mask: Tensor,
    itz_pred: Tensor,
    itz_edge_mask: Tensor,
    slot_indices: list,
) -> Tuple[Tensor, Tensor]:
    """Write ITZ predictions into ITZ-edge feature slots."""

    itz_idx = itz_edge_mask.nonzero(as_tuple=True)[0]
    injected = torch.zeros_like(edge_attr)
    inj_mask = torch.zeros_like(edge_attr)
    for k, idx in enumerate(slot_indices):
        injected[itz_idx, idx] = itz_pred[:, k]
        inj_mask[itz_idx, idx] = 1.0
    edge_new = edge_attr * (1.0 - inj_mask) + injected
    edge_mask_new = torch.maximum(edge_attr_mask, inj_mask)
    return edge_new, edge_mask_new


def build_itz_inputs(
    edge_index: Tensor,
    edge_type: Tensor,
    node_type: Tensor,
    x: Tensor,
    pos: Tensor,
    batch_per_node: Tensor,
    mix_itz_per_graph: Tensor,
    schema: SchemaSpec,
    input_dim: int,
) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
    """Assemble ITZ sub-model inputs for every ITZ edge."""

    itz_edge_mask = edge_type == EDGE_TYPE_ITZ
    src_all = edge_index[0]
    edge_graph = batch_per_node[src_all]
    itz_src = edge_index[0][itz_edge_mask]
    itz_dst = edge_index[1][itz_edge_mask]
    itz_graph = edge_graph[itz_edge_mask]
    n_itz = int(itz_edge_mask.sum().item())

    features = torch.zeros(n_itz, input_dim, device=x.device, dtype=x.dtype)
    masks = torch.zeros_like(features)
    if n_itz == 0:
        return features, masks, itz_edge_mask, itz_graph

    mix_dim = mix_itz_per_graph.size(1)
    features[:, :mix_dim] = mix_itz_per_graph[itz_graph]

    # Endpoint distance in millimetres.
    edge_dist = torch.linalg.norm(pos[itz_src] - pos[itz_dst], dim=-1)
    features[:, mix_dim] = edge_dist

    # Surface texture from the aggregate endpoint.
    src_is_agg = node_type[itz_src] == NODE_TYPE_AGGREGATE
    agg_nodes = torch.where(src_is_agg, itz_src, itz_dst)
    surface_texture_col = schema.aggregate.index("surface_texture")
    features[:, mix_dim + 1] = x[agg_nodes, surface_texture_col]

    masks[:, : mix_dim + 2] = 1.0
    # Preserve missing inputs in the observation mask.
    nan = torch.isnan(features)
    masks[nan] = 0.0
    features = torch.nan_to_num(features)
    return features, masks, itz_edge_mask, itz_graph


def aggregate_itz_per_graph(
    itz_pred: Tensor, itz_graph: Tensor, num_graphs: int
) -> Tensor:
    """Mean ITZ prediction per graph; zero for graphs without ITZ edges."""

    if itz_pred.numel() == 0:
        return torch.zeros(
            num_graphs, itz_pred.size(-1) if itz_pred.dim() == 2 else 0,
            device=itz_pred.device, dtype=itz_pred.dtype,
        )
    n_targets = itz_pred.size(-1)
    out = torch.zeros(num_graphs, n_targets, device=itz_pred.device, dtype=itz_pred.dtype)
    counts = torch.zeros(num_graphs, device=itz_pred.device, dtype=itz_pred.dtype)
    out.index_add_(0, itz_graph, itz_pred)
    counts.index_add_(0, itz_graph, torch.ones_like(itz_graph, dtype=itz_pred.dtype))
    return out / counts.clamp(min=1.0).unsqueeze(-1)


def _scatter_mean_1d(src: Tensor, index: Tensor, num_graphs: int) -> Tensor:
    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)


def _scatter_amin_1d(src: Tensor, index: Tensor, num_graphs: int) -> Tensor:
    base = src.new_full((num_graphs,), float("inf"))
    out = base.scatter_reduce(0, index, src, reduce="amin", include_self=True)
    return torch.where(torch.isinf(out), torch.zeros_like(out), out)


def _scatter_amin_2d(src: Tensor, index: Tensor, num_graphs: int) -> Tensor:
    base = src.new_full((num_graphs, src.size(-1)), float("inf"))
    out = base.scatter_reduce(
        0, index.unsqueeze(-1).expand_as(src), src, reduce="amin", include_self=True
    )
    return torch.where(torch.isinf(out), torch.zeros_like(out), out)


def build_itz_summary(
    itz_pred: Tensor,
    itz_graph: Tensor,
    itz_mean: Tensor,
    edge_index: Tensor,
    edge_type: Tensor,
    pos: Tensor,
    node_type: Tensor,
    x: Tensor,
    batch_per_node: Tensor,
    num_graphs: int,
    schema: SchemaSpec,
) -> Tensor:
    """Summarize ITZ properties, geometry and aggregate contacts per graph."""

    n_itz_targets = len(schema.itz_targets)
    if itz_pred.numel() > 0:
        itz_min = _scatter_amin_2d(itz_pred, itz_graph, num_graphs)
    else:
        itz_min = torch.zeros(num_graphs, n_itz_targets, device=x.device, dtype=x.dtype)

    # ITZ edge distances.
    itz_edge_mask = edge_type == EDGE_TYPE_ITZ
    isrc = edge_index[0][itz_edge_mask]
    idst = edge_index[1][itz_edge_mask]
    if isrc.numel() > 0:
        dist = torch.linalg.norm(pos[isrc] - pos[idst], dim=-1)
        dist_mean = _scatter_mean_1d(dist, itz_graph, num_graphs)
        dist_min = _scatter_amin_1d(dist, itz_graph, num_graphs)
        n_itz_per_graph = torch.zeros(num_graphs, device=x.device, dtype=x.dtype)
        n_itz_per_graph.index_add_(0, itz_graph, torch.ones_like(dist))
    else:
        dist_mean = torch.zeros(num_graphs, device=x.device, dtype=x.dtype)
        dist_min = torch.zeros(num_graphs, device=x.device, dtype=x.dtype)
        n_itz_per_graph = torch.zeros(num_graphs, device=x.device, dtype=x.dtype)

    # ITZ edges per aggregate node.
    agg_mask = node_type == NODE_TYPE_AGGREGATE
    agg_graph = batch_per_node[agg_mask]
    n_agg = torch.zeros(num_graphs, device=x.device, dtype=x.dtype)
    n_agg.index_add_(0, agg_graph, torch.ones(int(agg_mask.sum()), device=x.device, dtype=x.dtype))
    contact_density = n_itz_per_graph / n_agg.clamp(min=1.0)

    # Mean aggregate surface texture.
    st_col = schema.aggregate.index("surface_texture")
    st_mean = _scatter_mean_1d(x[agg_mask, st_col], agg_graph, num_graphs)

    return torch.cat(
        [
            itz_mean,
            itz_min,
            contact_density.unsqueeze(-1),
            dist_mean.unsqueeze(-1),
            dist_min.unsqueeze(-1),
            st_mean.unsqueeze(-1),
        ],
        dim=-1,
    )


class ITZWeakLink(nn.Module):
    """Compute ITZ efficiency from bond strength, thickness and interface state."""

    def __init__(self, itz_summary_dim: int, hidden_dim: int = 32):
        super().__init__()
        self.a = nn.Parameter(torch.zeros(()))
        self.c = nn.Parameter(torch.zeros(()))
        self.d = nn.Parameter(torch.tensor(2.0))  # eta_itz ~0.88 at init
        self.residual = nn.Sequential(
            nn.Linear(itz_summary_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 1)
        )

    def forward(self, itz_mean: Tensor, mortar_ref: Tensor, itz_summary: Tensor) -> Tensor:
        itz_thickness = itz_mean[:, 0]  # ITZ_TARGETS slot 0
        itz_strength = itz_mean[:, 2]   # ITZ_TARGETS slot 2 (tensile/bond, MPa)
        has_itz = (itz_thickness > 1.0).to(itz_mean.dtype)  # 0 thickness == no ITZ edges
        a = torch.nn.functional.softplus(self.a)
        c = torch.nn.functional.softplus(self.c)
        ratio_deficit = itz_strength / mortar_ref.clamp(min=0.1) - 1.0
        logit = a * ratio_deficit - c * itz_thickness / 70.0 + self.d
        res = torch.tanh(self.residual(itz_summary)).squeeze(-1)
        eta = torch.sigmoid(logit + res)
        return has_itz * eta + (1.0 - has_itz)  # eta=1 where there is no interface


class MortarCapacityStrengthHead(nn.Module):
    """Compute compression as mortar capacity times ITZ and matrix efficiencies."""

    MATRIX_GAIN_MAX = 1.5

    def __init__(
        self,
        feature_dim: int,
        mortar_other_dim: int,
        itz_summary_dim: int,
        hidden_dim: int = 64,
    ):
        super().__init__()
        self.mortar_eff = nn.Sequential(
            nn.Linear(mortar_other_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 1)
        )
        self.itz_link = ITZWeakLink(itz_summary_dim)
        self.matrix_eff = nn.Sequential(
            nn.Linear(feature_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 1)
        )

    def forward(
        self,
        mortar_pred: Tensor,
        itz_mean: Tensor,
        itz_summary: Tensor,
        graph_features: Tensor,
    ) -> Tensor:
        eta_mortar = torch.sigmoid(self.mortar_eff(mortar_pred[:, 1:])).squeeze(-1)
        capacity = mortar_pred[:, 0] * eta_mortar
        eta_itz = self.itz_link(itz_mean, mortar_pred[:, 1], itz_summary)
        eta_matrix = self.MATRIX_GAIN_MAX * torch.sigmoid(
            self.matrix_eff(graph_features)
        ).squeeze(-1)
        return capacity * eta_itz * eta_matrix


class PropertyCapacityHead(nn.Module):
    """Reduce a mortar property by ITZ and matrix efficiencies."""

    def __init__(self, feature_dim: int, itz_summary_dim: int, hidden_dim: int = 64):
        super().__init__()
        self.itz_link = ITZWeakLink(itz_summary_dim)
        self.matrix_eff = nn.Sequential(
            nn.Linear(feature_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 1)
        )

    def forward(
        self,
        mortar_property: Tensor,
        mortar_ref: Tensor,
        itz_mean: Tensor,
        itz_summary: Tensor,
        graph_features: Tensor,
    ) -> Tensor:
        eta_itz = self.itz_link(itz_mean, mortar_ref, itz_summary)
        eta_matrix = torch.sigmoid(self.matrix_eff(graph_features)).squeeze(-1)
        return mortar_property * eta_itz * eta_matrix


class IntegratedMultiscaleModel(nn.Module):
    """Published paste-mortar-ITZ-concrete hierarchy."""

    def __init__(
        self,
        schema: SchemaSpec = DEFAULT_SCHEMA,
        hidden_dim: int = 96,
        num_layers: int = 2,
        itz_input_dim: int = 32,
    ):
        super().__init__()
        self.schema = schema
        self.itz_input_dim = itz_input_dim
        self.mortar_slot_indices = _mortar_slot_indices(schema)
        self.itz_slot_indices = _itz_slot_indices(schema)
        self.mortar = MortarSubModel(schema=schema, hidden_dim=64, num_layers=2)
        self.itz = ITZSubModel(
            schema=schema,
            hidden_dim=96,
            input_dim=itz_input_dim,
        )
        self.mesoscale = ConcreteMesoGNN(
            schema=schema,
            hidden_dim=hidden_dim,
            num_layers=num_layers,
        )
        itz_summary_dim = 2 * len(schema.itz_targets) + 4
        self.strength_head = MortarCapacityStrengthHead(
            feature_dim=self.mesoscale.feature_dim,
            mortar_other_dim=len(schema.mortar_targets) - 1,
            itz_summary_dim=itz_summary_dim,
        )
        self.tensile_head = PropertyCapacityHead(
            feature_dim=self.mesoscale.feature_dim, itz_summary_dim=itz_summary_dim
        )
        self.flexural_head = PropertyCapacityHead(
            feature_dim=self.mesoscale.feature_dim, itz_summary_dim=itz_summary_dim
        )
        self.paste_abrams = AbramsPhaseStrength(schema=schema)
        self.paste_slot_indices = _paste_slot_indices(schema)
        if len(self.paste_slot_indices) != len(schema.paste_targets):
            raise ValueError("The paste schema must include every paste property slot")
        self.paste_comp_dim = _paste_composition_dim(schema)
        self.paste_sub = PasteSubModel(
            schema=schema, hidden_dim=64, input_dim=self.paste_comp_dim
        )

    def forward(self, batch) -> Dict[str, Tensor]:
        paste_abrams_out = self.paste_abrams(batch)
        paste_comp, paste_comp_mask = pool_paste_composition(
            batch.m_x,
            batch.m_mask,
            batch.m_node_type,
            batch.m_batch,
            batch.num_graphs,
            self.paste_comp_dim,
        )
        paste_aux = self.paste_sub(paste_comp, paste_comp_mask)
        paste_pred = torch.cat(
            [paste_abrams_out["f_paste"].unsqueeze(-1), paste_aux], dim=-1
        )
        m_x, m_mask = inject_paste(
            batch.m_x,
            batch.m_mask,
            paste_pred,
            batch.m_node_type,
            batch.m_batch,
            self.paste_slot_indices,
        )

        mortar_pred = self.mortar(
            x=m_x,
            x_mask=m_mask,
            edge_index=batch.m_edge_index,
            edge_attr=batch.m_edge_attr,
            node_type=batch.m_node_type,
            batch=batch.m_batch,
            num_graphs=batch.num_graphs,
            global_features=batch.mortar_global,
            global_mask=batch.mortar_global_mask,
        )

        x, x_mask = inject_mortar(
            batch.x,
            batch.x_mask,
            mortar_pred,
            batch.node_type,
            batch.batch,
            slot_indices=self.mortar_slot_indices,
        )

        itz_feat, itz_mask, itz_edge_mask, itz_graph = build_itz_inputs(
            edge_index=batch.edge_index,
            edge_type=batch.edge_type,
            node_type=batch.node_type,
            x=batch.x,
            pos=batch.pos,
            batch_per_node=batch.batch,
            mix_itz_per_graph=batch.mix_itz,
            schema=self.schema,
            input_dim=self.itz_input_dim,
        )
        itz_wb_penalty = None
        if itz_feat.size(0) > 0:
            itz_aux = self.itz(itz_feat, itz_mask)
            if torch.is_grad_enabled():
                itz_wb_penalty = self._itz_wb_mono_penalty(itz_feat, itz_mask)
            itz_strength = paste_abrams_out["f_itz"][itz_graph].unsqueeze(-1)
            itz_pred = torch.cat(
                [itz_aux[:, :2], itz_strength, itz_aux[:, 2:3]], dim=-1
            )
            # The ITZ remains more porous than its adjacent paste.
            paste_poro_edge = paste_pred[:, PASTE_POROSITY_SLOT][itz_graph]
            sigma_poro = ((itz_pred[:, 1] - 0.05) / 0.65).clamp(0.0, 1.0)
            r_itz = ITZ_PORO_R_MIN + (ITZ_PORO_R_MAX - ITZ_PORO_R_MIN) * sigma_poro
            itz_pred[:, 1] = (paste_poro_edge * r_itz).clamp(0.05, ITZ_PORO_CAP)
            edge_attr, edge_attr_mask = inject_itz(
                batch.edge_attr,
                batch.edge_attr_mask,
                itz_pred,
                itz_edge_mask,
                slot_indices=self.itz_slot_indices,
            )
        else:
            itz_pred = torch.zeros(
                0, len(self.schema.itz_targets), device=batch.edge_attr.device
            )
            edge_attr, edge_attr_mask = batch.edge_attr, batch.edge_attr_mask

        graph_features = self.mesoscale(
            x=x,
            x_mask=x_mask,
            edge_index=batch.edge_index,
            edge_attr=edge_attr,
            edge_attr_mask=edge_attr_mask,
            node_type=batch.node_type,
            edge_type=batch.edge_type,
            global_features=batch.global_features,
            global_mask=batch.global_mask,
            batch=batch.batch,
            num_graphs=batch.num_graphs,
        )
        itz_mean = aggregate_itz_per_graph(itz_pred, itz_graph, batch.num_graphs)
        itz_summary = build_itz_summary(
            itz_pred,
            itz_graph,
            itz_mean,
            batch.edge_index,
            batch.edge_type,
            batch.pos,
            batch.node_type,
            batch.x,
            batch.batch,
            batch.num_graphs,
            self.schema,
        )
        compressive = self.strength_head(
            mortar_pred, itz_mean, itz_summary, graph_features
        )
        tensile = self.tensile_head(
            mortar_pred[:, 1],
            mortar_pred[:, 1],
            itz_mean,
            itz_summary,
            graph_features,
        )
        flexural = self.flexural_head(
            mortar_pred[:, 2],
            mortar_pred[:, 1],
            itz_mean,
            itz_summary,
            graph_features,
        )
        concrete_pred = torch.stack([compressive, tensile, flexural], dim=-1)
        out = {
            "concrete_pred": concrete_pred,
            "mortar_pred": mortar_pred,
            "itz_pred": itz_pred,
            "itz_pred_per_graph": itz_mean,
            "paste_pred": paste_pred,
            "paste_abrams_penalty": paste_abrams_out["physics_penalty"],
            "itz_wb_penalty": itz_wb_penalty,
        }
        return out

    def _itz_wb_mono_penalty(self, itz_feat: Tensor, itz_mask: Tensor) -> Tensor:
        """Penalize decreasing ITZ thickness or porosity as w/b rises."""
        c = AbramsPhaseStrength.WB_COL
        wb = itz_feat[:, c : c + 1].detach().clone().requires_grad_(True)
        feat = torch.cat(
            [itz_feat[:, :c].detach(), wb, itz_feat[:, c + 1 :].detach()], dim=1
        )
        pred = self.itz(feat, itz_mask)
        pen = wb.new_zeros(())
        for ch in (0, 1):
            g = torch.autograd.grad(pred[:, ch].sum(), wb, create_graph=True)[0]
            pen = pen + torch.relu(-g).pow(2).mean()
        return pen


class TabularANN(nn.Module):
    """Structure-blind ANN over the publication mix descriptors."""

    def __init__(
        self,
        in_dim: int,
        schema: SchemaSpec = DEFAULT_SCHEMA,
        hidden_dim: int = 96,
    ):
        super().__init__()
        self.schema = schema
        self.in_dim = in_dim
        self.input_bn = nn.BatchNorm1d(in_dim)
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, len(schema.concrete_targets)),
        )

    def forward(self, batch) -> Dict[str, Tensor]:
        x = self.input_bn(batch.tabular_mix)
        raw = self.net(x)
        return {"concrete_pred": transform_concrete_outputs(raw)}
