import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric.data import Data

class MedicalGNN(nn.Module):
    def __init__(self, in_channels: int, hidden_channels: int = 256):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, in_channels)
        self.dropout = nn.Dropout(p=0.3)

    def forward(self, x: torch.tensor, edge_index: torch.tensor) -> torch.tensor:
        h = self.conv1(x, edge_index)
        h = F.relu(h)
        h = self.dropout(h)
        h = self.conv2(h, edge_index)
        return h
    
def get_structural_embeddings(pyg_data: Data, in_channels: int, hidden_channels: int = 256) -> torch.tensor:
    model = MedicalGNN(in_channels = in_channels, hidden_channels = hidden_channels)
    model.eval()
    with torch.no_grad():
        structural_embs = model(pyg_data.x, pyg_data.edge_index)
        
    print(f"GNN message passing finish: shape={structural_embs.shape}")
    return structural_embs