"""Tests for ttc.py (Tasks 1-3). Run: python -m pytest -q test_ttc.py"""

import importlib
import itertools
import os

import numpy as np
import pytest

from chain import sample_answers, verifier_scores

ttc = importlib.import_module("solutions.ttc" if os.environ.get("RL_SOLUTIONS") else "ttc")


# ------------------------------------------------------------ Task 1
def test_pass_at_n_worked_example():
    assert ttc.pass_at_n(0.3, 4) == pytest.approx(0.7599), "1 - 0.7^4 = 0.7599 (the chapter's example)"
    assert ttc.pass_at_n(0.3, 1) == pytest.approx(0.3), "one sample: pass@1 = p"
    assert ttc.pass_at_n(0.0, 100) == pytest.approx(0.0), "a problem you never solve stays unsolved"


def test_pass_at_n_vectorized():
    out = ttc.pass_at_n(np.array([0.1, 0.5]), 2)
    assert np.allclose(out, [0.19, 0.75]), f"expected [0.19, 0.75], got {out}"


def test_pass_at_k_hand_cases():
    assert ttc.pass_at_k(4, 1, 2) == pytest.approx(0.5), "1 - C(3,2)/C(4,2) = 1 - 3/6"
    assert ttc.pass_at_k(10, 0, 3) == pytest.approx(0.0), "no correct samples -> 0"
    assert ttc.pass_at_k(5, 4, 2) == pytest.approx(1.0), "n - c < k: every pair contains a correct one"
    assert ttc.pass_at_k(10, 3, 1) == pytest.approx(0.3), "k = 1 is just the fraction correct, c / n"


def test_pass_at_k_matches_subset_enumeration():
    # The estimator is exactly the fraction of size-k subsets that contain a correct sample.
    n, c, k = 7, 2, 3
    labels = [1] * c + [0] * (n - c)
    subsets = list(itertools.combinations(range(n), k))
    frac = sum(any(labels[i] for i in s) for s in subsets) / len(subsets)
    assert ttc.pass_at_k(n, c, k) == pytest.approx(frac)


def test_pass_at_k_is_unbiased():
    # Average the estimator over many simulated batches of n = 10 samples.
    rng = np.random.default_rng(0)
    p, n, k = 0.2, 10, 4
    est = [ttc.pass_at_k(n, int(c), k) for c in rng.binomial(n, p, size=20000)]
    assert np.mean(est) == pytest.approx(1 - (1 - p) ** k, abs=0.01), "E[estimate] should equal the true pass@k"


# ------------------------------------------------------------ Task 2
def test_majority_vote_accuracy_worked_examples():
    assert ttc.majority_vote_accuracy(0.6, 3) == pytest.approx(0.648), "0.6^3 + 3(0.6^2)(0.4)"
    assert ttc.majority_vote_accuracy(0.4, 3) == pytest.approx(0.352), "0.4^3 + 3(0.4^2)(0.6)"
    assert ttc.majority_vote_accuracy(0.6, 1) == pytest.approx(0.6)


def test_majority_vote_accuracy_ties_and_symmetry():
    assert ttc.majority_vote_accuracy(0.6, 2) == pytest.approx(0.6), "with coin-flip ties, N = 2 is no better than N = 1"
    assert ttc.majority_vote_accuracy(0.6, 4) == pytest.approx(ttc.majority_vote_accuracy(0.6, 3)), \
        "with coin-flip ties, an even N adds nothing over N - 1"
    for n in [1, 2, 5, 10]:
        assert ttc.majority_vote_accuracy(0.5, n) == pytest.approx(0.5), "p = 1/2 stays a coin flip"


def test_majority_vote_amplifies_and_suppresses():
    up = [ttc.majority_vote_accuracy(0.6, n) for n in (1, 11, 101)]
    down = [ttc.majority_vote_accuracy(0.4, n) for n in (1, 11, 101)]
    assert up[0] < up[1] < up[2] and up[2] > 0.97, f"p > 1/2: accuracy should climb toward 1, got {up}"
    assert down[0] > down[1] > down[2] and down[2] < 0.03, f"p < 1/2: accuracy should fall toward 0, got {down}"


# ------------------------------------------------------------ Task 3
def test_majority_vote_selector():
    assert ttc.majority_vote([7, 3, 3, 5]) == 3
    assert ttc.majority_vote([7, 3, 3, 7, 5]) == 7, "tie between 7 and 3: the one appearing first wins"
    assert ttc.majority_vote(["b", "a"]) == "b", "all tied: the first answer wins"
    assert ttc.majority_vote(np.array([2, 1, 1])) == 1, "should also accept numpy arrays"


def test_best_of_n_selector():
    assert ttc.best_of_n([7, 3, 3], [0.2, 0.9, 0.1]) == 3
    assert ttc.best_of_n([7, 3, 5], [0.5, 0.5, 0.1]) == 7, "tie on the top score: earliest sample wins"
    assert ttc.best_of_n(np.array([4, 2]), np.array([0.1, 0.3])) == 2


def test_weighted_vote_selector():
    assert ttc.weighted_vote([7, 3, 3], [0.8, 0.5, 0.4]) == 3, "3 totals 0.9 > 0.8"
    assert ttc.weighted_vote([7, 3, 3], [0.95, 0.5, 0.4]) == 7, "7 totals 0.95 > 0.9"
    assert ttc.weighted_vote([5, 6], [0.5, 0.5]) == 5, "tie: the answer appearing first wins"
    # With equal scores it is exactly majority vote.
    ans = [1, 2, 2, 3, 1, 2]
    assert ttc.weighted_vote(ans, [0.7] * 6) == ttc.majority_vote(ans) == 2


def test_majority_vote_matches_formula_on_binary_answers():
    # wrong_top_share = 1: every wrong sample gives the same wrong answer -> binary answers.
    rng = np.random.default_rng(1)
    p, n, trials = 0.6, 5, 4000
    hits = sum(ttc.majority_vote(list(sample_answers(p, n, 1.0, rng))) == 0 for _ in range(trials))
    assert hits / trials == pytest.approx(ttc.majority_vote_accuracy(p, n), abs=0.025)


def _accuracy(rule, p, w, d, f, n, trials=400, seed=0):
    rng = np.random.default_rng(seed)
    hits = 0
    for _ in range(trials):
        a = sample_answers(p, n, w, rng)
        s = verifier_scores(a, d, f, rng)
        hits += (rule(list(a)) if rule is ttc.majority_vote else rule(list(a), s)) == 0
    return hits / trials


def test_scaling_story():
    # p = 0.3, scattered wrong answers, a decent verifier that is fooled 5% of the time.
    p, w, d, f = 0.3, 0.3, 2.0, 0.05
    bon16 = _accuracy(ttc.best_of_n, p, w, d, f, 16)
    bon512 = _accuracy(ttc.best_of_n, p, w, d, f, 512)
    wv64 = _accuracy(ttc.weighted_vote, p, w, d, f, 64)
    maj64 = _accuracy(ttc.majority_vote, p, w, d, f, 64)
    assert bon16 > 0.6, f"best-of-16 should beat a single sample (0.3) by a lot, got {bon16:.2f}"
    assert bon512 < bon16 - 0.15, \
        f"with convincing wrong answers, best-of-N should DECLINE at large N: N=16 {bon16:.2f}, N=512 {bon512:.2f}"
    assert wv64 > 0.95, f"weighted vote should shrug off a few convincing wrong answers, got {wv64:.2f}"
    assert 0.7 < maj64 < 0.97, f"majority vote at N=64 should be climbing toward 1, got {maj64:.2f}"