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

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

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

    for _ in range(n_families):
        ga, gb = pid, pid + 1
        pid += 2  # grandparents
        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

kg_triples, kg_n_entities = family_kg()
_perm = np.random.default_rng(SEED).permutation(len(kg_triples))
kg_test, kg_train = (
    kg_triples[_perm[: len(_perm) // 10]],
    kg_triples[_perm[len(_perm) // 10 :]],
)
kg_known = {}
for _h, _r, _t in kg_triples:
    kg_known.setdefault((_h, _r), set()).add(_t)
mo.md(
    f"**{len(kg_triples):,} triples**, {kg_n_entities} people; per relation: "
    + ", ".join(
        f"{n} {c}"
        for n, c in zip(KG_RELATIONS, np.bincount(kg_triples[:, 1]))
    )
    + f". Train {len(kg_train):,} / test {len(kg_test):,}."
)

def train_kge(cls, epochs=300, dim=32, seed=SEED, **kw):
    torch.manual_seed(seed)
    model = cls(
        num_nodes=kg_n_entities,
        num_relations=len(KG_RELATIONS),
        hidden_channels=dim,
        **kw,
    )
    opt = torch.optim.Adam(model.parameters(), lr=0.01)
    h, r, t = (torch.from_numpy(kg_train[:, i]) for i in range(3))
    loader = model.loader(
        head_index=h,
        rel_type=r,
        tail_index=t,
        batch_size=512,
        shuffle=True,
    )
    for _ in range(epochs):
        for (
            hb,
            rb,
            tb,
        ) in (
            loader
        ):  # PyG samples corrupted (negative) triples inside model.loss
            opt.zero_grad()
            model.loss(hb, rb, tb).backward()
            opt.step()
    return model

def filtered_mrr_per_relation(model):
    out = {}
    all_t = torch.arange(kg_n_entities)
    with torch.no_grad():
        for ri, name in enumerate(KG_RELATIONS):
            rr = []
            for h, _, t in kg_test[kg_test[:, 1] == ri]:
                s = model(
                    torch.full((kg_n_entities,), int(h)),
                    torch.full((kg_n_entities,), ri),
                    all_t,
                )
                s[[x for x in kg_known[(h, ri)] if x != t]] = -float(
                    "inf"
                )  # filtered setting
                rr.append(1.0 / (1 + (s > s[t]).sum().item()))
            out[name] = np.mean(rr)
    return out

kge_models = {
    "TransE": (TransE, {"margin": 1.0}),
    "DistMult": (DistMult, {}),
    "ComplEx": (ComplEx, {}),
    "RotatE": (RotatE, {"margin": 9.0}),
}
kge_run = mo.ui.run_button(label="Train the four KGE models (≈ 15 s)")
kge_run

_rev = reveal(
    r"""
Measured (one seed; numbers move by a few points between runs):

* **TransE** fails both **symmetric** relations (spouse 0.46, sibling 0.44) yet is near-perfect on parent/child (inverse: 0.96/0.99) and
  **grandparent (composition: 1.00)**, exactly the table's TransE column.
* **DistMult** is mediocre everywhere (0.59–0.71). Because it scores $(h,r,t)$ and $(t,r,h)$ identically, it cannot tell
  "A parent_of B" from "B parent_of A", and that confusion spills over to every relation that shares those entities.
* **ComplEx** is solid everywhere (0.85–0.91), including composition, which it "can't" represent in general; on this small KG it
  approximates it.
* **RotatE** is best on composition (1.00) and directional relations (0.94–0.95), and decent but not perfect on symmetric ones
  (0.74–0.76): representable (θ = π) but harder to reach by training.

""",
    predict="Using the cheat-sheet table, mark for TransE, DistMult, ComplEx and RotatE which relations (symmetric spouse/sibling, inverse parent/child, composed grandparent) each should fail. Then check.",
)
_post = mo.md(r"""
Reading the table is the skill: check which patterns your KG has *before* choosing a model.

---

## 4 · Exercises

**Exercise 1 (TransE by hand).** $\mathbf h=(1,1)$, $\mathbf r=(0,2)$; candidates A $(1,3)$, B $(2,2)$, C $(0,3)$. Which candidate is ranked
**second**? → `ex1_second` (a string).

**Exercise 2 (RotatE composition).** RotatE relation $r_k$ is a vector of angles $\boldsymbol\theta_k$. Write `compose(theta1, theta2)`
returning the angles of $r_3$ with $r_3(x)=r_2(r_1(x))$, wrapped to $[0, 2\pi)$.

**Exercise 3 (filtered rank).** The true tail scores 0.6; the other candidates score 0.9 (*also a known true tail*), 0.7 and 0.2.
Raw rank is 3. What is the **filtered** rank? → `ex3_rank`.
""")
mo.vstack([_rev, _post])

ex1_second = None
ex3_rank = None

def compose(theta1, theta2):
    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_compose():
    try:
        g = np.random.default_rng(3)
        t1, t2 = g.uniform(0, 2 * np.pi, 5), g.uniform(0, 2 * np.pi, 5)
        got = compose(t1, t2)
        if got is None:
            return _verdict(False, False)
        h = g.normal(size=5) + 1j * g.normal(size=5)
        ok = np.allclose(
            h * np.exp(1j * t1) * np.exp(1j * t2),
            h * np.exp(1j * np.asarray(got)),
        ) and np.all(
            (np.asarray(got) >= 0) & (np.asarray(got) < 2 * np.pi)
        )
        return _verdict(ok, True)
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(ex1_second == "C", ex1_second is not None),
        mo.md("**Exercise 2**"),
        _check_compose(),
        mo.md("**Exercise 3**"),
        _verdict(ex3_rank == 2, ex3_rank is not None),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    r"$\mathbf h+\mathbf r=(1,3)$. Distances: A 0, C 1, B $\sqrt2\approx1.41$ → order A, **C**, B."
                ),
                "Solution 2": mo.md(
                    "```python\ndef compose(theta1, theta2):\n    return np.mod(np.asarray(theta1) + np.asarray(theta2), 2 * np.pi)\n```\nRotations compose by adding angles; this is why RotatE models composition (grandparent = parent ∘ parent: θ_gp = 2 θ_p)."
                ),
                "Solution 3": mo.md(
                    "Remove the other known true tail (0.9): only 0.7 beats 0.6 → filtered rank **2** (raw rank 3)."
                ),
            }
        ),
    ]
)

def mrr(model, rel, filtered):
    rr, all_t = [], torch.arange(kg_n_entities)
    with torch.no_grad():
        for h, _, t in kg_test[kg_test[:, 1] == rel]:
            s = model(
                torch.full((kg_n_entities,), int(h)),
                torch.full((kg_n_entities,), rel),
                all_t,
            )
            if filtered:
                s[[x for x in kg_known[(h, rel)] if x != t]] = -float(
                    "inf"
                )
            rr.append(1.0 / (1 + (s > s[t]).sum().item()))
    return np.mean(rr)

_m = train_kge(ComplEx)
pd.DataFrame(
    [
        {
            "relation": KG_RELATIONS[r],
            "raw MRR": mrr(_m, r, False),
            "filtered MRR": mrr(_m, r, True),
        }
        for r in range(5)
    ]
).round(3)