"""PyTorch Geometric message-passing layers used by HiCon-GNN."""

from __future__ import annotations

from typing import Optional

from torch import Tensor, nn
from torch_geometric.nn import GINEConv, NNConv
from torch_geometric.nn import global_mean_pool as _global_mean_pool


def global_mean_pool(
    x: Tensor, batch: Tensor, num_graphs: Optional[int] = None
) -> Tensor:
    """Mean-pool node embeddings into one row per graph."""

    return _global_mean_pool(x, batch, size=num_graphs)


def make_gine_conv(mlp: nn.Module, edge_dim: int) -> nn.Module:
    """Create the mortar-scale GINE convolution."""

    return GINEConv(mlp, train_eps=True, edge_dim=edge_dim)


def make_nn_conv(
    in_dim: int,
    out_dim: int,
    edge_nn: nn.Module,
    aggr: str = "mean",
) -> nn.Module:
    """Create the concrete-scale edge-conditioned convolution."""

    return NNConv(in_dim, out_dim, edge_nn, aggr=aggr)
