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 planetoid, temporal
from gl.nb import json_cache
import torch.nn as nn
import torch.nn.functional as F
import torch_geometric.transforms as T
from torch_geometric.data import Batch, Data
from torch_geometric.nn import (
    GCNConv,
    SAGEConv,
    global_add_pool,
    global_max_pool,
    global_mean_pool,
)
from torch_geometric.utils import from_networkx, negative_sampling
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    precision_score,
    recall_score,
    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,
    }
)

def cycles_graph(n_tri=60, n_sq=45, seed=SEED, relabel=False):
    """Disjoint triangles (label 1) and squares (label 0); half the nodes train, half test.
    relabel=True shuffles node IDs: the same kind of graph, but every ID now names a different node."""
    g = nx.disjoint_union_all(
        [nx.cycle_graph(3)] * n_tri + [nx.cycle_graph(4)] * n_sq
    )
    y = np.array([1] * (3 * n_tri) + [0] * (4 * n_sq))
    if relabel:
        mapping = dict(
            zip(
                g.nodes,
                np.random.default_rng(seed + 1).permutation(
                    g.number_of_nodes()
                ),
            )
        )
        g = nx.relabel_nodes(g, mapping)
        y = y[np.argsort(list(mapping.values()))]
    g_sorted = nx.Graph()
    g_sorted.add_nodes_from(sorted(g.nodes))  # tensor row i = node ID i
    g_sorted.add_edges_from(g.edges)
    g = g_sorted
    data = from_networkx(g)
    data.y = torch.tensor(y)
    perm = torch.randperm(
        data.num_nodes, generator=torch.Generator().manual_seed(seed)
    )
    data.train_mask = torch.zeros(data.num_nodes, dtype=torch.bool)
    data.train_mask[perm[: data.num_nodes // 2]] = True
    data.test_mask = ~data.train_mask
    return g, data

def node_features(g):
    A = torch.tensor(
        nx.to_numpy_array(g, nodelist=sorted(g.nodes)), dtype=torch.float
    )
    return {
        "constant (1)": torch.ones(len(A), 1),
        "one-hot node ID": torch.eye(len(A)),
        "degree": A.sum(1, keepdim=True),
        # closed walks of length 3 and 4 = diag(A^3), diag(A^4): "am I on a triangle / square?"
        "cycle counts diag(A³), diag(A⁴)": torch.stack(
            [torch.diag(A @ A @ A), torch.diag(A @ A @ A @ A)], 1
        ),
    }

class GCN2(nn.Module):
    def __init__(self, d_in, d_hid=32, n_out=2):
        super().__init__()
        self.c1, self.c2, self.head = (
            GCNConv(d_in, d_hid),
            GCNConv(d_hid, d_hid),
            nn.Linear(d_hid, n_out),
        )

    def forward(self, x, edge_index):
        return self.head(
            F.relu(self.c2(F.relu(self.c1(x, edge_index)), edge_index))
        )

def train_gcn2(x, data, epochs=150, seed=SEED):
    torch.manual_seed(seed)
    model = GCN2(x.size(1))
    opt = torch.optim.Adam(model.parameters(), lr=0.01)
    for _ in range(epochs):
        opt.zero_grad()
        F.cross_entropy(
            model(x, data.edge_index)[data.train_mask],
            data.y[data.train_mask],
        ).backward()
        opt.step()
    return model

def accuracy(model, x, data, mask):
    with torch.no_grad():
        return (
            (model(x, data.edge_index).argmax(1)[mask] == data.y[mask])
            .float()
            .mean()
            .item()
        )

cyc_g, cyc_data = cycles_graph()
new_g, new_data = cycles_graph(
    relabel=True
)  # a "new graph" of the same kind
_feats, _new_feats = node_features(cyc_g), node_features(new_g)
_rows = []
for _name, _x in _feats.items():
    _m = train_gcn2(_x, cyc_data)
    _rows.append(
        {
            "node features": _name,
            "dims": _x.size(1),
            "train acc": accuracy(_m, _x, cyc_data, cyc_data.train_mask),
            "test acc (same graph)": accuracy(
                _m, _x, cyc_data, cyc_data.test_mask
            ),
            "acc on a NEW graph": accuracy(
                _m,
                _new_feats[_name],
                new_data,
                torch.ones(new_data.num_nodes, dtype=torch.bool),
            ),
        }
    )
pd.DataFrame(_rows).round(3)

def path_graph_data(colour, length=20, virtual=False):
    """Path 0-1-...-(length-1); node 0 has a one-hot colour, all nodes a constant channel; optional virtual node."""
    x = torch.zeros(length, 3)
    x[0, colour] = 1.0
    x[:, 2] = 1.0
    src, dst = list(range(length - 1)), list(range(1, length))
    is_real = torch.ones(length, dtype=torch.bool)
    if virtual:  # one extra node, linked to every real node
        x = torch.cat([x, torch.zeros(1, 3)])
        src += [length] * length
        dst += list(range(length))
        is_real = torch.cat([is_real, torch.zeros(1, dtype=torch.bool)])
    edge_index = torch.tensor([src + dst, dst + src])
    return Data(
        x=x,
        edge_index=edge_index,
        y=torch.full((x.size(0),), colour),
        dist=torch.arange(x.size(0)),
        is_real=is_real,
    )

class StackGCN(nn.Module):
    def __init__(self, d_in, n_layers, d_hid=32, n_out=2):
        super().__init__()
        dims = [d_in] + [d_hid] * n_layers
        self.convs = nn.ModuleList(
            GCNConv(a, b) for a, b in zip(dims, dims[1:])
        )
        self.head = nn.Linear(d_hid, n_out)

    def forward(self, x, edge_index):
        for conv in self.convs:
            x = F.relu(conv(x, edge_index))
        return self.head(x)

def path_experiment(n_layers, virtual, epochs=200, seed=SEED):
    colours = np.random.default_rng(seed).integers(0, 2, 200)
    graphs = [path_graph_data(int(c), virtual=virtual) for c in colours]
    train, test = (
        Batch.from_data_list(graphs[:150]),
        Batch.from_data_list(graphs[150:]),
    )  # inductive: separate graphs
    torch.manual_seed(seed)
    model = StackGCN(3, n_layers)
    opt = torch.optim.Adam(model.parameters(), lr=0.01)
    for _ in range(epochs):
        opt.zero_grad()
        out = model(train.x, train.edge_index)
        F.cross_entropy(
            out[train.is_real], train.y[train.is_real]
        ).backward()  # loss on real nodes only
        opt.step()
    with torch.no_grad():
        correct = (model(test.x, test.edge_index).argmax(1) == test.y)[
            test.is_real
        ]
    dist = test.dist[test.is_real]
    return np.array(
        [correct[dist == k].float().mean().item() for k in range(20)]
    )

vn_layers = mo.ui.slider(
    1, 4, value=2, label="GCN layers", show_value=True
)
vn_layers

class SAGE2(nn.Module):
    def __init__(self, d_in, d_hid, n_out):
        super().__init__()
        self.c1, self.c2 = SAGEConv(d_in, d_hid), SAGEConv(d_hid, n_out)

    def forward(self, x, edge_index):
        return self.c2(
            F.dropout(F.relu(self.c1(x, edge_index)), 0.5, self.training),
            edge_index,
        )

def train_sampled(fanout, epochs=60, batch_size=32, seed=SEED):
    """Mini-batch training on sampled 2-hop subgraphs; full-graph inference at the end."""
    torch.manual_seed(seed)
    r = np.random.default_rng(seed)
    model = SAGE2(cora.num_features, 64, 7)
    opt = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
    train_idx = cora.train_mask.nonzero().squeeze().numpy()
    sizes, t0 = [], time.time()
    for _ in range(epochs):
        model.train()
        order = r.permutation(train_idx)  # new mini-batches every epoch
        for b in range(0, len(order), batch_size):
            seeds = order[b : b + batch_size]
            n_id, ei = sample_subgraph(
                seeds, [fanout, fanout], cora_csr, r
            )
            sizes.append(len(n_id))
            opt.zero_grad()
            out = model(cora.x[n_id], ei)[
                : len(seeds)
            ]  # predictions for the seeds only
            F.cross_entropy(
                out, cora.y[torch.from_numpy(seeds)]
            ).backward()
            opt.step()
    model.eval()
    with torch.no_grad():
        pred = model(cora.x, cora.edge_index).argmax(
            1
        )  # full-graph inference
    return {
        "fanout per hop": fanout,
        "mean nodes per mini-batch": np.mean(sizes),
        "test acc": (pred[cora.test_mask] == cora.y[cora.test_mask])
        .float()
        .mean()
        .item(),
        "train time (s)": time.time() - t0,
    }

sampling_run = mo.ui.run_button(
    label="Compare fanouts 1, 2, 5, 10, 25 (≈ 5 s)"
)
sampling_run

_rev = reveal(
    r"""
Measured (60 epochs, 2-layer GraphSAGE, batches of 32 labelled papers): with **one** neighbour per hop each mini-batch
touches ≈ 73 nodes and still reaches **0.78** test accuracy; fanout 5 touches ≈ 230 nodes and reaches **0.80**; fanout 25
touches ≈ 470 nodes for no further gain. Most of the accuracy comes from a few sampled neighbours, which is why PinSage-style
sampling scales to billions of nodes. (Single seed: differences under ~1–2 points are noise.)

""",
    predict="Sampling 1, 5 or 25 neighbours per hop: how many points of test accuracy does a fanout of 1 cost?",
)
_post = mo.md(r"""
## 4 · Prediction heads

* **Node:** $\hat{\mathbf y}_v = W^{(H)}\mathbf h_v^{(L)}$.
* **Edge:** $\text{Linear}(\mathbf h_u\|\mathbf h_v)$ (2d → k), or a dot product $\mathbf h_u^\top\mathbf h_v$ (1-way only), or $k$ bilinear forms $\mathbf h_u^\top W^{(i)}\mathbf h_v$.
* **Graph:** a global pooling over node embeddings — mean, max or sum. Simple, but lossy. The lecture's puzzle:
""")
mo.vstack([_rev, _post])

def train_link_predictor(message_edges, epochs=150, seed=SEED):
    """GCN encoder + dot-product decoder; positives = train supervision edges, fresh negatives every epoch."""
    torch.manual_seed(seed)
    enc = GCN2(cora.num_features, 64, 64)
    opt = torch.optim.Adam(enc.parameters(), lr=0.01)
    pos = lp_train.edge_label_index
    for _ in range(epochs):
        opt.zero_grad()
        z = enc(lp_train.x, message_edges)
        neg = negative_sampling(
            lp_train.edge_index,
            num_nodes=cora.num_nodes,
            num_neg_samples=pos.size(1),
        )
        logits = torch.cat(
            [
                (z[pos[0]] * z[pos[1]]).sum(1),
                (z[neg[0]] * z[neg[1]]).sum(1),
            ]
        )
        F.binary_cross_entropy_with_logits(
            logits,
            torch.cat([torch.ones(pos.size(1)), torch.zeros(neg.size(1))]),
        ).backward()
        opt.step()
    return enc

def test_auc(enc, message_edges):
    with torch.no_grad():
        z = enc(lp_test.x, message_edges)
        e = lp_test.edge_label_index
        return roc_auc_score(
            lp_test.edge_label.numpy(), (z[e[0]] * z[e[1]]).sum(1).numpy()
        )

_clean = train_link_predictor(lp_train.edge_index)
_leaky = train_link_predictor(
    cora.edge_index
)  # BUG: message passing over ALL edges, test edges included
mo.callout(
    mo.md(f"""**Correct split:** test AUC **{test_auc(_clean, lp_test.edge_index):.3f}**
(trained on train message edges, tested with train + val edges as messages).

**Leakage bug:** test AUC **{test_auc(_leaky, cora.edge_index):.3f}** (the full `edge_index`, including the very test edges,
was used for message passing). The model "predicts" edges it can see in its own neighbourhood.
**Fix:** always use `train_data.edge_index` / `val_data.edge_index` / `test_data.edge_index` from `RandomLinkSplit`."""),
    kind="warn",
)

ex1_auc = None

def my_auc(pos, neg):
    return 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",
    )

def _check_auc():
    try:
        r = np.random.default_rng(1)
        pos, neg = (
            np.round(r.random(30), 1),
            np.round(r.random(40), 1),
        )  # rounding creates ties
        got = my_auc(pos, neg)
        if got is None:
            return _verdict(False, False)
        return _verdict(
            abs(
                got
                - roc_auc_score(
                    np.r_[np.ones(30), np.zeros(40)], np.r_[pos, neg]
                )
            )
            < 1e-9,
            True,
        )
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(
            ex1_auc is not None and abs(ex1_auc - 4 / 6) < 2e-3,
            ex1_auc is not None,
        ),
        mo.md("**Exercise 2**"),
        _check_auc(),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    "Pairs: 0.8 beats both, 0.4 beats 0.2 only, 0.3 beats 0.2 only → 4 of 6 → **0.667**."
                ),
                "Solution 2": mo.md("""```python
def my_auc(pos, neg):
    pos, neg = np.asarray(pos)[:, None], np.asarray(neg)[None, :]
    return ((pos > neg).sum() + 0.5 * (pos == neg).sum()) / (pos.size * neg.size)
```"""),
                "Solution 3": mo.md(
                    "**Inductive**: each graph is one example with one label, so train / val / test must be different graphs. "
                    "A transductive split would need one graph to be partly train and partly test, which has no meaning for a graph-level label."
                ),
            }
        ),
    ]
)