"""LLM-prototype tests — the pure, always-runnable parts (task generation + verifier).

The model/LoRA/merge path is heavy (downloads a base model, trains on a GPU) and is validated by the
experiment run itself, not in CI. What *is* unit-testable — and worth locking, since it is the
prototype's "reality that says no" — is that tasks are well-formed and the exact-match verifier
accepts correct answers (including verbose model phrasings) and rejects wrong ones.
"""

from __future__ import annotations

import numpy as np

from llm.directed import sample_merge_weights, select_winners
from llm.moe import learned_routes
from llm.tasks import FAMILIES, make_tasks, verify


def test_make_tasks_wellformed_and_deterministic():
    for fam in FAMILIES:
        tasks = make_tasks(fam, 20, seed=0)
        assert len(tasks) == 20 and all(t.family == fam for t in tasks)
        assert all(t.prompt and t.answer for t in tasks)
    a = make_tasks("arith", 10, seed=3)
    b = make_tasks("arith", 10, seed=3)
    assert [t.answer for t in a] == [t.answer for t in b]      # deterministic in the seed


def test_hard_tasks_wellformed_verifiable_and_distinct():
    # The hard variant must stay well-formed, self-verifying (canonical answer passes its own verifier),
    # and genuinely different from the easy variant (harder content, same family labels + answer format).
    for fam in FAMILIES:
        hard = make_tasks(fam, 30, seed=7, hard=True)
        assert len(hard) == 30 and all(t.family == fam for t in hard)
        assert all(t.prompt and t.answer for t in hard)
        assert all(verify(t.answer, t) for t in hard)              # canonical answers verify
        easy = make_tasks(fam, 30, seed=7, hard=False)
        assert [t.prompt for t in hard] != [t.prompt for t in easy]  # hard != easy
    # a Caesar-cipher answer is a real transform of the input (not the identity)
    caesars = [t for t in make_tasks("strings", 60, seed=2, hard=True) if "Caesar" in t.prompt]
    assert caesars and any(t.answer not in t.prompt for t in caesars)


def test_verifier_accepts_correct_including_verbose():
    tasks = make_tasks("lists", 40, seed=1) + make_tasks("arith", 40, seed=2)
    assert all(verify(t.answer, t) for t in tasks)              # the canonical answer verifies
    # a verbose but correct model phrasing still verifies (the verifier extracts the answer)
    num_task = next(t for t in tasks if t.family == "arith")
    assert verify(f"The answer is {num_task.answer}.", num_task)
    list_task = next(t for t in tasks if t.family == "lists" and t.answer.startswith("["))
    assert verify(f"Here you go: {list_task.answer}", list_task)


def test_verifier_rejects_wrong():
    t = make_tasks("arith", 1, seed=5)[0]
    wrong = str(int(t.answer) + 1) if t.answer.lstrip("-").isdigit() else "zzz"
    assert not verify(wrong, t)
    lt = next(x for x in make_tasks("lists", 30, seed=6) if x.answer.startswith("["))
    assert not verify("[9, 9, 9]", lt) or lt.answer == "[9, 9, 9]"


def test_learned_router_assigns_nearest_centroid():
    # Three well-separated families in a 4-D "embedding" space; the nearest-centroid router
    # (the MoE expert-selection gene) must route each test prompt to its own family's specialist.
    rng = np.random.default_rng(0)
    fams = ["lists", "strings", "arith"]
    anchors = {"lists": [5, 0, 0, 0], "strings": [0, 5, 0, 0], "arith": [0, 0, 5, 0]}
    train_emb = np.array([anchors[f] for f in fams for _ in range(8)], dtype=float)
    train_emb += rng.normal(scale=0.1, size=train_emb.shape)
    train_fam = np.array([f for f in fams for _ in range(8)])
    test_fam = np.array(["arith", "lists", "strings", "arith"])
    test_emb = np.array([anchors[f] for f in test_fam], dtype=float) + rng.normal(scale=0.1, size=(4, 4))
    routes = learned_routes(train_emb, train_fam, test_emb, fams)
    assert [fams[r] for r in routes] == list(test_fam)          # each routed to its own family


def test_learned_router_is_cosine_scale_invariant():
    # Cosine routing must ignore prompt-embedding magnitude (long vs short prompts): a test point on a
    # family's ray routes there regardless of its norm.
    fams = ["a", "b"]
    train_emb = np.array([[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]])
    train_fam = np.array(["a", "a", "b", "b"])
    test_emb = np.array([[10.0, 0.0], [0.0, 0.01]])            # very different magnitudes
    routes = learned_routes(train_emb, train_fam, test_emb, fams)
    assert [fams[r] for r in routes] == ["a", "b"]


def test_merge_weights_population_pins_baselines_and_diversifies():
    # The offspring population must contain the two canonical baselines (uniform soup, unit task-arith)
    # and be diverse + reproducible for the rest.
    rng = np.random.default_rng(0)
    w = sample_merge_weights(3, 16, rng)
    assert w.shape == (16, 3)
    assert np.allclose(w[0], 1 / 3)                             # candidate 0 = uniform soup
    assert np.allclose(w[1], 1.0)                              # candidate 1 = task arithmetic
    assert np.unique(w[2:].round(3), axis=0).shape[0] > 5      # the random offspring are diverse
    assert np.allclose(sample_merge_weights(3, 16, np.random.default_rng(0)), w)  # deterministic


def test_select_winners_picks_argmax_per_objective():
    val_overall = np.array([0.5, 0.9, 0.7])
    val_worst = np.array([0.4, 0.1, 0.6])                       # a different candidate is most balanced
    w = select_winners(val_overall, val_worst)
    assert w == {"overall": 1, "balanced": 2}


def test_merge_weights_requires_two_candidates():
    import pytest
    with pytest.raises(ValueError):
        sample_merge_weights(3, 1, np.random.default_rng(0))