import torch
import torch.nn as nn

from model_gcn import GCNEncoder


class FusionModel(nn.Module):
    def __init__(
        self,
        text_dim: int,
        gcn_hidden_dim: int = 128,
        text_proj_dim: int = 128,
        fusion_hidden_dim: int = 128,
        num_labels: int = 2,
        dropout: float = 0.3,
    ) -> None:
        super().__init__()
        self.text_projection = nn.Sequential(
            nn.Linear(text_dim, text_proj_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
        self.graph_encoder = GCNEncoder(
            in_dim=text_dim,
            hidden_dim=gcn_hidden_dim,
            out_dim=text_proj_dim,
            dropout=dropout,
        )
        self.classifier = nn.Sequential(
            nn.Linear(text_proj_dim * 2, fusion_hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(fusion_hidden_dim, num_labels),
        )

    def forward(
        self,
        text_features: torch.Tensor,
        edge_index: torch.Tensor,
        edge_weight: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        h_text = self.text_projection(text_features)
        h_graph = self.graph_encoder(text_features, edge_index, edge_weight)
        fused = torch.cat([h_text, h_graph], dim=-1)
        logits = self.classifier(fused)
        return {
            "text_embeddings": h_text,
            "graph_embeddings": h_graph,
            "fused_embeddings": fused,
            "logits": logits,
        }