#!/usr/bin/env python3
"""Hands-on P08 — a recommender, assembled from eight lego blocks."""
import time
import numpy as np
from collections import defaultdict
from _harness import block, run_all
NI, D = 1200, 8

def generate(NU=800, beta=0.5, seed=8):
    """beta = fraction of choice mass that is pure popularity, 1-beta is taste."""
    rng = np.random.default_rng(seed)
    U = rng.normal(0, 1, (NU, D)); V = rng.normal(0, 1, (NI, D))
    pop = 1.0 / np.arange(1, NI + 1) ** 0.9; pop /= pop.sum()
    ev = {}
    for u in range(NU):
        n = int(rng.integers(12, 60))
        aff = U[u] @ V.T; sm = np.exp(aff - aff.max()); sm /= sm.sum()
        p = (1 - beta) * sm + beta * pop
        ev[u] = [int(i) for i in rng.choice(NI, size=n, replace=False, p=p / p.sum())]
    return NU, ev

def counts(tr):
    c = np.zeros(NI)
    for u in tr:
        for i in tr[u]: c[i] += 1
    return c

def train_bpr(NU, tr, alpha=None, d=32, epochs=30, lr=.05, reg=.05, bs=4096, seed=1):
    """Minibatch BPR-SGD. alpha=None -> uniform negatives, else q ~ (count+1)^alpha."""
    rng = np.random.default_rng(seed)
    P = rng.normal(0, .1, (NU, d)); Q = rng.normal(0, .1, (NI, d))
    us = np.array([u for u in tr for _ in tr[u]])
    it = np.array([i for uu in tr for i in tr[uu]])
    cum = None
    if alpha is not None:
        pj = (counts(tr) + 1.0) ** alpha; cum = np.cumsum(pj / pj.sum())
    for _ in range(epochs):
        perm = rng.permutation(len(us))
        for b in range(0, len(us), bs):
            idx = perm[b:b+bs]; u, i = us[idx], it[idx]
            j = (rng.integers(0, NI, len(idx)) if cum is None
                 else np.searchsorted(cum, rng.random(len(idx))))
            Pu, Qi, Qj = P[u], Q[i], Q[j]
            g = (1 / (1 + np.exp(np.einsum("ij,ij->i", Pu, Qi - Qj))))[:, None]
            np.add.at(P, u, lr * (g * (Qi - Qj) - reg * Pu))
            np.add.at(Q, i, lr * (g * Pu - reg * Qi))
            np.add.at(Q, j, lr * (-g * Pu - reg * Qj))
    return P, Q

def train_mse(NU, tr, d=32, epochs=30, lr=.05, reg=.05, bs=4096, seed=1):
    """Regression on implicit feedback: every observed pair has target 1.0."""
    rng = np.random.default_rng(seed)
    P = rng.normal(0, .1, (NU, d)); Q = rng.normal(0, .1, (NI, d))
    us = np.array([u for u in tr for _ in tr[u]])
    it = np.array([i for uu in tr for i in tr[uu]])
    for _ in range(epochs):
        perm = rng.permutation(len(us))
        for b in range(0, len(us), bs):
            idx = perm[b:b+bs]; u, i = us[idx], it[idx]
            Pu, Qi = P[u], Q[i]
            e = (1.0 - np.einsum("ij,ij->i", Pu, Qi))[:, None]
            np.add.at(P, u, lr * (e * Qi - reg * Pu))
            np.add.at(Q, i, lr * (e * Pu - reg * Qi))
    return P, Q

@block(1, "Synthesise a world whose answer you know", "you cannot debug a recommender on real data")
def b1(s, show):
    NU, ev = generate()
    if show:
        n = sum(len(v) for v in ev.values()); c = counts(ev); top = np.sort(c)[::-1]
        print(f"  {n} interactions, {NU} users, {NI} items, true latent dim {D}")
        print(f"  head mass: top 1% of items = {100*top[:12].sum()/n:>4.1f}% of events;"
              f"  top 10% = {100*top[:120].sum()/n:.1f}%")
        print(f"  {int((c==0).sum())} items ({100*(c==0).mean():.0f}%) never touched -- "
              "the cold tail is the default state")
        print("  Ground truth U, V and the popularity mixture beta are KNOWN here. On")
        print("  MovieLens they are not, so every bug looks like 'the model is bad'.")
    return {"NU": NU, "ev": ev}

@block(2, "The split decides the number", "a random split leaks the future and inflates everything")
def b2(s, show):
    def split(ev, mode="temporal", seed=3):
        rng = np.random.default_rng(seed)
        tr, te = defaultdict(set), {}
        for u, xs in ev.items():
            j = len(xs) - 1 if mode == "temporal" else int(rng.integers(0, len(xs)))
            te[u] = xs[j]; tr[u] = {x for k, x in enumerate(xs) if k != j}
        return tr, te
    tr, te = split(s["ev"], "temporal")
    rtr, rte = split(s["ev"], "random")
    if show:
        print("  Two splits of the SAME data, one held-out event per user:")
        print("    temporal — hold out each user's LAST event (honest)")
        print("    random   — hold out a uniformly chosen event (leaks the future)")
        print(f"  train sizes are identical ({sum(map(len,tr.values()))} vs "
              f"{sum(map(len,rtr.values()))} pairs), so any difference in score is")
        print("  purely the leak. Block 7 measures how big it is.")
    return {"split": split, "tr": tr, "te": te, "rtr": rtr, "rte": rte}

@block(3, "The baseline that embarrasses you", "popularity is not a strawman")
def b3(s, show):
    def evaluate(score, tr, te, k=20):
        rec = ndcg = 0.0
        for u in te:
            order = [i for i in score(u) if i not in tr[u]][:k]
            if te[u] in order:
                rec += 1; ndcg += 1 / np.log2(order.index(te[u]) + 2)
        return rec / len(te), ndcg / len(te)
    pop_order = np.argsort(-counts(s["tr"]))
    r, n = evaluate(lambda u: pop_order, s["tr"], s["te"])
    if show:
        print(f"  popularity, temporal split:   recall@20={r:.4f}   ndcg@20={n:.4f}")
        print(f"  uniform random guessing:      recall@20={20/NI:.4f}")
        print(f"  Popularity is {r/(20/NI):.0f}x random and costs one bincount. Any model")
        print("  that does not clear this line has learned popularity and nothing")
        print("  else -- and you only find out by computing this row.")
    return {"evaluate": evaluate, "pop_order": pop_order}

@block(4, "BPR: rank, don't predict", "the loss must match the task, or the metric punishes you")
def b4(s, show):
    NU, tr, te, ev = s["NU"], s["tr"], s["te"], s["evaluate"]
    P, Q = train_bpr(NU, tr); Pm, Qm = train_mse(NU, tr)
    if show:
        rp, np_ = ev(lambda u: s["pop_order"], tr, te)
        rm, nm = ev(lambda u: np.argsort(-(Pm[u] @ Qm.T)), tr, te)
        rb, nb = ev(lambda u: np.argsort(-(P[u] @ Q.T)), tr, te)
        print(f"  {'model':<28}{'recall@20':>11}{'ndcg@20':>10}{'vs popularity':>15}")
        for lbl, r, n in (("popularity", rp, np_), ("MF + squared error", rm, nm),
                          ("MF + BPR (pairwise)", rb, nb)):
            print(f"  {lbl:<28}{r:>11.4f}{n:>10.4f}{r/rp:>14.2f}x")
        print("  Same architecture, same d, same data, same epochs. Only the loss")
        print("  differs, and squared error lands near random. It asks 'what score?';")
        print("  BPR asks 'which of these two?' -- the question recall@20 grades.")
    return {"P": P, "Q": Q}

@block(5, "Regularisation: the cliff I fell off", "an under-regularised MF scores BELOW popularity")
def b5(s, show):
    NU, tr, te, ev = s["NU"], s["tr"], s["te"], s["evaluate"]
    if show:
        rp, _ = ev(lambda u: s["pop_order"], tr, te)
        print(f"  popularity baseline = {rp:.4f}. Recall@20 as training proceeds:")
        print(f"  {'epochs':>8}{'reg=0.01':>11}{'reg=0.05':>11}")
        for epochs in (10, 30, 60, 150):
            row = []
            for reg in (0.01, 0.05):
                P, Q = train_bpr(NU, tr, epochs=epochs, reg=reg)
                row.append(ev(lambda u: np.argsort(-(P[u] @ Q.T)), tr, te)[0])
            print(f"  {epochs:>8}{row[0]:>11.4f}{row[1]:>11.4f}"
                  f"{'   <-- below baseline' if row[0] < rp else ''}")
        print("  This block exists because the first version of this file shipped")
        print("  reg=0.01 and concluded 'MF cannot beat popularity'. It was not a")
        print("  fact about matrix factorisation; it was one hyperparameter. With 32")
        print("  free parameters per user fit from ~35 events, the model memorises")
        print("  the training set and pushes every unobserved item down -- including")
        print("  the held-out one. MORE training makes it WORSE, which is the")
        print("  signature of overfitting and not of a bad architecture.")
    return {}

@block(6, "Negative sampling: a prediction, then a test", "BPR's optimum ranks by p(i|u)/q(i)")
def b6(s, show):
    if show:
        print("  Theory: BPR with negatives drawn from q converges to a ranking by")
        print("  p(i|u)/q(i) -- the same importance-weighting that makes NCE work.")
        print("  PREDICTION: sampling negatives proportional to popularity DIVIDES OUT")
        print("  the popularity signal. If the truth is popularity-heavy that should")
        print("  be catastrophic; if the truth has no popularity component (beta=0)")
        print("  it should be harmless.\n")
        print(f"  {'beta (popularity mass)':<24}{'uniform q':>11}{'q ~ pop^0.75':>14}"
              f"{'damage':>9}")
        rows = {}
        for beta in (0.5, 0.2, 0.0):
            NU, ev_ = generate(beta=beta)
            tr, te = s["split"](ev_, "temporal")
            P1, Q1 = train_bpr(NU, tr)                 # uniform negatives
            P2, Q2 = train_bpr(NU, tr, alpha=0.75)     # q ~ popularity^0.75
            po = np.argsort(-counts(tr))
            a = s["evaluate"](lambda u: np.argsort(-(P1[u] @ Q1.T)), tr, te)[0]
            b = s["evaluate"](lambda u: np.argsort(-(P2[u] @ Q2.T)), tr, te)[0]
            pr = s["evaluate"](lambda u: po, tr, te)[0]
            rows[beta] = (pr, a, b)
            print(f"  {beta:<24.1f}{a:>11.4f}{b:>14.4f}{b/a:>8.2f}x")
        print("  VERDICT: partially confirmed. The damage shrinks monotonically as the")
        print("  popularity mass falls (0.42x -> 0.68x -> 0.75x), exactly as predicted,")
        print("  but it does not vanish at beta=0. A second mechanism is also present:")
        print("  under q ~ pop, tail items are almost never sampled as negatives, so")
        print("  their embeddings stay near random initialisation and rank spuriously.")
        print("  word2vec uses pop^0.75 because there discounting frequency is the")
        print("  GOAL. Copying the constant into a recommender inverts its purpose.")
        s["beta_rows"] = rows
    return {}

@block(7, "Two-stage: the ceiling you cannot re-rank past", "stage 2 can only reorder what stage 1 returned")
def b7(s, show):
    NU, tr, te, P, Q = s["NU"], s["tr"], s["te"], s["P"], s["Q"]
    pop_order = s["pop_order"]
    def two_stage(u, C, k=20):
        cands = [int(i) for i in pop_order[:C] if i not in tr[u]]   # cheap, no user model
        sc = P[u] @ Q[cands].T                                       # expensive, per-user
        return [cands[j] for j in np.argsort(-sc)][:k]
    if show:
        full = s["evaluate"](lambda u: np.argsort(-(P[u] @ Q.T)), tr, te)[0]
        print(f"  stage 1 = popularity top-C (one bincount, shared by all users)")
        print(f"  stage 2 = the BPR model, scoring only those C items\n")
        print(f"  {'C':>6}{'stage-1 ceiling':>17}{'after re-rank':>15}"
              f"{'scores/user':>13}")
        for C in (20, 50, 200, 600, NI):
            ceil = sum(1 for u in te if te[u] in set(int(i) for i in pop_order[:C])) / len(te)
            hit = sum(1 for u in te if te[u] in two_stage(u, C)) / len(te)
            print(f"  {C:>6}{ceil:>17.4f}{hit:>15.4f}{C:>13}")
        print(f"  single-stage full scan: {full:.4f} using {NI} scores per user")
        print("  Re-ranking never exceeds the ceiling -- it is a hard cap, not a")
        print("  tendency. Before blaming the ranker for a miss, check whether the")
        print("  item was in the candidate set at all. This is also the join to P02:")
        print("  swap popularity top-C for an HNSW query and the ceiling becomes")
        print("  recall@C of the index, which is the number that project measured.")
    return {"two_stage": two_stage}

@block(8, "How much headroom exists at all", "the data, not the model, sets the ceiling")
def b8(s, show):
    if show:
        print(f"  {'beta (popularity mass)':<24}{'popularity':>12}{'MF+BPR':>9}"
              f"{'model uplift':>14}")
        for beta in (0.5, 0.2, 0.0):
            pr, mf, _ = s["beta_rows"][beta]
            print(f"  {beta:<24.1f}{pr:>12.4f}{mf:>9.4f}{mf/pr:>13.2f}x")
        print("  Identical model, identical hyperparameters, three worlds. When half")
        print("  the choices are pure popularity there is a 1.07x model to be won;")
        print("  when none are, there is a 6x one. Personalisation uplift is a")
        print("  property of the DOMAIN. Before a quarter of modelling work, estimate")
        print("  the head mass -- it tells you the size of the prize.")
    return {}

def assembly(s):
    print("\nEight blocks = a recommender. One table, temporal split, honest rows.\n")
    NU, tr, te, ev = s["NU"], s["tr"], s["te"], s["evaluate"]
    P, Q = s["P"], s["Q"]
    Pm, Qm = train_mse(NU, tr)
    Pu, Qu = train_bpr(NU, tr, reg=0.01, epochs=150)
    Pa, Qa = train_bpr(NU, tr, alpha=0.75)
    base = ev(lambda u: s["pop_order"], tr, te)
    rows = [("uniform random", (20/NI, 0.0)),
            ("popularity", base),
            ("MF, squared error", ev(lambda u: np.argsort(-(Pm[u]@Qm.T)), tr, te)),
            ("MF, BPR, reg=0.01, 150ep", ev(lambda u: np.argsort(-(Pu[u]@Qu.T)), tr, te)),
            ("MF, BPR, q ~ pop^0.75", ev(lambda u: np.argsort(-(Pa[u]@Qa.T)), tr, te)),
            ("MF, BPR, tuned", ev(lambda u: np.argsort(-(P[u]@Q.T)), tr, te))]
    hit = sum(1 for u in te if te[u] in s["two_stage"](u, 600)) / len(te)
    rows.append(("  served two-stage, C=600", (hit, float("nan"))))
    print(f"  {'system':<28}{'recall@20':>11}{'ndcg@20':>10}{'vs popularity':>15}")
    for lbl, (r, n) in rows:
        nn = "   -- " if n != n else f"{n:.4f}"
        print(f"  {lbl:<28}{r:>11.4f}{nn:>10}{r/base[0]:>14.2f}x")
    rt = ev(lambda u: np.argsort(-(P[u]@Q.T)), tr, te)[0]
    Pr, Qr = train_bpr(NU, s["rtr"]); rr = ev(lambda u: np.argsort(-(Pr[u]@Qr.T)),
                                              s["rtr"], s["rte"])[0]
    print(f"\n  the same tuned model, scored on the RANDOM split: {rr:.4f} "
          f"({rr/rt:.2f}x)")
    print("  Nothing changed but which event was hidden. A number reported without")
    print("  naming its split is not comparable to anything.")
    print("\n  Read the table top to bottom. The largest single jump is random ->")
    print("  popularity, and it required no model at all. Three of the four MF rows")
    print("  score BELOW that baseline -- one for the wrong loss, one for weak")
    print("  regularisation, one for copying word2vec's sampling constant. The tuned")
    print("  row wins by 1.07x -- an unimpressive number until block 8, where the")
    print("  IDENTICAL model and hyperparameters win 6.11x on data with no popularity")
    print("  mass. The model was never the limiting factor here; the domain was.")
    print("\n  Built: synthetic ground truth -> split discipline -> popularity ->")
    print("  loss choice -> regularisation -> negative-sampling distribution ->")
    print("  two-stage retrieval -> headroom analysis.")
    print("  Missing, on the project page: real MovieLens/Amazon ingest (m1), item")
    print("  and user features for cold start (m7), a served HNSW index in place of")
    print("  popularity top-C (m9, reusing P02), latency budgets under load (m11),")
    print("  and E9 -- the diversity/accuracy tradeoff where recall@20 goes DOWN and")
    print("  the system gets better.")

if __name__ == "__main__":
    run_all(assembly, "HANDS-ON P08 — Recommender systems, block by block")