"""
Tests for experiments/shared/stats_utils.py

Run: pytest tests/test_stats_utils.py -v
"""
import sys
from pathlib import Path
import numpy as np
import pytest

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from experiments.shared.stats_utils import (
    bootstrap_ci,
    auroc_with_ci,
    auprc_with_ci,
    precision_at_fpr,
    bonferroni_threshold,
    clopper_pearson_ci,
    eval_report,
)
from sklearn.metrics import roc_auc_score


# ------------------------------------------------------------------ #
# bootstrap_ci                                                        #
# ------------------------------------------------------------------ #

def test_perfect_classifier_auroc_is_one():
    """Perfect classifier should give AUROC=1.0."""
    y_true = np.array([1, 1, 1, 0, 0, 0])
    y_score = np.array([0.9, 0.8, 0.7, 0.3, 0.2, 0.1])
    auroc, lo, hi = auroc_with_ci(y_true, y_score)
    assert auroc == pytest.approx(1.0, abs=1e-6)
    assert lo > 0.9


def test_random_classifier_auroc_near_half():
    """Random classifier should have AUROC ≈ 0.5, CI should span 0.5."""
    rng = np.random.RandomState(42)
    n = 200
    y_true = rng.randint(0, 2, n)
    y_score = rng.rand(n)
    auroc, lo, hi = auroc_with_ci(y_true, y_score)
    assert 0.35 < auroc < 0.65, f"Expected ~0.5, got {auroc}"
    assert lo < 0.55  # CI spans 0.5


def test_ci_width_shrinks_with_more_samples():
    """CI width should be smaller with more samples."""
    rng = np.random.RandomState(0)
    # Weak classifier on small vs large dataset
    y_true_small = np.array([1]*25 + [0]*25)
    y_true_large = np.array([1]*100 + [0]*100)
    y_score_small = np.concatenate([rng.uniform(0.4, 0.8, 25), rng.uniform(0.2, 0.6, 25)])
    y_score_large = np.concatenate([rng.uniform(0.4, 0.8, 100), rng.uniform(0.2, 0.6, 100)])

    _, lo_s, hi_s = auroc_with_ci(y_true_small, y_score_small, n_bootstrap=500)
    _, lo_l, hi_l = auroc_with_ci(y_true_large, y_score_large, n_bootstrap=500)

    width_small = hi_s - lo_s
    width_large = hi_l - lo_l
    assert width_large < width_small, f"Large CI {width_large:.3f} should be < small CI {width_small:.3f}"


def test_bootstrap_ci_uses_provided_metric():
    """bootstrap_ci should work with any metric_fn."""
    y_true = np.array([1, 1, 0, 0, 1, 0])
    y_score = np.array([0.9, 0.8, 0.3, 0.2, 0.7, 0.4])
    auroc, lo, hi = bootstrap_ci(y_true, y_score, roc_auc_score, n_bootstrap=100)
    assert 0.0 <= lo <= auroc <= hi <= 1.0


def test_bootstrap_ci_lo_le_point_le_hi():
    """CI bounds should bracket the point estimate."""
    rng = np.random.RandomState(7)
    y_true = rng.randint(0, 2, 80)
    y_score = rng.rand(80)
    point, lo, hi = auroc_with_ci(y_true, y_score, n_bootstrap=500)
    assert lo <= point <= hi


# ------------------------------------------------------------------ #
# bonferroni_threshold                                                #
# ------------------------------------------------------------------ #

def test_bonferroni_k1_unchanged():
    assert bonferroni_threshold(0.05, 1) == pytest.approx(0.05)


def test_bonferroni_k4():
    assert bonferroni_threshold(0.05, 4) == pytest.approx(0.0125)


def test_bonferroni_k3_architectures():
    """Three architectures compared → threshold = 0.05/3."""
    assert bonferroni_threshold(0.05, 3) == pytest.approx(0.05 / 3)


# ------------------------------------------------------------------ #
# clopper_pearson_ci                                                  #
# ------------------------------------------------------------------ #

def test_clopper_pearson_90_of_100():
    """90/100 successes: CI should bracket 0.90."""
    lo, hi = clopper_pearson_ci(90, 100)
    assert lo < 0.90 < hi
    # Standard result: ~[0.823, 0.957]
    assert lo > 0.80
    assert hi < 0.97


def test_clopper_pearson_perfect():
    """100/100: hi should be 1.0, lo should be high."""
    lo, hi = clopper_pearson_ci(100, 100)
    assert hi == pytest.approx(1.0, abs=1e-3)
    assert lo > 0.95


def test_clopper_pearson_zero():
    """0 successes: lo should be 0."""
    lo, hi = clopper_pearson_ci(0, 100)
    assert lo == pytest.approx(0.0, abs=1e-6)
    assert hi < 0.05


def test_clopper_pearson_zero_trials():
    """0 trials: should return (0, 1) without error."""
    lo, hi = clopper_pearson_ci(0, 0)
    assert lo == 0.0
    assert hi == 1.0


# ------------------------------------------------------------------ #
# precision_at_fpr                                                    #
# ------------------------------------------------------------------ #

def test_precision_at_fpr_perfect_signal():
    """Perfect signal: precision should be 1.0 at any FPR."""
    y_true = np.array([1, 1, 1, 0, 0, 0])
    y_score = np.array([0.9, 0.8, 0.7, 0.3, 0.2, 0.1])
    prec, threshold = precision_at_fpr(y_true, y_score, target_fpr=0.10)
    assert prec == pytest.approx(1.0, abs=1e-6)


def test_precision_at_fpr_returns_threshold():
    """Should return a valid threshold in [0, 1]."""
    rng = np.random.RandomState(1)
    y_true = rng.randint(0, 2, 50)
    y_score = rng.rand(50)
    prec, threshold = precision_at_fpr(y_true, y_score, target_fpr=0.10)
    assert 0.0 <= threshold <= 1.0
    assert 0.0 <= prec <= 1.0


# ------------------------------------------------------------------ #
# eval_report                                                         #
# ------------------------------------------------------------------ #

def test_eval_report_structure():
    """eval_report should return dict with all required keys."""
    y_true = np.array([1]*50 + [0]*50)
    y_score = np.concatenate([np.random.rand(50) * 0.4 + 0.6, np.random.rand(50) * 0.4])
    report = eval_report(y_true, y_score, domain="hotpotqa", experiment_id="A1", architecture="P")
    required_keys = [
        "experiment_id", "architecture", "domain", "n", "n_incorrect",
        "auroc", "auroc_lo", "auroc_hi",
        "auprc", "auprc_lo", "auprc_hi",
        "precision_at_fpr10", "beats_random", "passes_gate",
    ]
    for k in required_keys:
        assert k in report, f"Missing key: {k}"


def test_eval_report_passes_gate_when_lo_above_threshold():
    """passes_gate should be True when auroc_lo > gate_threshold."""
    y_true = np.array([1]*100 + [0]*100)
    # Strong signal
    y_score = np.concatenate([np.ones(100) * 0.9, np.ones(100) * 0.1])
    report = eval_report(y_true, y_score, domain="hotpotqa",
                         experiment_id="test", architecture="P",
                         gate_threshold=0.70)
    assert report["passes_gate"] is True
    assert report["auroc"] > 0.98


def test_eval_report_fails_gate_when_lo_below_threshold():
    """passes_gate should be False when auroc_lo ≤ gate_threshold."""
    rng = np.random.RandomState(42)
    y_true = rng.randint(0, 2, 60)
    y_score = rng.rand(60)
    report = eval_report(y_true, y_score, domain="triviaqa",
                         experiment_id="test", architecture="FB",
                         gate_threshold=0.90)
    assert report["passes_gate"] is False


def test_eval_report_no_gate():
    """passes_gate should be None when gate_threshold not provided."""
    y_true = np.array([1, 0, 1, 0])
    y_score = np.array([0.8, 0.2, 0.7, 0.3])
    report = eval_report(y_true, y_score, domain="nq",
                         experiment_id="test", architecture="P")
    assert report["passes_gate"] is None
    assert report["gate_threshold"] is None


# ------------------------------------------------------------------ #
# chain_schema                                                        #
# ------------------------------------------------------------------ #

def test_chain_record_roundtrip():
    """ChainRecord should serialize and deserialize cleanly."""
    from experiments.shared.chain_schema import ChainRecord, StepRecord, make_chain_id
    steps = [
        StepRecord(thought="Looking for X", action="Search[X]", observation="Found X."),
        StepRecord(thought="Got it", action="Finish[X is the answer]", observation=""),
    ]
    chain = ChainRecord(
        question="What is X?",
        steps=steps,
        final_answer="X is the answer",
        label=1,
        domain="hotpotqa",
        model="haiku",
        chain_id=make_chain_id("What is X?", "hotpotqa", "haiku"),
        metadata={"gold": "X"},
    )
    jsonl = chain.to_jsonl()
    chain2 = ChainRecord.from_jsonl(jsonl)
    assert chain2.question == chain.question
    assert chain2.label == chain.label
    assert chain2.domain == chain.domain
    assert len(chain2.steps) == 2
    assert chain2.steps[0].thought == "Looking for X"
    assert chain2.steps[1].action == "Finish[X is the answer]"
    assert chain2.metadata["gold"] == "X"


def test_chain_id_is_stable():
    """Same inputs should always produce the same chain_id."""
    from experiments.shared.chain_schema import make_chain_id
    id1 = make_chain_id("Who was president?", "nq", "haiku")
    id2 = make_chain_id("Who was president?", "nq", "haiku")
    assert id1 == id2


def test_chain_id_differs_by_domain():
    from experiments.shared.chain_schema import make_chain_id
    id_nq = make_chain_id("Q", "nq", "haiku")
    id_hp = make_chain_id("Q", "hotpotqa", "haiku")
    assert id_nq != id_hp
