class MyGCNConv(MessagePassing):
    """GCN: message = W h_u / sqrt(d_u d_v), aggregation = sum, self-loops included."""

    def __init__(self, d_in, d_out):
        super().__init__(aggr="add")
        self.lin = nn.Linear(d_in, d_out, bias=False)
        self.bias = nn.Parameter(torch.zeros(d_out))

    def forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        src, dst = edge_index
        deg = degree(dst, x.size(0))
        norm = deg[src].pow(-0.5) * deg[dst].pow(-0.5)
        return (
            self.propagate(edge_index, x=self.lin(x), norm=norm)
            + self.bias
        )

    def message(self, x_j, norm):
        return norm.unsqueeze(-1) * x_j

class MySAGEConv(MessagePassing):
    """GraphSAGE-mean: h_v' = W_l · mean_{u∈N(v)} h_u + W_r · h_v   (= W [h_v || h_N(v)])."""

    def __init__(self, d_in, d_out):
        super().__init__(aggr="mean")
        self.lin_l = nn.Linear(d_in, d_out)  # neighbour summary
        self.lin_r = nn.Linear(d_in, d_out, bias=False)  # self

    def forward(self, x, edge_index):
        return self.lin_l(self.propagate(edge_index, x=x)) + self.lin_r(x)

    def message(self, x_j):
        return x_j

class MyGATConv(MessagePassing):
    """Single-head GAT: e_vu = LeakyReLU(a_src·Wh_u + a_dst·Wh_v), α = softmax over N(v)∪{v}, h_v' = Σ α_vu W h_u."""

    def __init__(self, d_in, d_out):
        super().__init__(aggr="add")
        self.lin = nn.Linear(d_in, d_out, bias=False)
        self.att_src = nn.Parameter(torch.randn(d_out) * 0.1)
        self.att_dst = nn.Parameter(torch.randn(d_out) * 0.1)
        self.bias = nn.Parameter(torch.zeros(d_out))

    def forward(self, x, edge_index, return_attention=False):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        h = self.lin(x)
        a_src = (h * self.att_src).sum(
            -1, keepdim=True
        )  # [N, 1] per-node halves of the score
        a_dst = (h * self.att_dst).sum(-1, keepdim=True)
        out = (
            self.propagate(edge_index, h=h, a_src=a_src, a_dst=a_dst)
            + self.bias
        )
        return (
            (out, (edge_index, self._alpha)) if return_attention else out
        )

    def message(self, h_j, a_src_j, a_dst_i, index):
        e = F.leaky_relu(
            a_src_j + a_dst_i, 0.2
        )  # one raw score per edge u→v
        alpha = softmax(
            e, index
        )  # normalise over each target's incoming edges
        self._alpha = alpha.squeeze(-1)
        return alpha * h_j

cora = planetoid("Cora")[0]

def compare_with_pyg():
    torch.manual_seed(SEED)
    rows = []
    pairs = [
        (
            "GCN",
            MyGCNConv(cora.num_features, 16),
            GCNConv(cora.num_features, 16),
        ),
        (
            "GraphSAGE-mean",
            MySAGEConv(cora.num_features, 16),
            SAGEConv(cora.num_features, 16),
        ),
        (
            "GAT (1 head)",
            MyGATConv(cora.num_features, 16),
            GATConv(cora.num_features, 16, heads=1),
        ),
    ]
    for name, mine, ref in pairs:
        with torch.no_grad():
            for n_ref, p_ref in ref.named_parameters():
                p_ref.copy_(
                    dict(mine.named_parameters())[n_ref].view_as(p_ref)
                )  # identical weights
            gap = (
                (
                    mine(cora.x, cora.edge_index)
                    - ref(cora.x, cora.edge_index)
                )
                .abs()
                .max()
                .item()
            )
        rows.append(
            {
                "layer": name,
                "parameters": sum(p.numel() for p in mine.parameters()),
                "max |mine − PyG|": gap,
            }
        )
    return pd.DataFrame(rows)

compare_with_pyg()

class FlexGNN(nn.Module):
    """pre-MLP → L × [conv → BatchNorm → ReLU → Dropout (+ skip)] → post-MLP head."""

    def __init__(
        self,
        d_in,
        d_hid,
        n_cls,
        n_layers=2,
        conv="GCN",
        skip=False,
        bn=True,
        dropout=0.5,
        heads=4,
    ):
        super().__init__()
        make = {
            "GCN": lambda: GCNConv(d_hid, d_hid),
            "SAGE": lambda: SAGEConv(d_hid, d_hid),
            "GAT": lambda: GATConv(d_hid, d_hid // heads, heads=heads),
        }[conv]
        self.pre = nn.Linear(d_in, d_hid)
        self.convs = nn.ModuleList(make() for _ in range(n_layers))
        self.norms = nn.ModuleList(
            nn.BatchNorm1d(d_hid) if bn else nn.Identity()
            for _ in range(n_layers)
        )
        self.post = nn.Sequential(
            nn.Linear(d_hid, d_hid), nn.ReLU(), nn.Linear(d_hid, n_cls)
        )
        self.skip, self.p = skip, dropout

    def forward(self, x, edge_index, return_hidden=False):
        h = F.relu(self.pre(x))
        for conv, norm in zip(self.convs, self.norms):
            out = F.dropout(
                F.relu(norm(conv(h, edge_index))), self.p, self.training
            )
            h = h + out if self.skip else out  # F(x) + x  vs  F(x)
        return (self.post(h), h) if return_hidden else self.post(h)

def fit(model, data, epochs=200, lr=0.01, wd=5e-4):
    """Full-batch training; returns test accuracy at the best-validation epoch and the model."""
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=wd)
    best_val, best_test = -1, 0
    for _ in range(epochs):
        model.train()
        opt.zero_grad()
        F.cross_entropy(
            model(data.x, data.edge_index)[data.train_mask],
            data.y[data.train_mask],
        ).backward()
        opt.step()
        model.eval()
        with torch.no_grad():
            pred = model(data.x, data.edge_index).argmax(1)
        val = (
            (pred[data.val_mask] == data.y[data.val_mask])
            .float()
            .mean()
            .item()
        )
        if val >= best_val:
            best_val, best_test = (
                val,
                (pred[data.test_mask] == data.y[data.test_mask])
                .float()
                .mean()
                .item(),
            )
    return best_test, model

layer_run = mo.ui.run_button(label="Train GCN / SAGE / GAT (≈ 20 s)")
layer_run

_rev = reveal(
    r"""
Measured (2 seeds, 64 hidden, pre/post MLPs + BatchNorm): **GCN ≈ 0.77, GAT ≈ 0.76, SAGE ≈ 0.74**, each ≈ 3 s.
On a small, homophilous citation graph the three are within a few points — the *layer* matters less than people expect.
(The plain 2-layer GCN of L03 scored ≈ 0.79: with only 140 labels the extra pre/post layers and BatchNorm add parameters
without adding much. Design choices interact with data size, which is the point of L19's design-space study.)

""",
    predict="GCN, GAT and GraphSAGE with identical pre/post MLPs on Cora: will they be 1, 5 or 15 points apart?",
)
_post = mo.md(r"""
## 5 · Stacking layers: receptive fields and over-smoothing

The **receptive field** of a $K$-layer GNN is the $K$-hop neighbourhood. On Cora, measured for 200 random papers:
""")
mo.vstack([_rev, _post])

def depth_experiment(depths=(1, 2, 4, 8, 16, 32), seed=SEED):
    key = hashlib.sha1(
        json.dumps({"depths": list(depths), "seed": seed, "v": 1}).encode()
    ).hexdigest()[:10]
    path = CACHE_DIR / f"L04_depth_{key}.json"
    if path.exists():
        return pd.read_json(path), True
    sample = torch.from_numpy(
        np.random.default_rng(seed).choice(
            cora.num_nodes, 500, replace=False
        )
    )
    rows = []
    for L in depths:
        for skip in (False, True):
            torch.manual_seed(seed)
            t = time.time()
            acc, model = fit(
                FlexGNN(
                    cora.num_features,
                    64,
                    7,
                    n_layers=L,
                    conv="GCN",
                    skip=skip,
                ),
                cora,
            )
            model.eval()
            with torch.no_grad():
                h = F.normalize(
                    model(cora.x, cora.edge_index, return_hidden=True)[1][
                        sample
                    ],
                    dim=1,
                )
            rows.append(
                {
                    "layers": L,
                    "skip": skip,
                    "test acc": acc,
                    "mean cosine of hidden vectors": (h @ h.T)
                    .mean()
                    .item(),
                    "seconds": time.time() - t,
                }
            )
    df = pd.DataFrame(rows)
    df.to_json(path)
    return df, False

depth_run = mo.ui.run_button(
    label="Run depth experiment (≈ 3 min first time, then cached)"
)
depth_run

_rev = reveal(
    r"""
Measured (one seed, so ±2–3 points of noise per cell):

* **Plain GCN** peaks at **4 layers (0.80)**, then collapses: 0.70 at 8, 0.50 at 16, **0.32 at 32**. The hidden vectors of
  random nodes become near-parallel (mean cosine 0.76 at 32 layers vs. ≈ 0.4 for shallow models). That is over-smoothing,
  together with the general difficulty of training very deep stacks.
* **With skip connections** accuracy stays at **0.70–0.75 all the way to 32 layers**. Skips keep each node's own signal alive.
* Skips are *not* free: at 1–2 layers they score lower here (0.70 vs. 0.77), because the residual path keeps a lot of
  the pre-MLP's feature-only signal. Depth and skips are design choices to tune, not defaults.

""",
    predict="At 32 layers, what happens to a plain GCN's accuracy, and what do skip connections change? Guess both numbers before you open this.",
)
_post = mo.md(r"""
---

## 6 · Exercises

**Exercise 1 (by hand).** Raw GAT scores for A's neighbours B, C, D are $(2, 2, 0)$. What attention weight does D get?
→ `ex1_alpha_D` (3 decimals).

**Exercise 2 (parameters).** How many trainable parameters does `SAGEConv(64, 64)` have (PyG: `lin_l` with bias, `lin_r` without)? → `ex2_params`.

**Exercise 3 (code).** Implement `MyMaxSAGE`, GraphSAGE with **max** aggregation, using `MessagePassing`. It must
match `SAGEConv(d_in, d_out, aggr="max")` when the weights are copied (parameter names `lin_l`, `lin_r`).
""")
mo.vstack([_rev, _post])

ex1_alpha_D = None
ex2_params = None

class MyMaxSAGE(MessagePassing):
    def __init__(self, d_in, d_out):
        super().__init__(aggr="mean")  # <- change things here
        self.lin_l = nn.Linear(d_in, d_out)
        self.lin_r = nn.Linear(d_in, d_out, bias=False)

    def forward(self, x, edge_index):
        return None  # <- and here

    def message(self, x_j):
        return x_j

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_max_sage():
    try:
        torch.manual_seed(SEED)
        mine, ref = (
            MyMaxSAGE(cora.num_features, 8),
            SAGEConv(cora.num_features, 8, aggr="max"),
        )
        with torch.no_grad():
            for n, p in ref.named_parameters():
                p.copy_(dict(mine.named_parameters())[n])
            out = mine(cora.x, cora.edge_index)
            if out is None:
                return _verdict(False, False)
            return _verdict(
                torch.allclose(
                    out, ref(cora.x, cora.edge_index), atol=1e-5
                ),
                True,
            )
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(
            ex1_alpha_D is not None
            and abs(ex1_alpha_D - 1 / (2 * np.e**2 + 1)) < 2e-3,
            ex1_alpha_D is not None,
        ),
        mo.md("**Exercise 2**"),
        _verdict(ex2_params == 64 * 64 * 2 + 64, ex2_params is not None),
        mo.md("**Exercise 3**"),
        _check_max_sage(),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    r"$\alpha_D=e^0/(e^2+e^2+e^0)=1/(2e^2+1)=1/15.78\approx 0.063$; B and C get 0.468 each."
                ),
                "Solution 2": mo.md(
                    r"`lin_l`: $64\times64+64$, `lin_r`: $64\times64$ → **8,256**. Twice a GCN layer's $64\times64+64=4{,}160$."
                ),
                "Solution 3": mo.md("""```python
class MyMaxSAGE(MessagePassing):
    def __init__(self, d_in, d_out):
        super().__init__(aggr="max")
        self.lin_l = nn.Linear(d_in, d_out)
        self.lin_r = nn.Linear(d_in, d_out, bias=False)
    def forward(self, x, edge_index):
        return self.lin_l(self.propagate(edge_index, x=x)) + self.lin_r(x)
    def message(self, x_j):
        return x_j
```
Only `aggr` changes: the framework separates *what is sent* (`message`) from *how it is combined* (`aggr`)."""),
            }
        ),
    ]
)

_x = torch.randn(cora.num_nodes, 64)
pd.DataFrame(
    [
        {
            "layer": "GATConv(64, 64, heads=4)",
            "output shape": tuple(
                GATConv(64, 64, heads=4)(_x, cora.edge_index).shape
            ),
        },
        {
            "layer": "GATConv(64, 16, heads=4)  ← width per head = 64/4",
            "output shape": tuple(
                GATConv(64, 16, heads=4)(_x, cora.edge_index).shape
            ),
        },
        {
            "layer": "GATConv(64, 64, heads=4, concat=False)  ← average heads",
            "output shape": tuple(
                GATConv(64, 64, heads=4, concat=False)(
                    _x, cora.edge_index
                ).shape
            ),
        },
    ]
)