"""
demo_d1_graph_node_embeddings.py
================================

Session 08 — Graphs and node embeddings.

Concept
-------
A node embedding turns each node into a vector so that **graph proximity
becomes geometric proximity**. Same headline trick as the user/movie demo
in session 07 — only now the "co-occurrence" comes from random walks on
the graph instead of co-rated items.

Recipe (DeepWalk-style):

    1. Generate many random walks on the graph.
    2. Treat each walk as a "sentence", each node as a "word".
    3. Train a skip-gram model (with negative sampling) on (target, context)
       pairs from the walks.
    4. The learned embedding has the property that nodes which appear in
       similar neighbourhoods land near each other in vector space.

We use a synthetic graph with 3 planted communities (stochastic block
model) so we can verify that the embeddings recover community structure
**without ever telling the model which node belongs to which community**.

The demo produces two figures:

    fig1_graph_representations.png  — the same graph shown three ways
    fig2_embeddings_evolution.png   — embeddings learning over training

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

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

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

# ---------------------------------------------------------------------------
# 1. A synthetic graph with 3 planted communities
# ---------------------------------------------------------------------------
N_COMMUNITIES = 3
NODES_PER_COMM = 20
P_IN, P_OUT = 0.20, 0.015          # within / between community edge prob
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])
node_color = [COMMUNITY_COLORS[c] for c in community]

print(f"Graph: {N} nodes, {G.number_of_edges()} edges, "
      f"{N_COMMUNITIES} planted communities")
print(f"  avg degree = {2 * G.number_of_edges() / N:.2f}")

# ---------------------------------------------------------------------------
# 2. Figure 1 — three representations of the same graph
# ---------------------------------------------------------------------------
fig1, axes = plt.subplots(1, 3, figsize=(15, 4.5))

# 2a. Force-directed layout, coloured by community
pos = nx.spring_layout(G, seed=0, k=0.5, iterations=80)
ax = axes[0]
nx.draw_networkx_edges(G, pos, ax=ax, alpha=0.25, width=0.7)
nx.draw_networkx_nodes(G, pos, ax=ax, node_color=node_color,
                       node_size=110, edgecolors='black', linewidths=0.5)
ax.set_title("(a) Picture: force-directed layout\n"
             "structure is visible to the eye")
ax.set_axis_off()
ax.set_aspect('equal')

# 2b. Adjacency matrix, nodes ordered by community → block structure
order = np.argsort(community, kind='stable')
A = nx.to_numpy_array(G)[order][:, order]
ax = axes[1]
ax.imshow(A, cmap='Greys', interpolation='nearest', aspect='equal')
# mark community boundaries
boundaries = np.cumsum([NODES_PER_COMM] * N_COMMUNITIES)[:-1]
for b in boundaries:
    ax.axhline(b - 0.5, color=COMMUNITY_COLORS[0], lw=0.7, alpha=0.7)
    ax.axvline(b - 0.5, color=COMMUNITY_COLORS[0], lw=0.7, alpha=0.7)
ax.set_title("(b) Adjacency matrix\n"
             "block structure when nodes ordered by community")
ax.set_xlabel("node"); ax.set_ylabel("node")

# 2c. Degree distribution
degrees = np.array([d for _, d in G.degree()])
ax = axes[2]
for c in range(N_COMMUNITIES):
    ax.hist(degrees[community == c], bins=range(0, degrees.max() + 2),
            alpha=0.65, color=COMMUNITY_COLORS[c],
            edgecolor='black', linewidth=0.4,
            label=f'community {c}')
ax.axvline(degrees.mean(), color='black', lw=1, linestyle='--',
           label=f'mean = {degrees.mean():.1f}')
ax.set_title("(c) Degree distribution\n"
             "ki = number of edges adjacent to node i")
ax.set_xlabel("degree"); ax.set_ylabel("count")
ax.legend(fontsize=8)
ax.grid(alpha=0.3)

fig1.suptitle("The same graph, three representations", fontsize=13)
fig1.tight_layout()
fig1.savefig("fig1_graph_representations.png", dpi=130, bbox_inches='tight')
print("Saved → fig1_graph_representations.png")

# ---------------------------------------------------------------------------
# 3. Random walks  (the "sentences" for skip-gram)
# ---------------------------------------------------------------------------
def random_walks(G, walks_per_node=10, walk_length=20, seed=0):
    rng = np.random.default_rng(seed)
    walks = []
    nodes = list(G.nodes())
    for _ in range(walks_per_node):
        rng.shuffle(nodes)
        for v in nodes:
            walk = [v]
            for _ in range(walk_length - 1):
                neighbours = list(G.neighbors(walk[-1]))
                if not neighbours:
                    break
                walk.append(int(rng.choice(neighbours)))
            walks.append(walk)
    return walks

walks = random_walks(G)
print(f"Generated {len(walks)} random walks of length up to 20")

# Build (target, context) pairs with a window of size W on either side.
W = 3
pairs = []
for walk in walks:
    for i, target in enumerate(walk):
        lo, hi = max(0, i - W), min(len(walk), i + W + 1)
        for j in range(lo, hi):
            if i != j:
                pairs.append((target, walk[j]))
pairs = np.array(pairs, dtype=np.int64)
print(f"{len(pairs):,} (target, context) pairs")

target = torch.tensor(pairs[:, 0], dtype=torch.long)
context = torch.tensor(pairs[:, 1], dtype=torch.long)

# ---------------------------------------------------------------------------
# 4. Skip-gram with negative sampling
# ---------------------------------------------------------------------------
DIM = 2          # 2D so we can plot directly
K_NEG = 5        # negatives per positive

class NodeEmb(nn.Module):
    def __init__(self, n, d):
        super().__init__()
        self.E = nn.Embedding(n, d)
        nn.init.normal_(self.E.weight, std=0.1)

    def forward(self, target, context, neg):
        t = self.E(target)             # (B, d)
        c = self.E(context)            # (B, d)
        n = self.E(neg)                # (B, K, d)
        pos = (t * c).sum(-1)                       # (B,)
        ng  = torch.einsum('bd,bkd->bk', t, n)       # (B, K)
        # Skip-gram with negative sampling: maximise log σ(t·c) + Σ log σ(-t·n)
        loss = -(torch.nn.functional.logsigmoid(pos).mean()
                 + torch.nn.functional.logsigmoid(-ng).mean())
        return loss

model = NodeEmb(N, DIM)
opt = torch.optim.Adam(model.parameters(), lr=0.05)

# ---------------------------------------------------------------------------
# 5. Train and snapshot
# ---------------------------------------------------------------------------
SNAPSHOTS = [0, 5, 25, 100, 300]
N_EPOCHS = SNAPSHOTS[-1] + 1

snaps = {}
losses = []
rng_neg = np.random.default_rng(1)
for epoch in range(N_EPOCHS):
    if epoch in SNAPSHOTS:
        snaps[epoch] = model.E.weight.detach().numpy().copy()
    # fresh negatives every epoch
    neg = torch.tensor(rng_neg.integers(0, N, size=(len(pairs), K_NEG)),
                       dtype=torch.long)
    loss = model(target, context, neg)
    opt.zero_grad()
    loss.backward()
    opt.step()
    losses.append(loss.item())

print(f"final skip-gram loss = {losses[-1]:.3f}")

# ---------------------------------------------------------------------------
# 6. Figure 2 — embeddings evolution + loss curve
# ---------------------------------------------------------------------------
fig2 = plt.figure(figsize=(4 * len(SNAPSHOTS), 5.6))
gs = fig2.add_gridspec(2, len(SNAPSHOTS), height_ratios=[4, 1], hspace=0.55)

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 = fig2.add_subplot(gs[0, col])
    E = snaps[ep]
    for c in range(N_COMMUNITIES):
        mask = community == c
        ax.scatter(E[mask, 0], E[mask, 1],
                   c=COMMUNITY_COLORS[c], s=60,
                   edgecolor='black', linewidth=0.4, alpha=0.85,
                   label=f'community {c}' 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=8, loc='lower left', framealpha=0.85)

ax_loss = fig2.add_subplot(gs[1, :])
ax_loss.plot(losses, c='#264653')
for ep in SNAPSHOTS:
    ax_loss.axvline(ep, c='gray', lw=0.6, alpha=0.6)
ax_loss.set_xlabel("epoch")
ax_loss.set_ylabel("skip-gram loss")
ax_loss.set_title("Training loss")
ax_loss.grid(alpha=0.3)

fig2.suptitle(
    "Node embeddings learn community structure from random walks alone\n"
    f"({N} nodes, {N_COMMUNITIES} planted communities, dim={DIM}, "
    "DeepWalk-style skip-gram with negative sampling)",
    fontsize=13,
)
fig2.savefig("fig2_embeddings_evolution.png", dpi=130, bbox_inches='tight')
print("Saved → fig2_embeddings_evolution.png")

# ---------------------------------------------------------------------------
# 7. Sanity check — k-NN classification accuracy from embeddings alone
# ---------------------------------------------------------------------------
E_final = snaps[SNAPSHOTS[-1]]
# normalise so distances ≡ cosine
En = E_final / (np.linalg.norm(E_final, axis=1, keepdims=True) + 1e-9)
sim = En @ En.T
np.fill_diagonal(sim, -np.inf)
nearest = sim.argmax(axis=1)
acc = (community[nearest] == community).mean()
print(f"1-NN community-prediction accuracy from embeddings: {acc:.1%}")

plt.show()