#A Gentle Introduction to Graph Neural Networks in Python
#https://machinelearningmastery.com/a-gentle-introduction-to-graph-neural-networks-in-python/ 


# Defining a Graph Neural Network in Python

# In this introductory example of building a GNN, we will consider a small graph dataset 
# associated with a social media platform, where each node represents a person and each edge 
# connecting any two nodes is a friendship between persons. Furthermore, each node (person) 
# has associated features like the person’s age, their interests, etc.

# The target task of the GNN we will build is classifying people on either popular or not 
# popular in the social network (binary classification), depending on whether having more
# than two or less than two friends in it, and taking into account:
#     The person’s features, such as their interests
#     The person’s connections with other persons

# Therefore, GNNs give an extra layer of sophistication to predictive tasks, because they 
# not only look at the target instance’s features to make a prediction but also at its 
# relationship with other data instances, unlike classical classification and regression models.

# Do first:
# pip install ogb
# pip install torch_geometric
# pip install networkx

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

#This is our “mini-social network” dataset or graph:

edge_index = torch.tensor([
    [0, 1, 0, 2, 0, 4, 2, 4],
    [1, 0, 2, 0, 4, 0, 4, 2],
], dtype=torch.long)

# Basically, edge_index is a matrix of edges or connections between users. There are 5 
# users, numbered 0 to 4. The first connection is from user 0 to user 1, and we know this 
# by looking at the first element in each row of the matrix. The second connection is the 
# reciprocal of the previous one: user 1 to user 0. Then comes user 0 to user 2, and 
# so on. User 3 seems not to be connected to anyone yet!

# Now we model two numerical features for each person, in a tensor node_features: the 
# person’s age, and their interest in sports, with 1 indicating interest and 0 indicating
#  no interest.

# Define data features
node_features = torch.tensor([
    [25, 1],  # Person 0 (25 years old, likes sports)
    [30, 0],  # Person 1 (30 years old, does not like sports)
    [22, 1],  # Person 2 (22 years old, likes sports)
    [35, 0],  # Person 3 (35 years old, does not like sports)
    [27, 1],  # Person 4 (27 years old, likes sports)
], dtype=torch.float)


# Visualizing a Graph Neural Network in Python

# One way to visualize our graph neural network in Python can be accomplished by using 
# the NetworkX library. It will create a graph from the edge list and Matplotlib to 
# display it. An example of this is below.

import networkx as nx
import matplotlib.pyplot as plt

# Convert the edge_index tensor to a list of edge tuples
edge_list = edge_index.t().tolist()

# Create a NetworkX graph from the edge list
G = nx.Graph()
G.add_edges_from(edge_list)

# Optionally, include nodes that might be isolated (e.g., person 3)
G.add_nodes_from(range(node_features.size(0)))

# Generate a layout for the nodes
pos = nx.spring_layout(G, seed=42)  # fixed seed for reproducibility

# Draw the graph with labels
plt.figure(figsize=(6, 6))
nx.draw_networkx(G, pos, with_labels=True, node_color='lightblue', edge_color='gray', node_size=800)
plt.title("Visualization of the Social Network Graph")
plt.axis('off')
plt.show()


# Building a Graph Neural Network Model in Python

# Now we define labels for the dataset of users, i.e. whether a person is popular or 
# not, based on whether the person has more than 2 friends or not. The process entails 
# calculating the number of friends of each person (ground truth) based on the adjacency matrix.

# Define dataset labels
num_friends = torch.tensor([3, 1, 2, 0, 3])
labels = (num_friends >= 2).long()

# Using the following mask, we will indicate that the first three people will be 
# used as training data to build the GNN, and the other two will be used later 
# for inference. Finally, we also wrap everything into a Data object.

# Mask for separating training and testing data
train_mask = torch.tensor([1, 1, 1, 0, 0], dtype=torch.bool)
data = Data(x=node_features, edge_index=edge_index, y=labels, train_mask=train_mask)

# The next piece of code is crucial. It defines the GNN architecture and instantiates 
# the model. In PyTorch, GNN models can be built by using graph convolutional layers, 
# such as the ones implemented by the GCNConv class in torch_geometric.nn. Graph 
# convolutional layers aggregate information from a node’s neighbors, helping 
# learn representations that capture not only node features but also structural 
# relationships in the graph.

# Define model
class GNN(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(GNN, self).__init__()
        self.conv1 = GCNConv(input_dim, hidden_dim) 
        self.conv2 = GCNConv(hidden_dim, output_dim)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        x = F.relu(x)  # Activation function
        x = self.conv2(x, edge_index)
        return x

# Instantiate model
model = GNN(input_dim=2, hidden_dim=4, output_dim=2)


# Training a Graph Neural Network in Python

# The training model is reasonably similar to training other types of neural network 
# models in PyTorch:
    
# Define optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# Train model
for epoch in range(100):
    model.train()
    optimizer.zero_grad()
    
    out = model(data)
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    
    loss.backward()
    optimizer.step()

    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {loss.item():.4f}")


# Graph Neural Network Inference in Python

# Once the GNN has been trained, the inference process is straightforward. We pass 
# the full dataset to calculate popularity predictions, including the two users 
# that were not seen during training, and print the results. Notice that the 
# argmax function is used to obtain the class with the highest probability for 
# each user, from among the two available classes: this is the essence of 
# binary classifiers like logistic regressors.

# Test model
model.eval()
with torch.no_grad():
    predictions = model(data).argmax(dim=1)

print("\nFinal Predictions (1=Popular, 0=Not Popular):", predictions.tolist())


# Wrapping Up

# To sum up, we have built a very simple GNN that uses a graph representation of a 
# dataset to perform predictions based not only on the features of instances 
# (represented by nodes) but also by looking at the relationships or connections 
# with other instances.


   
