import torch_geometric as tg
from torch_geometric.nn import GCNConv, global_mean_pool

class GNN(torch.nn.Module):
    def __init__(self, node_features, hidden_dim, num_classes):
        super().__init__()
        self.conv1 = GCNConv(node_features, hidden_dim)
        self.conv2 = GCNConv(hidden_dim, hidden_dim)
        self.classifier = nn.Linear(hidden_dim, num_classes)
        
    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch
        
        # Message passing
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        
        # Graph-level pooling
        x = global_mean_pool(x, batch)
        
        # Classification
        return self.classifier(x)

# Example usage
dataset = tg.datasets.Planetoid(root='/tmp/Cora', name='Cora')
model = GNN(node_features=dataset.num_node_features, 
           hidden_dim=64, 
           num_classes=dataset.num_classes).to(device)

# Specialized DataLoader
loader = tg.data.DataLoader(dataset, batch_size=32, shuffle=True)

# Graph Attention Network (GAT)
class GAT(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = tg.nn.GATConv(in_channels, 8, heads=8, dropout=0.6)
        self.conv2 = tg.nn.GATConv(8*8, out_channels, heads=1, concat=False, dropout=0.6)
        
    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = F.dropout(x, p=0.6, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

# Graph Isomorphism Network (GIN)
class GIN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = tg.nn.GINConv(
            nn.Sequential(
                nn.Linear(in_channels, hidden_channels),
                nn.ReLU(),
                nn.Linear(hidden_channels, hidden_channels)
            ), train_eps=True)
        self.conv2 = tg.nn.GINConv(
            nn.Sequential(
                nn.Linear(hidden_channels, hidden_channels),
                nn.ReLU(),
                nn.Linear(hidden_channels, out_channels)
            ), train_eps=True)
        
    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        return x