from __future__ import annotations

from typing import Literal

import torch
import torch.nn as nn
from torch_geometric.nn import GINEConv, GPSConv


class GraphSpatialEncoder(nn.Module):
    """Per-frame spatial graph encoder built on top of PyG."""

    def __init__(
        self,
        node_in_dim: int,
        edge_in_dim: int,
        hidden_dim: int = 256,
        num_layers: int = 4,
        dropout: float = 0.1,
        backbone: Literal["gine", "gps"] = "gine",
        num_heads: int = 8,
    ):
        super().__init__()
        self.backbone = backbone
        self.node_proj = nn.Linear(node_in_dim, hidden_dim)
        self.edge_proj = nn.Linear(edge_in_dim, hidden_dim)
        self.dropout = nn.Dropout(dropout)

        self.layers = nn.ModuleList()
        for _ in range(num_layers):
            mlp = nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim),
                nn.SiLU(),
                nn.Linear(hidden_dim, hidden_dim),
            )
            if backbone == "gine":
                conv = GINEConv(mlp, edge_dim=hidden_dim)
            elif backbone == "gps":
                conv = GPSConv(
                    channels=hidden_dim,
                    conv=GINEConv(mlp, edge_dim=hidden_dim),
                    heads=num_heads,
                    attn_type="multihead",
                    attn_kwargs={"dropout": dropout},
                )
            else:
                raise ValueError(f"Unsupported graph backbone: {backbone}")
            self.layers.append(conv)

        self.norm = nn.LayerNorm(hidden_dim)

    def forward(self, data) -> torch.Tensor:
        x = self.node_proj(data.x)
        edge_attr = self.edge_proj(data.edge_attr)

        for layer in self.layers:
            residual = x
            if self.backbone == "gps":
                x = layer(x, data.edge_index, data.batch, edge_attr=edge_attr)
            else:
                x = layer(x, data.edge_index, edge_attr=edge_attr)
            x = self.norm(x + self.dropout(x) + residual)

        return x