try:
    dblp = DBLP(str(DATA_DIR / "DBLP"))[0]
except Exception as _e:  # no network, or a blocked download host
    dblp = None
    _err = _e
mo.stop(
    dblp is None,
    mo.callout(
        mo.md(
            f"Could not download DBLP (`{type(_err).__name__ if dblp is None else ''}`). PyG fetches it from "
            "Dropbox on first use; run this notebook once with internet access, or copy a `data/DBLP/` folder "
            "from another machine."
        ),
        kind="warn",
    ),
)
dblp["conference"].x = torch.eye(
    dblp["conference"].num_nodes
)  # featureless type: one-hot IDs (only 20 nodes)
mo.vstack(
    [
        pd.DataFrame(
            [
                {
                    "node type": t,
                    "nodes": dblp[t].num_nodes,
                    "feature dim": dblp[t].x.size(1),
                }
                for t in dblp.node_types
            ]
        ),
        pd.DataFrame(
            [
                {
                    "relation (src, edge, dst)": str(r),
                    "edges": dblp[r].edge_index.size(1),
                }
                for r in dblp.edge_types
            ]
        ),
        mo.md(
            f"Labelled authors: train {int(dblp['author'].train_mask.sum())}, val {int(dblp['author'].val_mask.sum())}, "
            f"test {int(dblp['author'].test_mask.sum())}; classes: {dblp['author'].y.bincount().tolist()}"
        ),
    ]
)

class RGCNLayer(nn.Module):
    """Heterogeneous R-GCN layer on a HeteroData dict. Optional basis decomposition over the relation weights."""

    def __init__(self, relations, node_types, d_in, d_out, num_bases=None):
        super().__init__()
        self.relations = relations
        self.key = {r: "__".join(r) for r in relations}
        if num_bases is None:
            self.W = nn.ParameterDict(
                {
                    self.key[r]: nn.Parameter(
                        torch.randn(d_in, d_out) * d_in**-0.5
                    )
                    for r in relations
                }
            )
        else:
            self.V = nn.Parameter(
                torch.randn(num_bases, d_in, d_out) * d_in**-0.5
            )  # shared bases
            self.a = nn.ParameterDict(
                {
                    self.key[r]: nn.Parameter(
                        torch.randn(num_bases) / num_bases
                    )
                    for r in relations
                }
            )
        self.W0 = nn.ModuleDict(
            {t: nn.Linear(d_in, d_out) for t in node_types}
        )  # self-loop per node type
        self.num_bases = num_bases

    def weight(self, r):
        k = self.key[r]
        return (
            self.W[k]
            if self.num_bases is None
            else torch.einsum("b,bio->io", self.a[k], self.V)
        )

    def forward(self, x_dict, edge_index_dict):
        out = {t: self.W0[t](x) for t, x in x_dict.items()}
        for r in self.relations:
            src_t, _, dst_t = r
            src, dst = edge_index_dict[r]
            msg = x_dict[src_t][src] @ self.weight(
                r
            )  # W_r h_u for every edge
            agg = torch.zeros(
                x_dict[dst_t].size(0), msg.size(1)
            ).index_add_(0, dst, msg)
            deg = (
                torch.bincount(dst, minlength=x_dict[dst_t].size(0))
                .clamp(min=1)
                .unsqueeze(1)
            )
            out[dst_t] = out[dst_t] + agg / deg  # mean per relation
        return out

class HeteroNet(nn.Module):
    """Per-type input projection → 2 message-passing layers (R-GCN, HeteroConv-SAGE or HGT) → author classifier."""

    def __init__(self, kind, d=64, num_bases=None):
        super().__init__()
        self.kind = kind
        self.proj = nn.ModuleDict(
            {t: Linear(-1, d) for t in dblp.node_types}
        )
        if kind == "R-GCN":
            self.layers = nn.ModuleList(
                RGCNLayer(
                    dblp.edge_types, dblp.node_types, d, d, num_bases
                )
                for _ in range(2)
            )
        elif kind == "HeteroConv (SAGE per relation)":
            self.layers = nn.ModuleList(
                HeteroConv(
                    {r: SAGEConv((-1, -1), d) for r in dblp.edge_types},
                    aggr="sum",
                )
                for _ in range(2)
            )
        elif kind == "HGT":
            self.layers = nn.ModuleList(
                HGTConv(d, d, dblp.metadata(), heads=4) for _ in range(2)
            )
        self.out = nn.Linear(d, 4)

    def forward(self, data):
        x = {t: F.relu(self.proj[t](data[t].x)) for t in data.node_types}
        for layer in self.layers:
            x = {
                t: F.dropout(F.relu(h), 0.3, self.training)
                for t, h in layer(x, data.edge_index_dict).items()
            }
        return self.out(x["author"])

def n_params(model):
    return sum(p.numel() for p in model.parameters())

with torch.no_grad():
    _m = HeteroNet("R-GCN")
    _m(dblp)  # lazy Linear(-1, d) layers need one forward pass
    _mb = HeteroNet("R-GCN", num_bases=2)
    _mb(dblp)
mo.md(
    f"R-GCN (2 layers, 6 relations, d=64): **{n_params(_m):,}** parameters (incl. input projections); "
    f"with **2 bases**: **{n_params(_mb):,}**. Message-passing part only: "
    f"{sum(p.numel() for p in _m.layers.parameters()):,} vs {sum(p.numel() for p in _mb.layers.parameters()):,}."
)

class SharedRGCNLayer(RGCNLayer):
    """Relation-blind variant: every relation uses the same weight matrix."""

    def __init__(self, relations, node_types, d_in, d_out):
        super().__init__(relations, node_types, d_in, d_out)
        self.W_shared = nn.Parameter(torch.randn(d_in, d_out) * d_in**-0.5)
        del self.W

    def weight(self, r):
        return self.W_shared

class MLPNet(nn.Module):
    def __init__(self, d=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dblp["author"].x.size(1), d),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(d, 4),
        )

    def forward(self, data):
        return self.net(data["author"].x)

def make_model(name):
    if name == "MLP (author features only)":
        return MLPNet()
    if name == "relation-blind GNN":
        m = HeteroNet("R-GCN")
        m.layers = nn.ModuleList(
            SharedRGCNLayer(dblp.edge_types, dblp.node_types, 64, 64)
            for _ in range(2)
        )
        return m
    if name == "R-GCN + 2 bases":
        return HeteroNet("R-GCN", num_bases=2)
    return HeteroNet(name)

def fit_dblp(name, epochs=100, seed=SEED):
    torch.manual_seed(seed)
    model = make_model(name)
    with torch.no_grad():
        model(dblp)  # initialise lazy layers
    opt = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=1e-4)
    a = dblp["author"]
    best_val, best_test, t0 = 0, 0, time.time()
    for _ in range(epochs):
        model.train()
        opt.zero_grad()
        F.cross_entropy(
            model(dblp)[a.train_mask], a.y[a.train_mask]
        ).backward()
        opt.step()
        model.eval()
        with torch.no_grad():
            pred = model(dblp).argmax(1)
        val = (pred[a.val_mask] == a.y[a.val_mask]).float().mean().item()
        if val >= best_val:
            best_val, best_test = (
                val,
                (pred[a.test_mask] == a.y[a.test_mask])
                .float()
                .mean()
                .item(),
            )
    return {
        "model": name,
        "test acc (best val)": best_test,
        "parameters": n_params(model),
        "seconds": time.time() - t0,
    }

dblp_models = [
    "MLP (author features only)",
    "relation-blind GNN",
    "R-GCN",
    "R-GCN + 2 bases",
    "HeteroConv (SAGE per relation)",
    "HGT",
]
dblp_run = mo.ui.run_button(label="Train the six models (≈ 3 min)")
dblp_run

_rev = reveal(
    r"""
Measured (one seed, 100 epochs):

| model | test acc |
|---|---|
| MLP (author features only) | 0.78 |
| relation-blind GNN | **0.87** |
| R-GCN | 0.86 |
| R-GCN + 2 bases | 0.85 |
| HeteroConv (SAGE per relation) | 0.86 |
| HGT | 0.79 (slowest: 73 s; under-trained at 100 epochs) |

* Using the graph is worth **+8–9 points** over the author features alone.
* The **relation-blind** model is as good as the relation-aware ones here. In DBLP every pair of node types is linked by exactly one
  relation, and the per-type projections and self-loops already tell the model *which* type a message came from. Per-relation
  weights pay off when the **same** pair of types is linked by several relations with different meanings (a knowledge graph's
  "treats" vs "causes", or follows vs. retweets vs. replies between users).
* Differences of 1–2 points are within seed noise; the conclusion "heterogeneous ≫ homogeneous" is not automatic.

""",
    predict="On DBLP, will the relation-aware models (R-GCN, HeteroConv, HGT) beat a relation-blind GNN? How many points does the graph add over author features alone?",
)
_post = mo.md(r"""
## 5 · Ranking metrics for link prediction (Hits@k, MRR)

For a test edge $(h, r, t)$, score the true tail against candidate tails; the **rank** of the true tail gives Hits@k
$=\mathbb 1[\text{rank}\le k]$ and reciprocal rank $1/\text{rank}$, averaged over test edges (MRR). Lecture example: true
$(E,r_3,D)$ scores 0.7; candidates $B$: 0.9, $F$: 0.4 → rank 2 → Hits@1 = 0, Hits@2 = 1, RR = 0.5.
""")
mo.vstack([_rev, _post])

def rank_metrics(true_scores, candidate_scores, ks=(1, 3, 10)):
    """true_scores: [n]; candidate_scores: [n, m] for m negatives each. Rank = 1 + #negatives scoring higher (ties count half)."""
    t = np.asarray(true_scores)[:, None]
    c = np.asarray(candidate_scores)
    ranks = 1 + (c > t).sum(1) + 0.5 * (c == t).sum(1)
    return {
        "mean rank": ranks.mean(),
        "MRR": (1 / ranks).mean(),
        **{f"Hits@{k}": (ranks <= k).mean() for k in ks},
    }

pd.DataFrame(
    [
        {
            "case": "lecture example (one test edge)",
            **rank_metrics([0.7], [[0.9, 0.4]], ks=(1, 2)),
        },
        {
            "case": "random scorer, 1000 edges × 99 negatives",
            **rank_metrics(
                rng.random(1000), rng.random((1000, 99)), ks=(1, 2)
            ),
        },
    ]
).round(3)

ex1_params = None
ex2_mrr = None

def block_diag_weight(blocks):
    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_block():
    try:
        g = torch.Generator().manual_seed(0)
        blocks = [torch.randn(4, 4, generator=g) for _ in range(3)]
        got = block_diag_weight(blocks)
        if got is None:
            return _verdict(False, False)
        want = torch.zeros(12, 12)
        for i, b in enumerate(blocks):
            want[4 * i : 4 * i + 4, 4 * i : 4 * i + 4] = b
        return _verdict(torch.allclose(got, want), True)
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(
            ex1_params == 4 * 64 * 64 * 3 + 20 * 4 * 3,
            ex1_params is not None,
        ),
        mo.md("**Exercise 2**"),
        _verdict(
            ex2_mrr is not None
            and abs(ex2_mrr - (1 / 2 + 1 + 1 / 4) / 3) < 2e-3,
            ex2_mrr is not None,
        ),
        mo.md("**Exercise 3**"),
        _check_block(),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    r"Bases $B\,d^2L=4\cdot4096\cdot3=49{,}152$ plus coefficients $|R|\,B\,L=240$ → **49,392** (plain: $20\cdot4096\cdot3=245{,}760$)."
                ),
                "Solution 2": mo.md(
                    r"Ranks 2, 1, 4 → RR 0.5, 1, 0.25 → MRR $=1.75/3\approx 0.583$. Check with `rank_metrics([0.8, 0.5, 0.3], [[0.9, 0.7, 0.1], [0.2, 0.3, 0.4], [0.6, 0.9, 0.8]])`."
                ),
                "Solution 3": mo.md(
                    "```python\ndef block_diag_weight(blocks):\n    return torch.block_diag(*blocks)\n```\nParameters: $k\\,(d/k)^2 = d^2/k$ per relation, but dimension $i$ only mixes with dimensions in its own block."
                ),
            }
        ),
    ]
)