src/test_pairs.py

"""
Synthetic tests for E3/E4 preference-pair construction. Catches logic errors now
rather than four hours into the pipeline.

Most important: test_divpo_prob_picks_least_probable_as_chosen. DivPO's
probability variant inverts the usual intuition -- the CHOSEN story is the one
the model found LEAST likely. Getting that backwards would train the policy
straight into the greedy mode while the logs looked entirely normal.
"""
import numpy as np

from build_pairs import build_divpo, build_multipos, normalize_weights


def make_item(pid, idx, quality, dev, meanlp, passed=True, text=None):
    return {"prompt_id": pid, "idx": idx, "prompt": f"prompt-{pid}",
            "text": text or f"story-{pid}-{idx}", "quality": quality,
            "deviation": dev, "mean_logprob": meanlp, "gate_passed": passed,
            "marginal": -1.0, "n_words": 300}


def simple_pool():
    """One prompt, 6 stories spanning quality/deviation/logprob."""
    return {"p0": [
        make_item("p0", 0, 8.0, 0.40, -1.50),   # high q, MOST diverse, least probable
        make_item("p0", 1, 7.0, 0.20, -0.80),   # high q, low dev  -> "competent cliche"
        make_item("p0", 2, 6.5, 0.30, -1.10),
        make_item("p0", 3, 3.0, 0.35, -0.60),   # low q, HIGH dev  -> must NOT be chosen
        make_item("p0", 4, 2.0, 0.05, -0.40),   # low q, least diverse, MOST probable
        make_item("p0", 5, 1.0, 0.10, -0.90, passed=False),
    ]}


# --------------------------------------------------------------------- E4
def test_divpo_emb_chosen_is_most_diverse_above_rho():
    rows, st = build_divpo(simple_pool(), "emb", rho=6.0)
    assert st["n_rows"] == 1
    r = rows[0]
    assert r["chosen"] == "story-p0-0", r
    assert r["chosen_quality"] >= 6.0


def test_divpo_emb_never_chooses_a_diverse_but_low_quality_story():
    """Story 3 has high deviation but quality 3.0. Quality gating must exclude
    it -- this is the 'different because it is worse' failure."""
    rows, _ = build_divpo(simple_pool(), "emb", rho=6.0)
    assert rows[0]["chosen"] != "story-p0-3"


def test_divpo_emb_rejected_is_least_diverse_below_rho():
    rows, _ = build_divpo(simple_pool(), "emb", rho=6.0)
    assert rows[0]["rejected"] == "story-p0-4", rows[0]


def test_divpo_prob_picks_least_probable_as_chosen():
    """Most diverse == LOWEST length-normalized logprob; least diverse ==
    HIGHEST (the near-greedy sample). Inverting this trains toward the mode."""
    rows, _ = build_divpo(simple_pool(), "prob", rho=6.0)
    r = rows[0]
    assert r["chosen"] == "story-p0-0", f"chosen should be least probable: {r}"
    assert r["rejected"] == "story-p0-4", f"rejected should be most probable: {r}"
    assert r["chosen_meanlp"] < r["rejected_meanlp"]


def test_divpo_skips_prompt_with_no_qualifying_chosen():
    pool = {"p0": [make_item("p0", i, 3.0, 0.2, -1.0) for i in range(4)]}
    rows, st = build_divpo(pool, "emb", rho=6.0)
    assert rows == [] and st["no_chosen"] == 1 and st["skip_rate"] == 1.0


def test_divpo_skips_prompt_with_no_rejected():
    pool = {"p0": [make_item("p0", i, 9.0, 0.2 + 0.01 * i, -1.0) for i in range(4)]}
    rows, st = build_divpo(pool, "emb", rho=6.0)
    assert rows == [] and st["no_rejected"] == 1


def test_divpo_gate_failures_are_eligible_rejects():
    """A gate-failing story is legitimately in the low-quality tail."""
    pool = {"p0": [make_item("p0", 0, 8.0, 0.4, -1.5),
                   make_item("p0", 1, 9.0, 0.3, -1.2),
                   make_item("p0", 2, 0.0, 0.01, -0.3, passed=False)]}
    rows, st = build_divpo(pool, "emb", rho=6.0)
    assert st["n_rows"] == 1 and rows[0]["rejected"] == "story-p0-2"


def test_divpo_is_one_row_per_prompt():
    pool = {f"p{i}": simple_pool()["p0"] for i in range(5)}
    rows, st = build_divpo(pool, "emb", rho=6.0)
    assert len(rows) == 5 == st["n_rows"]
    assert all(r["weight"] == 1.0 for r in rows), "DivPO must be unweighted"


# --------------------------------------------------------------------- E3
def _emb_for(pool):
    """Deterministic embeddings + row index aligned to pool iteration order."""
    idx, vecs, i = {}, [], 0
    rng = np.random.default_rng(0)
    for pid, items in pool.items():
        for it in items:
            idx[(pid, it["idx"])] = i
            v = rng.standard_normal(32)
            vecs.append(v / np.linalg.norm(v))
            i += 1
    return np.array(vecs, dtype=np.float32), idx


def test_multipos_emits_multiple_chosen_rows_per_prompt():
    pool = simple_pool()
    emb, idx = _emb_for(pool)
    rows, st = build_multipos(pool, emb, idx, q_keep=5.0, k=4, lam=1.0)
    assert st["ok"] == 1
    assert 2 <= len(rows) <= 4, f"expected up to 4 chosen rows, got {len(rows)}"
    assert all(r["chosen_quality"] >= 5.0 for r in rows)


def test_multipos_rotates_negatives():
    """Don't hammer one rejected across all rows."""
    pool = simple_pool()
    emb, idx = _emb_for(pool)
    rows, _ = build_multipos(pool, emb, idx, q_keep=5.0, k=4, lam=1.0)
    if len(rows) >= 2:
        assert len(set(r["neg_type"] for r in rows)) == 2, \
            f"negatives not rotated: {[r['neg_type'] for r in rows]}"


def test_multipos_weight_is_chosen_deviation():
    pool = simple_pool()
    emb, idx = _emb_for(pool)
    rows, _ = build_multipos(pool, emb, idx, q_keep=5.0, k=4, lam=1.0)
    for r in rows:
        assert abs(r["weight"] - r["chosen_dev"]) < 1e-9


def test_multipos_skips_prompt_with_too_few_survivors():
    pool = {"p0": [make_item("p0", i, 2.0, 0.2, -1.0) for i in range(6)]}
    emb, idx = _emb_for(pool)
    rows, st = build_multipos(pool, emb, idx, q_keep=5.0, k=4, lam=1.0)
    assert rows == [] and st["skipped_no_survivors"] == 1


def test_normalize_weights_gives_mean_one():
    rows = [{"weight": w} for w in (0.1, 0.2, 0.3, 0.8)]
    normalize_weights(rows)
    w = np.array([r["weight"] for r in rows])
    assert abs(w.mean() - 1.0) < 1e-9
    # relative ordering preserved -> it reweights, it does not rescale the LR
    assert np.all(np.diff(w) > 0)


def test_normalize_weights_handles_all_zero():
    rows = [{"weight": 0.0} for _ in range(3)]
    normalize_weights(rows)
    assert all(r["weight"] == 0.0 for r in rows)


if __name__ == "__main__":
    import sys, traceback
    fns = [(n, f) for n, f in sorted(globals().items())
           if n.startswith("test_") and callable(f)]
    bad = 0
    for n, f in fns:
        try:
            f(); print(f"  PASS  {n}")
        except Exception:
            bad += 1; print(f"  FAIL  {n}"); traceback.print_exc()
    print(f"\n{len(fns)-bad}/{len(fns)} passed")
    sys.exit(1 if bad else 0)