"""Pure-schedule tests for the curriculum society (prereg v4 §7): no GPU, no model."""

import numpy as np
import pytest

from llm.curriculum import (complementarity, cumulative_accuracy, forgetting, latin_square,
                            partner_for, replay_split, resolve_orders)

FAMS = ["mnli", "arc", "hellaswag", "squad", "boolq", "winogrande"]
DECOR = [["mnli", "arc", "hellaswag", "squad", "boolq", "winogrande"],
         ["mnli", "squad", "boolq", "winogrande", "arc", "hellaswag"],
         ["mnli", "winogrande", "arc", "hellaswag", "squad", "boolq"]]


def test_latin_square_gives_disjoint_sets_early_and_identical_sets_at_the_end():
    orders = latin_square(3, 9)
    assert len(orders) == 3 and all(sorted(o) == list(range(9)) for o in orders)
    # at generation 3 each lineage has seen 3 families and they are disjoint
    seen3 = [set(o[:3]) for o in orders]
    assert seen3 == [{0, 1, 2}, {3, 4, 5}, {6, 7, 8}]
    assert set.intersection(*seen3) == set()
    # by the last generation everyone has seen everything
    assert all(set(o[:9]) == set(range(9)) for o in orders)


def test_latin_square_requires_even_division():
    with pytest.raises(ValueError):
        latin_square(4, 9)


def test_complementarity_is_one_when_disjoint_and_zero_when_identical():
    orders = latin_square(3, 9)
    assert complementarity(orders, 3) == pytest.approx(1.0)     # disjoint
    assert complementarity(orders, 9) == pytest.approx(0.0)     # identical
    mid = complementarity(orders, 6)
    assert 0.0 < mid < 1.0                                       # decays in between
    assert complementarity(orders, 0) == 0.0


def test_complementarity_decays_monotonically_after_its_peak():
    orders = latin_square(3, 9)
    vals = [complementarity(orders, t) for t in range(3, 10)]
    assert all(b <= a + 1e-9 for a, b in zip(vals, vals[1:]))


def test_replay_split_is_a_fixed_total_spread_over_families_seen():
    assert sum(replay_split([0, 1, 2], 150).values()) == 150
    assert sum(replay_split(list(range(7)), 150).values()) == 150   # remainder distributed
    # per-family protection thins as the curriculum grows — the point of a fixed total
    assert max(replay_split([0, 1], 150).values()) > max(replay_split(list(range(8)), 150).values())
    assert replay_split([], 150) == {} and replay_split([0], 0) == {}


def test_partner_selection_per_arm():
    assert partner_for("isolated", 0, 5, 3) is None
    assert partner_for("society", 0, 5, 3) == ("contemporary", 1)
    assert partner_for("society", 2, 5, 3) == ("contemporary", 0)      # wraps
    assert partner_for("society_dry", 1, 5, 3) == ("contemporary", 2)
    assert partner_for("seed_bank", 0, 5, 3) == ("ancestor", 2)        # t - depth
    assert partner_for("seed_bank", 0, 2, 3) is None                   # no ancestor yet
    assert partner_for("society", 0, 5, 1) is None                     # nobody to pair with


def test_merge_until_is_a_forced_stop_for_every_recombining_arm():
    # the fixed "merge early, then stop" schedule: partners exist only at t < merge_until
    assert partner_for("society", 0, 2, 3, merge_until=3) == ("contemporary", 1)
    assert partner_for("society", 0, 3, 3, merge_until=3) is None
    assert partner_for("seed_bank", 0, 4, 3, ancestor_depth=3, merge_until=3) is None
    assert partner_for("society", 0, 5, 3, merge_until=None) == ("contemporary", 1)  # default: unchanged
    assert partner_for("isolated", 0, 1, 3, merge_until=3) is None


def test_resolve_orders_defaults_to_the_latin_square_and_validates_custom_orders():
    assert resolve_orders({}, FAMS, 3, 6) == latin_square(3, 6)
    orders = resolve_orders({"orders": DECOR}, FAMS, 3, 6)
    assert orders[0] == [0, 1, 2, 3, 4, 5] and orders[1][1] == FAMS.index("squad")
    with pytest.raises(ValueError):
        resolve_orders({"orders": DECOR[:2]}, FAMS, 3, 6)                     # wrong lineage count
    with pytest.raises(ValueError):
        resolve_orders({"orders": [["mnli", "nope"] + FAMS[2:]] + DECOR[1:]}, FAMS, 3, 6)  # bad name
    with pytest.raises(ValueError):
        resolve_orders({"orders": [o[:5] for o in DECOR]}, FAMS, 3, 6)        # too short


def test_decorrelated_curriculum_has_non_monotone_complementarity():
    # Latin square: complementarity falls with generation, so it is collinear with adapter age
    ls = [complementarity(latin_square(3, 6), t) for t in range(1, 7)]
    assert ls == pytest.approx([1.0, 1.0, 0.8, 2 / 3, 1 / 3, 0.0], abs=1e-9)
    # shared-start curriculum: zero, rises to a peak at generation 3, then falls back to zero
    orders = resolve_orders({"orders": DECOR}, FAMS, 3, 6)
    dc = [complementarity(orders, t) for t in range(1, 7)]
    assert dc[0] == 0.0 and dc[-1] == 0.0
    assert max(dc) == pytest.approx(0.7, abs=1e-9) and dc.index(max(dc)) == 2
    assert dc[2] > dc[0] and dc[2] > dc[5]            # non-monotone: breaks the collinearity


def test_cumulative_accuracy_only_counts_families_the_model_should_know():
    acc = {"a": 0.9, "b": 0.5, "c": 0.1}
    assert cumulative_accuracy(acc, ["a", "b"]) == pytest.approx(0.7)
    assert cumulative_accuracy(acc, ["a"]) == pytest.approx(0.9)
    assert np.isnan(cumulative_accuracy(acc, []))
    assert cumulative_accuracy({"a": float("nan"), "b": 0.4}, ["a", "b"]) == pytest.approx(0.4)


def test_forgetting_is_positive_when_a_skill_decays():
    hist = [{"a": 0.9}, {"a": 0.7}, {"a": 0.4}]
    assert forgetting(hist, "a", 0, 2) == pytest.approx(0.5)
    assert forgetting(hist, "a", 0, 0) == pytest.approx(0.0)
    assert np.isnan(forgetting(hist, "b", 0, 2))


# ------------------------------------------------------------ manuscript revision 2026-09-12

import yaml
from pathlib import Path

from llm.curriculum import cull_step, inherit_slot
from llm.speciation import adapter_root

CONFIGS = Path(__file__).resolve().parents[1] / "configs" / "llm"
CONFLICT = {"boolq", "winogrande"}          # the two families whose answer conventions clash


def _orders(name):
    cfg = yaml.safe_load(open(CONFIGS / f"{name}.yaml"))
    return resolve_orders(cfg, FAMS, 3, 6), cfg


def test_cull_step_replaces_the_worst_slot_with_the_best_and_leaves_ties_alone():
    assert cull_step([0.5, 0.8, 0.3]) == (2, 1)           # (culled, source)
    assert cull_step([0.6, 0.6, 0.6]) is None             # a population of equals is not reshuffled
    assert cull_step([float("nan"), 0.4, 0.5]) == (0, 2)  # an unmeasured slot counts as the worst
    assert cull_step([0.4]) is None                        # N = 1: nothing to select between


def test_inherit_slot_aliases_the_genome_and_its_ancestry_but_not_the_slot():
    adapters = ["a0", "a1", "a2"]
    history = [["mnli"], ["arc", "squad"], ["boolq"]]
    budget = [300, 750, 300]
    archive = {(1, 0): "g0/lin1", (1, 1): "g1/lin1", (0, 0): "g0/lin0", (1, 5): "future"}
    inherit_slot(0, 1, adapters, history, budget, archive, t=1)
    assert adapters[0] == "a1" and budget[0] == 750
    assert history[0] == ["arc", "squad"] and history[0] is not history[1]   # a copy, not a view
    # the seed-bank ancestry follows the genome up to the current generation only
    assert archive[(0, 0)] == "g0/lin1" and archive[(0, 1)] == "g1/lin1" and (0, 5) not in archive
    assert adapters[1] == "a1" and budget[1] == 750                          # the source is untouched


def test_conflict_early_and_late_curricula_move_only_the_arrival_of_the_conflicting_pair():
    early, ecfg = _orders("curriculum_v5_early")
    late, lcfg = _orders("curriculum_v5_late")
    for orders in (early, late):
        assert all(sorted(o) == list(range(6)) for o in orders)              # every lineage sees all six
    names = lambda o: [FAMS[k] for k in o]
    assert all(set(names(o)[:2]) == CONFLICT for o in early)                 # conflict in generations 1-2
    assert all(set(names(o)[4:]) == CONFLICT for o in late)                  # conflict in generations 5-6
    # skill count rises identically (one new family per generation) so only conflict arrival differs
    assert all(len(set(o[:t])) == t for o in early + late for t in range(1, 7))
    # a contemporary partner still brings something new mid-curriculum in both designs
    assert complementarity(early, 3) > 0 and complementarity(late, 3) > 0
    # the obligate variants share the orders and only drop the veto
    for name, orders in (("early", early), ("late", late)):
        obl, ocfg = _orders(f"curriculum_v5_{name}_obl")
        assert obl == orders and ocfg["allow_veto"] is False and ocfg["arms"] == ["society"]
    assert ecfg["allow_veto"] is True and set(ecfg["arms"]) == {"isolated", "society"}


def test_cull_config_turns_selection_on_for_both_arms_of_the_latin_square():
    cfg = yaml.safe_load(open(CONFIGS / "curriculum_v5_cull.yaml"))
    assert cfg["cull"] is True and cfg["allow_veto"] is True and "orders" not in cfg
    assert set(cfg["arms"]) == {"isolated", "society"}


def test_speciation_adapter_root_is_seed_specific_so_array_elements_do_not_collide():
    a = adapter_root({"seed": 2})
    b = adapter_root({"seed": 3, "adapters_dir": "models/elsewhere"})
    assert a.name == "speciation_s2" and b == Path("models/elsewhere/speciation_s3")