"""Concrete-scale edge-conditioned graph encoder."""

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_nn_conv


class ConcreteMesoGNN(nn.Module):
    """Encode the concrete proxy graph into a graph-level representation."""

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

        self.node_encoder = MissingFeatureEncoder(
            schema.node_pad_dim,
            hidden_dim,
            num_stat_groups=schema.num_node_types,
        )
        self.node_type_embed = nn.Embedding(schema.num_node_types, hidden_dim)
        self.edge_encoder = MissingFeatureEncoder(
            schema.edge_pad_dim,
            hidden_dim,
            num_stat_groups=schema.num_edge_types,
        )
        self.edge_type_embed = nn.Embedding(schema.num_edge_types, hidden_dim)
        self.global_encoder = MissingFeatureEncoder(schema.global_dim, hidden_dim)

        self.layers = nn.ModuleList()
        self.norms = nn.ModuleList()
        for _ in range(num_layers):
            edge_nn = nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim),
                nn.SiLU(),
                nn.Linear(hidden_dim, hidden_dim * hidden_dim),
            )
            self.layers.append(
                make_nn_conv(hidden_dim, hidden_dim, edge_nn, aggr="mean")
            )
            self.norms.append(nn.LayerNorm(hidden_dim))

        self.feature_dim = hidden_dim * 2

    def forward(
        self,
        x: Tensor,
        x_mask: Tensor,
        edge_index: Tensor,
        edge_attr: Tensor,
        edge_attr_mask: Tensor,
        node_type: Tensor,
        edge_type: Tensor,
        global_features: Tensor,
        global_mask: Tensor,
        batch: Tensor,
        num_graphs: int,
    ) -> Tensor:
        h = self.node_encoder(x, x_mask, node_type) + self.node_type_embed(node_type)
        e = self.edge_encoder(
            edge_attr, edge_attr_mask, edge_type
        ) + self.edge_type_embed(edge_type)

        for conv, norm in zip(self.layers, self.norms):
            h = torch.nn.functional.silu(norm(h + conv(h, edge_index, e)))

        pooled = global_mean_pool(h, batch, num_graphs=num_graphs)
        global_state = self.global_encoder(global_features, global_mask)
        return torch.cat([pooled, global_state], dim=-1)
