import networkx as nx

G = nx.karate_club_graph()

!pip install torch-geometric

from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import NormalizeFeatures
from torch_geometric.utils.convert import from_networkx

data = from_networkx(G)

from torch_geometric.utils import degree, to_undirected
from torch_geometric.transforms import ClusterData

# Convert the graph to an undirected graph, so we can aggregate neighbors in both directions.
G = to_undirected(G)

# Compute the degree of each node.
deg = degree(data.edge_index[0], num_nodes=data.num_nodes)

# Use METIS algorithm to partition the graph into multiple subgraphs.
cluster_data = ClusterData(data, num_parts=2, recursive=False, transform=NormalizeFeatures(),
                           degree=deg)

from torch_geometric.utils import random_walk

# Perform random walk sampling to obtain node neighbor samples.
walk_length = 20  # The length of random walk trail.
num_steps = 4     # The number of nodes to sample from each step.
data.batch = None
data.edge_index = to_undirected(data.edge_index)  # Use undirected edge for random walk.

rw_data = random_walk(data.edge_index, walk_length=walk_length, num_steps=num_steps)

import torch.nn.functional as F
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.nn import global_mean_pool

class GraphSageConv(MessagePassing):
    def __init__(self, in_channels, out_channels, aggr='mean'):
        super(GraphSageConv, self).__init__(aggr=aggr)
        self.lin = nn.Linear(in_channels, out_channels)
        
    def forward(self, x, edge_index):
        return self.propagate(edge_index, x=x)
    
    def message(self, x_j):
        return x_j
    
    def update(self, aggr_out, x):
        return F.relu(self.lin(torch.cat([x, aggr_out], dim=1)))

class MeanAggregator(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(MeanAggregator, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.lin = nn.Linear(input_dim, output_dim)
        
    def forward(self, neigh_mean):
        out = F.relu(self.lin(neigh_mean))
        return out

class GraphSAGE(nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers=2):
        super(GraphSAGE, self).__init__()
        self.conv1 = GraphSageConv(in_channels, hidden_channels)
        self.aggreg1 = MeanAggregator(hidden_channels, hidden_channels)
        self.conv2 = GraphSageConv(hidden_channels, out_channels)
        
    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = global_mean_pool(x, edge_index)  # Compute global mean over nodes.
        x = self.aggreg1(x)
        x = self.conv2(x, edge_index)
        return x

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

num_epochs = 100

for epoch in range(num_epochs):
    model.train()
    
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = criterion(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    
    print('Epoch {:03d}, Loss: {:.4f}'.format(epoch, loss.item()))

model.eval()

with torch.no_grad():
    pred = model(data.x, data.edge_index)
    pred = pred.argmax(dim=1)

acc = (pred[data.test_mask] == data.y[data.test_mask]).sum().item() / data.test_mask.sum().item()
print('Test accuracy: {:.4f}'.format(acc))