"""
demo_d1_gcn_node_classification.py
==================================

Session 10 — Graph Neural Networks: a 2-layer GCN, from scratch.

Concept
-------
A Graph Convolutional Network learns node embeddings by **message passing**:
each node's new representation is a function of its own features *and*
its neighbours' features. Stack two such layers and every node sees its
2-hop neighbourhood.

The forward pass for a layer is just one line of linear algebra:

    H' = sigma(  A_hat @ H @ W  )

where `A_hat` is the symmetrically normalised adjacency with self-loops
(Kipf & Welling 2017):

    A_tilde = A + I
    A_hat   = D_tilde^(-1/2) A_tilde D_tilde^(-1/2)

We re-use the same 3-community stochastic block model from the session 08
DeepWalk demo so you can compare directly:

  - DeepWalk (session 08): unsupervised, learns from random walks.
    Node features are *not used*. Inductive on new nodes? **No.**
  - GCN (this demo):       semi-supervised, learns from features + structure.
    A handful of labels is enough. Inductive on new nodes? **Yes.**

We give nodes **no real features at all** — `X = I` (identity matrix),
the canonical featureless-GCN setup — and label only **4 nodes per
community** (12 of 60). The GCN still learns to classify the remaining
48 nodes correctly: the signal comes entirely from message-passing over
the graph.

Run
---
    pip install torch numpy matplotlib networkx
    python demo_d1_gcn_node_classification.py
"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import networkx as nx

torch.manual_seed(0)
rng = np.random.default_rng(0)

# ---------------------------------------------------------------------------
# 1. The same 3-community graph used in the session-08 DeepWalk demo
# ---------------------------------------------------------------------------
N_COMMUNITIES = 3
NODES_PER_COMM = 20
P_IN, P_OUT = 0.20, 0.015
COMMUNITY_COLORS = ["#e63946", "#2a9d8f", "#457b9d"]

sizes = [NODES_PER_COMM] * N_COMMUNITIES
probs = np.full((N_COMMUNITIES, N_COMMUNITIES), P_OUT)
np.fill_diagonal(probs, P_IN)
G = nx.stochastic_block_model(sizes, probs, seed=0)
N = G.number_of_nodes()
community = np.array([G.nodes[v]["block"] for v in G.nodes])
print(f"Graph: {N} nodes, {G.number_of_edges()} edges, "
      f"{N_COMMUNITIES} planted communities")

# ---------------------------------------------------------------------------
# 2. Normalised adjacency  A_hat = D_tilde^(-1/2) (A + I) D_tilde^(-1/2)
# ---------------------------------------------------------------------------
A = nx.to_numpy_array(G) + np.eye(N)
D = A.sum(axis=1)
D_inv_sqrt = np.diag(1.0 / np.sqrt(D))
A_hat = torch.tensor(D_inv_sqrt @ A @ D_inv_sqrt, dtype=torch.float32)

# ---------------------------------------------------------------------------
# 3. Identity (featureless) GCN — no node features, structure is the signal
# ---------------------------------------------------------------------------
# When nodes carry no real features, the canonical fallback is to use
# X = I (one-hot per node).  Layer-1 weights then act as a *learnable
# embedding lookup* — but unlike DeepWalk, the lookups are immediately
# averaged with their neighbours by the GCN aggregation, so message
# passing kicks in from the very first layer.
F_DIM = N
X = torch.eye(N, dtype=torch.float32)

# ---------------------------------------------------------------------------
# 4. A 2-layer GCN
# ---------------------------------------------------------------------------
class GCN(nn.Module):
    def __init__(self, in_dim, hidden_dim, n_classes):
        super().__init__()
        self.W1 = nn.Linear(in_dim, hidden_dim, bias=False)
        self.W2 = nn.Linear(hidden_dim, n_classes, bias=False)

    def forward(self, A_hat, X):
        # Layer 1: aggregate neighbours, project, ReLU
        H = F.relu(A_hat @ self.W1(X))
        # Layer 2: aggregate, project to class logits
        logits = A_hat @ self.W2(H)
        return H, logits

HIDDEN = 16
model = GCN(F_DIM, hidden_dim=HIDDEN, n_classes=N_COMMUNITIES)
opt = torch.optim.Adam(model.parameters(), lr=0.05, weight_decay=5e-4)


def pca_2d(H, basis=None):
    """Project H to 2D with PCA. If basis is given, reuse it (so all snapshots
    share the same axes — otherwise PCA flips signs/axes between snapshots and
    the animation looks nonsense)."""
    Hc = H - H.mean(0, keepdims=True)
    if basis is None:
        _, _, Vt = np.linalg.svd(Hc, full_matrices=False)
        basis = Vt[:2].T
    return Hc @ basis, basis

# ---------------------------------------------------------------------------
# 5. Semi-supervised split: 4 labels per community
# ---------------------------------------------------------------------------
labeled_mask = np.zeros(N, dtype=bool)
for c in range(N_COMMUNITIES):
    in_comm = np.where(community == c)[0]
    chosen = rng.choice(in_comm, size=4, replace=False)
    labeled_mask[chosen] = True

train_idx = torch.tensor(np.where(labeled_mask)[0], dtype=torch.long)
labels = torch.tensor(community, dtype=torch.long)
print(f"Training on {labeled_mask.sum()} labelled nodes "
      f"({labeled_mask.sum() / N:.0%} of the graph)")

# ---------------------------------------------------------------------------
# 6. Train, snapshotting the hidden-layer embeddings
# ---------------------------------------------------------------------------
SNAPSHOTS = [0, 5, 25, 100, 300]
N_EPOCHS = SNAPSHOTS[-1] + 1

snaps_H = {}
losses, test_accs = [], []
for epoch in range(N_EPOCHS):
    if epoch in SNAPSHOTS:
        with torch.no_grad():
            H, _ = model(A_hat, X)
            snaps_H[epoch] = H.numpy().copy()

    H, logits = model(A_hat, X)
    loss = F.cross_entropy(logits[train_idx], labels[train_idx])
    opt.zero_grad()
    loss.backward()
    opt.step()
    losses.append(loss.item())

    with torch.no_grad():
        pred = logits.argmax(-1).numpy()
    test_accs.append((pred[~labeled_mask] == community[~labeled_mask]).mean())

# Project every snapshot onto the PCA basis of the *final* hidden state, so
# all panels share an axis system and the geometry is comparable across epochs.
_, basis = pca_2d(snaps_H[SNAPSHOTS[-1]])
snaps = {ep: pca_2d(H, basis=basis)[0] for ep, H in snaps_H.items()}

print(f"final loss = {losses[-1]:.3f}, "
      f"test accuracy on unlabelled nodes = {test_accs[-1]:.1%}")

# ---------------------------------------------------------------------------
# 7. Visualise: hidden-layer embeddings + loss / accuracy curves
# ---------------------------------------------------------------------------
fig = plt.figure(figsize=(4 * len(SNAPSHOTS), 6.0))
gs = fig.add_gridspec(2, len(SNAPSHOTS), height_ratios=[4, 1.1], hspace=0.6)

all_pts = np.vstack(list(snaps.values()))
lim = max(abs(all_pts.min()), abs(all_pts.max())) * 1.15

for col, ep in enumerate(SNAPSHOTS):
    ax = fig.add_subplot(gs[0, col])
    Z = snaps[ep]
    for c in range(N_COMMUNITIES):
        mask = (community == c) & ~labeled_mask
        ax.scatter(Z[mask, 0], Z[mask, 1],
                   c=COMMUNITY_COLORS[c], s=55, alpha=0.75,
                   edgecolor='black', linewidth=0.4,
                   label=f'community {c} (unlabelled)' if col == 0 else None)
        # labelled nodes drawn as bigger stars with a thick black border
        mask_lab = (community == c) & labeled_mask
        ax.scatter(Z[mask_lab, 0], Z[mask_lab, 1],
                   c=COMMUNITY_COLORS[c], s=180, marker='*',
                   edgecolor='black', linewidth=1.2,
                   label=f'community {c} (labelled)' if col == 0 else None)
    ax.set_title(f"Epoch {ep}", fontsize=11)
    ax.set_xlim(-lim, lim); ax.set_ylim(-lim, lim)
    ax.set_aspect('equal'); ax.grid(alpha=0.25)
    ax.axhline(0, c='gray', lw=0.5); ax.axvline(0, c='gray', lw=0.5)
    if col == 0:
        ax.legend(fontsize=6.5, loc='lower left', framealpha=0.85)

# loss + accuracy on twin axes
ax_loss = fig.add_subplot(gs[1, :])
ax_loss.plot(losses, c='#264653', label='train loss')
ax_loss.set_xlabel("epoch")
ax_loss.set_ylabel("cross-entropy loss", color='#264653')
ax_loss.tick_params(axis='y', labelcolor='#264653')
ax_loss.grid(alpha=0.3)

ax_acc = ax_loss.twinx()
ax_acc.plot(test_accs, c='#e76f51', label='test acc (unlabelled)')
ax_acc.set_ylabel("test accuracy", color='#e76f51')
ax_acc.tick_params(axis='y', labelcolor='#e76f51')
ax_acc.set_ylim(0, 1.05)

for ep in SNAPSHOTS:
    ax_loss.axvline(ep, c='gray', lw=0.6, alpha=0.6)

ax_loss.set_title("Training loss (left) and held-out community accuracy (right)")

fig.suptitle(
    "GCN learns community structure from graph topology alone\n"
    "(featureless setup: X = identity, only 12 labels)\n"
    f"({N} nodes, {N_COMMUNITIES} communities, "
    f"{labeled_mask.sum()} labels, hidden={HIDDEN}, 2-D view = PCA, "
    f"2-layer GCN, Kipf & Welling 2017)",
    fontsize=13,
)
fig.savefig("fig_gcn_evolution.png", dpi=130, bbox_inches='tight')
print("Saved → fig_gcn_evolution.png")
plt.show()