"""Pure-operator tests for the v2 society (prereg §10): no GPU, no model."""

import numpy as np
import pytest

from llm.families import ALL_CANDIDATES, EXTRA_FAMILIES
from llm.society_ops import (arm_settings, choose_single_parent, families_alive, mating_plan,
                             novelty, pooled_survival, route_union)
from llm.tasks import make_tasks, verify


def test_extra_families_are_verifier_safe_and_deterministic():
    for fam in ALL_CANDIDATES:
        ts = make_tasks(fam, 200, seed=3)
        assert all(verify(t.answer, t) for t in ts), fam          # canonical answer verifies
        assert [t.prompt for t in make_tasks(fam, 200, seed=3)] == [t.prompt for t in ts]
        assert all(t.family == fam for t in ts)
    assert len(set(EXTRA_FAMILIES)) == 17 and len(set(ALL_CANDIDATES)) == 20


def test_pseudo_word_families_have_a_large_prompt_space():
    # The 20-word vocabulary gave sortletters 40 unique prompts; training would cover the test set.
    for fam in ("sortletters", "caesar", "charfreq"):
        assert len({t.prompt for t in make_tasks(fam, 600, seed=1)}) > 500, fam


def test_pooled_survival_is_e11_rule_and_greedy_at_lambda_zero():
    scores = np.array([0.9, 0.5, 0.5, 0.1])
    # agent 2 is behaviourally distant from everyone; agent 1 is a clone of agent 0
    dist = np.array([[0, 0.0, 0.9, 0.9],
                     [0.0, 0, 0.9, 0.9],
                     [0.9, 0.9, 0, 0.9],
                     [0.9, 0.9, 0.9, 0]], dtype=float)
    assert pooled_survival(scores, dist, 2, lam=0.0) == [0, 1]       # greedy: top-2 by score
    keep = pooled_survival(scores, dist, 2, lam=0.5)                  # QD: novelty lifts agent 2
    assert keep[0] == 0 and 2 in keep and 1 not in keep
    assert novelty(dist).argmax() == 2


def test_mating_plan_caps_use_and_prefers_distant_pairs():
    dist = np.array([[0, 0.9, 0.1, 0.2],
                     [0.9, 0, 0.3, 0.8],
                     [0.1, 0.3, 0, 0.7],
                     [0.2, 0.8, 0.7, 0]], dtype=float)
    plan = mating_plan(dist, 4, max_use=2)
    assert plan[0] == (0, 1)                                          # most distant pair first
    use = np.bincount(np.array(plan).ravel(), minlength=4)
    assert use.max() <= 2 and len(plan) == 4
    # every agent breeds at least once with N pairs and cap 2 — no allele is truncated at gen 1
    assert use.min() >= 1


def test_mating_plan_never_empty_when_cap_exhausts():
    dist = np.array([[0, 0.5], [0.5, 0]], dtype=float)
    plan = mating_plan(dist, 5, max_use=1)
    assert len(plan) == 5 and all(p == (0, 1) for p in plan)


def test_route_union_takes_the_more_confident_parent_and_is_deterministic_on_ties():
    a, ca = ["1", "2", "3"], np.array([0.9, 0.2, 0.5])
    b, cb = ["x", "y", "z"], np.array([0.1, 0.8, 0.5])
    out, src = route_union(a, ca, b, cb)
    assert out == ["1", "y", "3"] and src.tolist() == [0, 1, 0]


def test_choose_single_parent_is_score_proportional():
    rng = np.random.default_rng(0)
    picks = [choose_single_parent(np.array([0.0, 0.0, 1.0]), rng) for _ in range(300)]
    assert picks.count(2) > 250                                       # the fit parent dominates
    assert set(picks) <= {0, 1, 2}


def test_arm_settings_v2_table():
    assert arm_settings("full", 0.85) == {"g": 0.85, "sex": "union", "diversity": True}
    assert arm_settings("no_grounding", 0.85)["g"] == 0.0
    assert arm_settings("no_sex", 0.85)["sex"] is None
    assert arm_settings("no_diversity", 0.85)["diversity"] is False
    assert arm_settings("sex_linear", 0.85)["sex"] == "linear"
    with pytest.raises(ValueError):
        arm_settings("elitism", 0.85)


def test_families_alive_counts_competent_families_once():
    accs = [{"a": 0.9, "b": 0.1}, {"a": 0.7, "b": 0.2}, {"a": 0.0, "b": 0.61}]
    assert families_alive(accs, ["a", "b"]) == 2
    assert families_alive(accs, ["a", "b"], threshold=0.8) == 1