import warnings

warnings.filterwarnings("ignore", category=FutureWarning)
# harmless Intel-GPU probe message printed by torch on this machine
warnings.filterwarnings("ignore", message=".*Level Zero.*")

import os
import time

import marimo as mo
from gl.nb import reveal
import numpy as np
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
import torch
from gl.datasets import adjusted_homophily, graph, node_data, planetoid
from gl.nb import json_cache
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric.nn.models import CorrectAndSmooth
from torch_geometric.utils import (
    negative_sampling,
    subgraph,
    to_undirected,
)
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score

SEED = 0
rng = np.random.default_rng(SEED)
torch.manual_seed(SEED)
# several open notebooks share the CPU; more threads = contention
torch.set_num_threads(4)
# shared datasets (Cora etc. are cached here) and this folder's disk cache
DATA_DIR = mo.notebook_dir().parents[1] / "data"
CACHE_DIR = mo.notebook_dir() / "cache"
CACHE_DIR.mkdir(exist_ok=True)
plt.rcParams.update(
    {
        "figure.dpi": 100,
        "axes.spines.top": False,
        "axes.spines.right": False,
    }
)

cora = planetoid("Cora")[0]
_und = cora.edge_index[:, cora.edge_index[0] < cora.edge_index[1]]
_perm = torch.randperm(
    _und.size(1), generator=torch.Generator().manual_seed(SEED)
)
_n_test = _und.size(1) // 10
_test_pos, _train_und = _und[:, _perm[:_n_test]], _und[:, _perm[_n_test:]]
torch.manual_seed(SEED)
_test_neg = negative_sampling(
    cora.edge_index, num_nodes=cora.num_nodes, num_neg_samples=_n_test
)
cora_train_graph = nx.Graph()
cora_train_graph.add_nodes_from(range(cora.num_nodes))
cora_train_graph.add_edges_from(_train_und.t().tolist())
_pairs = [
    tuple(p) for p in torch.cat([_test_pos, _test_neg], 1).t().tolist()
]
_labels = np.r_[np.ones(_n_test), np.zeros(_n_test)]
_A_tr = nx.to_scipy_sparse_array(
    cora_train_graph, nodelist=range(cora.num_nodes), format="csr"
).astype(float)
_katz_sum, _walks = 0, _A_tr.copy()  # truncated Katz: walks of length 1..4
for _l in range(1, 5):
    _katz_sum = _katz_sum + (0.05**_l) * _walks
    _walks = (
        _walks @ _A_tr
    )  # matrix power (scipy's ** would be elementwise!)
_heur = {
    "common neighbours": [
        len(set(cora_train_graph[u]) & set(cora_train_graph[v]))
        for u, v in _pairs
    ],
    "Jaccard": [
        s for *_, s in nx.jaccard_coefficient(cora_train_graph, _pairs)
    ],
    "Adamic–Adar": [
        s for *_, s in nx.adamic_adar_index(cora_train_graph, _pairs)
    ],
    "Katz (β=0.05, walks ≤ 4)": [_katz_sum[u, v] for u, v in _pairs],
}
pd.DataFrame(
    [
        {
            "heuristic": k,
            "link-prediction AUC on Cora": roc_auc_score(_labels, v),
        }
        for k, v in _heur.items()
    ]
).round(3)

_rev = reveal(
    r"""
Measured (10% held-out edges vs. as many non-edges): common neighbours, Jaccard and Adamic–Adar all reach **AUC ≈ 0.75**: about a
quarter of Cora's held-out edges have no common neighbour left, and those all score 0. **Katz reaches 0.85**: longer walks rescue those
pairs, matching L02's node2vec (0.85) with no training at all.

""",
    predict="What fraction of Cora's held-out edges have no common neighbour left, and what does that do to the AUC of common neighbours vs. Katz?",
)
_post = mo.md(r"""
## 3 · Labels are signals too: label propagation and Correct & Smooth

**Label propagation:** clamp the known labels, and repeatedly set every other node's class distribution to the average of its
neighbours' (with symmetric normalisation and a restart weight $\alpha$): $Y^{(t+1)}=\alpha\,\tilde A Y^{(t)}+(1-\alpha)Y^{(0)}$.

**Correct & Smooth** (Huang et al. 2021): take *any* base predictor (here an MLP on the features), then (1) **correct**: spread its
*errors* on the training nodes along the edges; (2) **smooth**: spread the corrected predictions, treating training labels as known.
""")
mo.vstack([_rev, _post])

def norm_adj(edge_index, n):
    A = torch.zeros(n, n)
    A[edge_index[0], edge_index[1]] = 1.0
    d = A.sum(1).clamp(min=1).pow(-0.5)
    return d[:, None] * A * d[None, :]

cora_A_norm = norm_adj(cora.edge_index, cora.num_nodes)
_y_onehot = F.one_hot(cora.y, 7).float()
_tr, _te = cora.train_mask, cora.test_mask

def label_propagation(alpha=0.9, iters=50):
    Y0 = torch.zeros(cora.num_nodes, 7)
    Y0[_tr] = _y_onehot[_tr]
    Y = Y0.clone()
    for _ in range(iters):
        Y = alpha * cora_A_norm @ Y + (1 - alpha) * Y0
    return Y

def train_mlp(epochs=200, seed=SEED):
    torch.manual_seed(seed)
    mlp = nn.Sequential(
        nn.Linear(cora.num_features, 64),
        nn.ReLU(),
        nn.Dropout(0.5),
        nn.Linear(64, 7),
    )
    opt = torch.optim.Adam(mlp.parameters(), lr=0.01, weight_decay=5e-4)
    for _ in range(epochs):
        mlp.train()
        opt.zero_grad()
        F.cross_entropy(mlp(cora.x)[_tr], cora.y[_tr]).backward()
        opt.step()
    mlp.eval()
    with torch.no_grad():
        return F.softmax(mlp(cora.x), 1)

def train_gcn(epochs=200, seed=SEED):
    torch.manual_seed(seed)
    c1, c2 = GCNConv(cora.num_features, 64), GCNConv(64, 7)
    opt = torch.optim.Adam(
        list(c1.parameters()) + list(c2.parameters()),
        lr=0.01,
        weight_decay=5e-4,
    )
    f = lambda training: c2(
        F.dropout(F.relu(c1(cora.x, cora.edge_index)), 0.5, training),
        cora.edge_index,
    )
    for _ in range(epochs):
        opt.zero_grad()
        F.cross_entropy(f(True)[_tr], cora.y[_tr]).backward()
        opt.step()
    with torch.no_grad():
        return F.softmax(f(False), 1)

acc = lambda P: (P.argmax(1)[_te] == cora.y[_te]).float().mean().item()
_lp = label_propagation()
_mlp = train_mlp()
_cs = CorrectAndSmooth(
    num_correction_layers=50,
    correction_alpha=0.9,
    num_smoothing_layers=50,
    smoothing_alpha=0.8,
    autoscale=True,
)
_corrected = _cs.correct(_mlp, cora.y[_tr], _tr, cora.edge_index)
_smoothed = _cs.smooth(_corrected, cora.y[_tr], _tr, cora.edge_index)
_gcn = train_gcn()
pd.DataFrame(
    [
        {
            "method": "label propagation (labels + graph, no features)",
            "test acc": acc(_lp),
        },
        {"method": "MLP (features only)", "test acc": acc(_mlp)},
        {"method": "MLP + Correct (C)", "test acc": acc(_corrected)},
        {
            "method": "MLP + Correct & Smooth (C&S)",
            "test acc": acc(_smoothed),
        },
        {"method": "2-layer GCN", "test acc": acc(_gcn)},
    ]
).round(3)

_rev = reveal(
    r"""
Measured (Cora, 140 labels):

* **Label propagation alone: 0.71**, without looking at a single word of the papers, purely labels + homophily.
* **MLP: 0.57** (features only). The *correct* step alone changes nothing here: the MLP fits its 140 training nodes perfectly, so there are
  almost no training errors to propagate. **Smoothing** lifts it to **0.75**.
* **GCN: 0.80.** C&S closes most of the gap with a model that never does message passing during training. On larger benchmarks
  (OGB-Products, Arxiv) the C&S paper even beats GNNs at a fraction of the cost.

""",
    predict="Rank label propagation alone, an MLP, MLP + Correct & Smooth, and a GCN on Cora.",
)
_rev

def cluster_gcn(clusters_per_batch=4, epochs=60, seed=SEED):
    """Partition Cora with Louvain, train a 2-layer GCN on random groups of clusters (edges between groups are dropped)."""
    g = nx.Graph()
    g.add_nodes_from(range(cora.num_nodes))
    g.add_edges_from(cora.edge_index.t().tolist())
    parts = [
        sorted(c) for c in nx.community.louvain_communities(g, seed=seed)
    ]
    torch.manual_seed(seed)
    r = np.random.default_rng(seed)
    c1, c2 = GCNConv(cora.num_features, 64), GCNConv(64, 7)
    opt = torch.optim.Adam(
        list(c1.parameters()) + list(c2.parameters()),
        lr=0.01,
        weight_decay=5e-4,
    )
    kept = []
    for _ in range(epochs):
        order = r.permutation(len(parts))
        for b in range(0, len(order), clusters_per_batch):
            nodes = torch.tensor(
                sorted(
                    sum(
                        (
                            parts[i]
                            for i in order[b : b + clusters_per_batch]
                        ),
                        [],
                    )
                )
            )
            ei, _ = subgraph(
                nodes,
                cora.edge_index,
                relabel_nodes=True,
                num_nodes=cora.num_nodes,
            )
            kept.append(ei.size(1))
            mask = cora.train_mask[nodes]
            if mask.sum() == 0:
                continue
            out = c2(
                F.dropout(F.relu(c1(cora.x[nodes], ei)), 0.5, True), ei
            )
            opt.zero_grad()
            F.cross_entropy(out[mask], cora.y[nodes][mask]).backward()
            opt.step()
    with torch.no_grad():
        pred = c2(
            F.relu(c1(cora.x, cora.edge_index)), cora.edge_index
        ).argmax(1)  # full-graph inference
    n_batches = int(np.ceil(len(parts) / clusters_per_batch))
    return (
        len(parts),
        (pred[cora.test_mask] == cora.y[cora.test_mask])
        .float()
        .mean()
        .item(),
        1 - sum(kept[-n_batches:]) / cora.edge_index.size(1),
    )

def sgc(K=2):
    X = cora.x.clone()
    A_hat = norm_adj(
        torch.cat(
            [cora.edge_index, torch.arange(cora.num_nodes).repeat(2, 1)], 1
        ),
        cora.num_nodes,
    )  # with self-loops
    for _ in range(K):
        X = A_hat @ X  # precomputed once
    clf = LogisticRegression(max_iter=2000, C=1.0).fit(
        X[cora.train_mask].numpy(), cora.y[cora.train_mask].numpy()
    )
    return clf.score(
        X[cora.test_mask].numpy(), cora.y[cora.test_mask].numpy()
    )

_t = time.time()
_n_parts, _cg_acc, _dropped = cluster_gcn()
_t_cg = time.time() - _t
_t = time.time()
_sgc_acc = sgc()
_t_sgc = time.time() - _t
_t = time.time()
_gcn_acc = acc(train_gcn())
_t_gcn = time.time() - _t
pd.DataFrame(
    [
        {
            "method": "full-batch GCN",
            "test acc": _gcn_acc,
            "seconds": _t_gcn,
            "note": "",
        },
        {
            "method": f"Cluster-GCN ({_n_parts} Louvain clusters, 4 per batch)",
            "test acc": _cg_acc,
            "seconds": _t_cg,
            "note": f"{_dropped:.0%} of edges dropped in an epoch",
        },
        {
            "method": "SGC (²X precomputed + logistic regression)",
            "test acc": _sgc_acc,
            "seconds": _t_sgc,
            "note": "no non-linearity",
        },
    ]
).round(3)

_rev = reveal(
    r"""
Measured: **Cluster-GCN 0.796** vs. full-batch **0.798**, although each mini-batch only sees 4 of 102 Louvain clusters and 11% of edges
are never used in an epoch (communities keep most edges inside). **SGC 0.825**: two propagation steps, no non-linearity, a logistic
regression, and it *beats* the GCN here; Cora's labels are smooth over the graph. All three take ≈ 1.5 s on Cora; the difference is
memory and scalability on graphs with $10^8$ nodes.

""",
    predict="Cluster-GCN sees 4 of 102 clusters per batch: how much accuracy does it lose against full batch? Will SGC (no non-linearity) beat the GCN?",
)
_post = mo.md(r"""
---

## 5 · Exercises (the notes' practice problems)

**Exercise 1 (star centrality).** In a star with centre $h$ and $n-1$ leaves, give the betweenness of $h$ for $n=10$ → `ex1_betweenness`,
and the closeness $1/\sum d$ of a **leaf** for $n=10$ → `ex1_leaf_closeness`.

**Exercise 2 (label propagation by hand).** Triangle u–v–w plus an edge w–x; u is labelled 1, x labelled 0; v and w start at 0.5 and are
updated synchronously to the **plain average of their neighbours** until convergence. Final value of w? → `ex2_w`. (Check by coding it.)

**Exercise 3 (SGC).** Why does SGC add self-loops? Think about a bipartite graph with $K=1$, then open the solution.
""")
mo.vstack([_rev, _post])

ex1_betweenness = None
ex1_leaf_closeness = None
ex2_w = None

def _verdict(ok, answered):
    if not answered:
        return mo.callout(mo.md("Not answered yet."), kind="neutral")
    return mo.callout(
        mo.md("Correct!" if ok else "Not yet — try again or open a hint."),
        kind="success" if ok else "warn",
    )

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(
            ex1_betweenness == 36
            and ex1_leaf_closeness is not None
            and abs(ex1_leaf_closeness - 1 / 17) < 1e-6,
            ex1_betweenness is not None or ex1_leaf_closeness is not None,
        ),
        mo.md("**Exercise 2**"),
        _verdict(
            ex2_w is not None and abs(ex2_w - 0.6) < 1e-3,
            ex2_w is not None,
        ),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    r"Every pair of the 9 leaves is connected only through $h$: $\binom92=36$. A leaf is 1 from $h$ and 2 from the other 8 leaves: $1/(1+16)=1/17$."
                ),
                "Solution 2": mo.md(r"""Fixed point: $v=(1+w)/2$, $w=(1+v+0)/3$ → $6w = 3 + w$ → **w = 0.6**, $v=0.8$.
```python
v, w = 0.5, 0.5
for _ in range(100):
    v, w = (1 + w) / 2, (1 + v + 0) / 3
```"""),
                "Solution 3": mo.md(
                    "Without self-loops, $\\tilde A^K X$ only mixes features of nodes at walk distance exactly $K$: for odd $K$ on a bipartite graph a "
                    "node's own features vanish. Self-loops add walks that 'stay', so every node keeps part of its own features and all "
                    "nodes within distance $K$ contribute; they also make the filter more low-pass."
                ),
            }
        ),
    ]
)