FAMILY_RELS = [
    "spouse",
    "parent_of",
    "child_of",
    "grandparent_of",
    "sibling_of",
]

def family_kg(n_families, seed):
    r = np.random.default_rng(seed)
    triples, pid = set(), 0

    def add(h, rel, t):
        triples.add((h, FAMILY_RELS.index(rel), t))

    for _ in range(n_families):
        ga, gb = pid, pid + 1
        pid += 2
        add(ga, "spouse", gb)
        add(gb, "spouse", ga)
        kids = []
        for _ in range(int(r.integers(2, 4))):
            c, partner = pid, pid + 1
            pid += 2
            kids.append(c)
            add(c, "spouse", partner)
            add(partner, "spouse", c)
            for p in (ga, gb):
                add(p, "parent_of", c)
                add(c, "child_of", p)
            grandkids = []
            for _ in range(int(r.integers(1, 4))):
                k = pid
                pid += 1
                grandkids.append(k)
                for p in (c, partner):
                    add(p, "parent_of", k)
                    add(k, "child_of", p)
                for gp in (ga, gb):
                    add(gp, "grandparent_of", k)
            for a in grandkids:
                for b in grandkids:
                    if a != b:
                        add(a, "sibling_of", b)
        for a in kids:
            for b in kids:
                if a != b:
                    add(a, "sibling_of", b)
    return np.array(sorted(triples)), pid

def split_triples(T, seed, frac=0.2):
    p = np.random.default_rng(seed).permutation(len(T))
    k = int(frac * len(T))
    return T[p[k:]], T[p[:k]]  # (message triples, query triples)

def relation_graph(msg, n_rel, n_ent):
    """4 x n_rel x n_rel adjacency: h2h, t2t, h2t, t2h (relations sharing an entity in those roles)."""
    Eh = np.zeros((n_ent, n_rel))
    Et = np.zeros((n_ent, n_rel))
    Eh[msg[:, 0], msg[:, 1]] = 1
    Et[msg[:, 2], msg[:, 1]] = 1
    A = np.stack([Eh.T @ Eh, Et.T @ Et, Eh.T @ Et, Et.T @ Eh]) > 0
    return torch.tensor(A, dtype=torch.float)

# KG "A" for training; KG "B": different people AND relation IDs shuffled (a new KG to the model)
kg_A, n_ent_A = family_kg(40, seed=0)
kg_B_raw, n_ent_B = family_kg(40, seed=1)
REL_SHUFFLE = np.array([3, 0, 4, 1, 2])
kg_B = kg_B_raw.copy()
kg_B[:, 1] = REL_SHUFFLE[kg_B_raw[:, 1]]
msg_A, query_A = split_triples(kg_A, seed=0)
msg_B, query_B = split_triples(kg_B, seed=1)
mo.md(
    f"KG A: {len(kg_A):,} triples, {n_ent_A} entities. KG B: {len(kg_B):,} triples, {n_ent_B} **different** entities, "
    f"relation IDs shuffled by {REL_SHUFFLE.tolist()} (old id → new id). Each KG: 80% message triples, 20% queries."
)

class ULTRALite(nn.Module):
    def __init__(
        self, d=32, n_layers=3, relation_inductive=True, n_relation_ids=5
    ):
        super().__init__()
        self.d, self.relation_inductive = d, relation_inductive
        self.rel_edge = nn.Parameter(
            torch.randn(4, d) * 0.1
        )  # one vector per relation-graph edge type
        self.rel_layers = nn.ModuleList(
            nn.Linear(2 * d, d) for _ in range(n_layers)
        )
        self.rel_id_emb = nn.Embedding(
            n_relation_ids, d
        )  # baseline only: a vector per relation ID
        self.ent_layers = nn.ModuleList(
            nn.Linear(2 * d, d) for _ in range(n_layers)
        )
        self.score = nn.Sequential(
            nn.Linear(d, d), nn.ReLU(), nn.Linear(d, 1)
        )

    def relation_reps(self, A_rel, q):
        B, n_rel = len(q), A_rel.size(1)
        if not self.relation_inductive:
            R = (
                self.rel_id_emb.weight[:n_rel]
                .unsqueeze(0)
                .expand(B, -1, -1)
            )
            return R, R[torch.arange(B), q]
        h = torch.zeros(B, n_rel, self.d)
        h[torch.arange(B), q] = 1.0  # query relation = all-ones
        for lin in self.rel_layers:
            msg = sum(
                torch.einsum("ij,bjd->bid", A_rel[k], h) * self.rel_edge[k]
                for k in range(4)
            )
            h = F.relu(lin(torch.cat([h, msg], -1))) + h
        return h, h[torch.arange(B), q]

    def forward(self, msg, n_ent, A_rel, heads, q):
        R, r_q = self.relation_reps(A_rel, q)
        x = torch.zeros(len(heads), n_ent, self.d)
        x[torch.arange(len(heads)), heads] = r_q  # labelling trick
        src, rel, dst = (torch.from_numpy(msg[:, i]) for i in range(3))
        for lin in self.ent_layers:
            m = x[:, src] * R[:, rel]  # x_u ⊙ R_r along every triple
            x = (
                F.relu(
                    lin(
                        torch.cat(
                            [x, torch.zeros_like(x).index_add_(1, dst, m)],
                            -1,
                        )
                    )
                )
                + x
            )
        return self.score(x).squeeze(-1)  # [batch, n_ent] tail scores

def train_ultra(model, epochs=30, bs=64, seed=SEED):
    torch.manual_seed(seed)
    r = np.random.default_rng(seed)
    A_rel = relation_graph(msg_A, 5, n_ent_A)
    opt = torch.optim.Adam(model.parameters(), lr=5e-3)
    for _ in range(epochs):
        order = r.permutation(len(query_A))
        for i in range(0, len(order), bs):
            b = query_A[order[i : i + bs]]
            logits = model(
                msg_A,
                n_ent_A,
                A_rel,
                torch.from_numpy(b[:, 0]),
                torch.from_numpy(b[:, 1]),
            )
            opt.zero_grad()
            F.cross_entropy(logits, torch.from_numpy(b[:, 2])).backward()
            opt.step()
    return model

def filtered_mrr(model, msg, queries, n_ent, all_triples):
    A_rel = relation_graph(msg, 5, n_ent)
    known = {}
    for h, r, t in all_triples:
        known.setdefault((h, r), set()).add(t)
    rr = []
    with torch.no_grad():
        for i in range(0, len(queries), 128):
            b = queries[i : i + 128]
            logits = model(
                msg,
                n_ent,
                A_rel,
                torch.from_numpy(b[:, 0]),
                torch.from_numpy(b[:, 1]),
            )
            for j, (h, r, t) in enumerate(b):
                s = logits[j].clone()
                s[[x for x in known[(h, r)] if x != t]] = -float("inf")
                rr.append(1.0 / (1 + (s > s[t]).sum().item()))
    return float(np.mean(rr))

ultra_run = mo.ui.run_button(
    label="Train both models on KG A and test on KG B (≈ 40 s)"
)
ultra_run

_rev = reveal(
    r"""
Measured (one seed, ≈ 17 s per model):

| | KG A (training) | KG B: new entities | KG B: new entities + shuffled relation IDs |
|---|---|---|---|
| **ULTRA-lite** | 0.963 | 0.962 | **0.962** |
| **relation-ID baseline** | 0.973 | 0.967 | **0.105** |

* Both models are **entity-inductive**: on 482 people they have never seen, MRR stays ≈ 0.96–0.97. Entities are computed from
  structure (the labelling trick + message passing), as in R-GCN / NBFNet.
* Shuffle the relation IDs, which is what a new KG with its own relation vocabulary looks like, and the ID-based model collapses
  to **0.105**: its "parent_of" vector now means "grandparent_of". ULTRA-lite is **unchanged (0.962)**, because it never looks at relation
  IDs, only at how relations co-occur. That is double equivariance, and it is what lets ULTRA be pretrained on a few KGs and
  used zero-shot on dozens of others.

""",
    predict="Shuffle the relation IDs of the new KG. What happens to the relation-ID model's MRR, and to ULTRA-lite's?",
)
_post = mo.md(r"""
---

## 5 · Exercises

**Exercise 1 (relation graph).** Triples $(a,r_1,b)$, $(a,r_2,c)$, $(d,r_3,a)$. How many *pairs* of distinct relations are linked
in the relation graph by any of the four edge types? → `ex1_pairs`.

**Exercise 2 (tail part).** For the movie KG of §2, compute the tail-side weight $A_t[\text{actedIn},\text{directed}]$ with
$A_t=E_t^\top D_t^{-2}E_t$ ($E_t[e,r]=1$ if $e$ is a tail of $r$; $D_t$ = number of triples in which $e$ is the tail). → `ex2_At`.
""")
mo.vstack([_rev, _post])

ex1_pairs = None
ex2_At = 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_pairs == 3, ex1_pairs is not None),
        mo.md("**Exercise 2**"),
        _verdict(
            ex2_At is not None and abs(ex2_At - 0.5) < 1e-6,
            ex2_At is not None,
        ),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    r"$r_1$–$r_2$ (head–head via $a$); $r_1$–$r_3$ and $r_2$–$r_3$ (head–tail / tail–head via $a$, head of $r_1,r_2$ and tail of $r_3$) → **3** pairs."
                ),
                "Solution 2": mo.md(
                    r"Shared tails: *Poor Things* (tail of actedIn and directed, tail degree 2 → $1/4$) and *The Favourite* (same → $1/4$); *La La Land* is only an actedIn tail → **0.5**."
                ),
            }
        ),
    ]
)