"""GINEConv mortar model over sand and paste nodes."""

from __future__ import annotations

import torch
from torch import Tensor, nn

from ..missing_features import MissingFeatureEncoder
from ..schema import DEFAULT_SCHEMA, SchemaSpec
from .layers import global_mean_pool, make_gine_conv
from .outputs import transform_mortar_outputs


class MortarSubModel(nn.Module):
    """Sand/paste GNN that produces a mortar property vector per micrograph."""

    def __init__(
        self,
        schema: SchemaSpec = DEFAULT_SCHEMA,
        hidden_dim: int = 64,
        num_layers: int = 2,
    ):
        super().__init__()
        self.schema = schema

        self.node_encoder = MissingFeatureEncoder(
            schema.mortar_node_pad_dim, hidden_dim, num_stat_groups=schema.num_mortar_node_types
        )
        self.type_embed = nn.Embedding(schema.num_mortar_node_types, hidden_dim)

        self.edge_encoder = nn.Sequential(
            nn.Linear(2, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.SiLU(),
        )
        self.layers = nn.ModuleList(
            [
                make_gine_conv(
                    nn.Sequential(
                        nn.Linear(hidden_dim, hidden_dim),
                        nn.SiLU(),
                        nn.Linear(hidden_dim, hidden_dim),
                    ),
                    edge_dim=hidden_dim,
                )
                for _ in range(num_layers)
            ]
        )
        self.norms = nn.ModuleList(
            [nn.LayerNorm(hidden_dim) for _ in range(num_layers)]
        )
        self.global_encoder = MissingFeatureEncoder(schema.mortar_global_dim, hidden_dim)

        self.head = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, len(schema.mortar_targets)),
        )

    def forward(
        self,
        x: Tensor,
        x_mask: Tensor,
        edge_index: Tensor,
        edge_attr: Tensor,
        node_type: Tensor,
        batch: Tensor,
        num_graphs: int,
        global_features: Tensor,
        global_mask: Tensor,
    ) -> Tensor:
        h = self.node_encoder(x, x_mask, node_type) + self.type_embed(node_type)
        e = self.edge_encoder(edge_attr)
        for conv, norm in zip(self.layers, self.norms):
            h_new = conv(h, edge_index, e)
            h = norm(h + h_new)
            h = torch.nn.functional.silu(h)
        pooled = global_mean_pool(h, batch, num_graphs=num_graphs)
        g = self.global_encoder(global_features, global_mask)
        raw = self.head(torch.cat([pooled, g], dim=-1))
        return transform_mortar_outputs(raw)
