# tests/test_mechanism_layer.py
import sys
import numpy as np
import pytest
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from llm_guard.features.mechanism_layer import MechanismFeatureExtractor

# ── Fixtures ─────────────────────────────────────────────────────────────────

@pytest.fixture(scope="module")
def extractor():
    return MechanismFeatureExtractor()

Q = "What year did the Berlin Wall fall?"
A_CORRECT = "The Berlin Wall fell in 1989, marking a pivotal moment in history."
A_HALLUC  = "The Berlin Wall fell in 1975 after a series of protests in East Germany."
A_EMPTY   = ""
CONTEXT   = "The Berlin Wall was a barrier that divided Berlin from 1961 to 1989."

def test_feature_shape(extractor):
    feat = extractor.extract_features(Q, A_CORRECT)
    assert feat.shape == (12,), f"Expected (12,), got {feat.shape}"
    assert feat.dtype in [np.float32, np.float64]

def test_feature_no_label_param(extractor):
    """Signature must not accept a label parameter."""
    import inspect
    sig = inspect.signature(extractor.extract_features)
    assert 'label' not in sig.parameters

def test_feature_with_context(extractor):
    feat_no_ctx = extractor.extract_features(Q, A_CORRECT)
    feat_ctx    = extractor.extract_features(Q, A_CORRECT, context=CONTEXT)
    assert feat_ctx[11] != 0.0, "f12 should be nonzero when context is provided"
    assert feat_no_ctx[11] == 0.0, "f12 should be 0.0 when context is empty"

def test_feature_empty_answer_no_crash(extractor):
    feat = extractor.extract_features(Q, A_EMPTY)
    assert feat.shape == (12,)
    assert not np.any(np.isnan(feat)), "No NaN values on empty answer"
    assert not np.any(np.isinf(feat)), "No Inf values on empty answer"

def test_feature_empty_question_no_crash(extractor):
    feat = extractor.extract_features("", A_CORRECT)
    assert feat.shape == (12,)
    assert not np.any(np.isnan(feat))

def test_f1_sem_sim_range(extractor):
    feat = extractor.extract_features(Q, A_CORRECT)
    assert 0.0 <= feat[0] <= 1.0, f"f1 out of range: {feat[0]}"

def test_f3_certainty_density_higher_for_certain(extractor):
    certain = extractor.extract_features(Q, "This is definitely certainly always true.")
    hedged  = extractor.extract_features(Q, "This might perhaps be possibly true.")
    assert certain[2] > hedged[2], "Certainty density should be higher for certain text"

def test_f4_hedge_density_higher_for_hedged(extractor):
    certain = extractor.extract_features(Q, "This is definitely certainly always true.")
    hedged  = extractor.extract_features(Q, "This might perhaps be possibly true.")
    assert hedged[3] > certain[3], "Hedge density should be higher for hedged text"

def test_f8_polarity_shift(extractor):
    neg_q    = "Is this treatment harmful and dangerous?"
    pos_ans  = "Yes, it is very good and beneficial."
    feat     = extractor.extract_features(neg_q, pos_ans)
    same_ans = "Yes, it is dangerous and harmful."
    feat2    = extractor.extract_features(neg_q, same_ans)
    assert feat[7] > feat2[7], "Polarity shift should be larger for opposite-polarity answer"

def test_f9_type_token_ratio(extractor):
    repetitive = extractor.extract_features(Q, "yes yes yes yes yes yes yes yes")
    varied     = extractor.extract_features(Q, A_CORRECT)
    assert repetitive[8] < varied[8], "Repetitive answer should have lower TTR"

def test_f12_zero_when_context_absent(extractor):
    feat = extractor.extract_features(Q, A_CORRECT, context="")
    assert feat[11] == 0.0

def test_batch_matches_single(extractor):
    records = [
        {"question": Q, "answer": A_CORRECT, "context": ""},
        {"question": Q, "answer": A_HALLUC,  "context": CONTEXT},
    ]
    batch = extractor.extract_features_batch(records)
    single0 = extractor.extract_features(Q, A_CORRECT, context="")
    single1 = extractor.extract_features(Q, A_HALLUC,  context=CONTEXT)
    np.testing.assert_allclose(batch[0], single0, rtol=1e-5)
    np.testing.assert_allclose(batch[1], single1, rtol=1e-5)


# ── Dataset Loader Tests ───────────────────────────────────────────────────────

from llm_guard.features.mechanism_layer import (
    load_halueval_qa,
    load_halueval_summarization,
    load_halueval_dialogue,
    load_truthfulqa,
    load_sycophancy,
)

BENCH_DIR = Path(__file__).resolve().parent.parent / "results" / "benchmark_datasets"

def test_load_halueval_qa_schema():
    path = BENCH_DIR / "hallucination_qa" / "halueval2_qa.jsonl"
    records = load_halueval_qa(path)
    assert len(records) > 0
    for r in records[:5]:
        assert "question" in r and "answer" in r and "label" in r and "context" in r
        assert r["label"] in (0, 1)
        assert isinstance(r["question"], str)

def test_load_halueval_qa_balance():
    path = BENCH_DIR / "hallucination_qa" / "halueval2_qa.jsonl"
    records = load_halueval_qa(path)
    labels = [r["label"] for r in records]
    pos_rate = sum(labels) / len(labels)
    assert 0.35 <= pos_rate <= 0.65, f"Expected ~50/50 balance, got {pos_rate:.2f}"

def test_load_halueval_summarization_uses_first_sentence_as_question():
    path = BENCH_DIR / "hallucination_qa" / "halueval2_summarization.jsonl"
    records = load_halueval_summarization(path)
    assert len(records) > 0
    for r in records[:5]:
        assert r["question"] != ""  # first sentence extracted
        assert r["context"] != ""   # document used as context

def test_load_halueval_summarization_label_correctness():
    path = BENCH_DIR / "hallucination_qa" / "halueval2_summarization.jsonl"
    records = load_halueval_summarization(path)
    for r in records[:5]:
        assert r["label"] in (0, 1)
    labels = [r["label"] for r in records]
    pos_rate = sum(labels) / len(labels)
    assert 0.35 <= pos_rate <= 0.65, f"Expected ~50/50 label balance, got {pos_rate:.2f}"

def test_load_halueval_dialogue_extracts_last_human_utterance():
    path = BENCH_DIR / "hallucination_qa" / "halueval2_dialogue.jsonl"
    records = load_halueval_dialogue(path)
    assert len(records) > 0
    for r in records[:5]:
        assert r["question"] != ""
        assert r["context"] != ""  # knowledge field

def test_load_truthfulqa_expands_mc_choices():
    path = BENCH_DIR / "confident_wrong" / "truthfulqa.jsonl"
    records = load_truthfulqa(path)
    assert len(records) > 817  # more than source questions
    for r in records[:5]:
        assert r["label"] in (0, 1)
    labels = [r["label"] for r in records]
    pos_rate = sum(labels) / len(labels)
    assert 0.60 <= pos_rate <= 0.95, f"Expected ~80% wrong (label=1), got {pos_rate:.2f}"

def test_load_sycophancy_extracts_full_choice_text():
    path = BENCH_DIR / "sycophancy" / "anthropic_sycophancy_eval.jsonl"
    records = load_sycophancy(path)
    assert len(records) > 0
    for r in records[:10]:
        assert len(r["answer"]) > 3, f"Answer too short: '{r['answer']}'"
        assert r["label"] in (0, 1)
        # question stem must not contain choice markers at (A) position
        assert "(A)" not in r["question"] or len(r["question"]) < 5

def test_load_sycophancy_shuffled_seed42():
    """Same seed=42 produces same shuffle both times."""
    path = BENCH_DIR / "sycophancy" / "anthropic_sycophancy_eval.jsonl"
    r1 = load_sycophancy(path, seed=42)
    r2 = load_sycophancy(path, seed=42)
    for a, b in zip(r1[:10], r2[:10]):
        assert a["label"] == b["label"]
        assert a["answer"] == b["answer"]


from llm_guard.features.mechanism_layer import bootstrap_ci, run_cv_experiment

# ── Bootstrap CI tests ────────────────────────────────────────────────────────

def test_bootstrap_ci_perfect_separation():
    labels = np.array([0]*50 + [1]*50)
    scores = np.array([0.0]*50 + [1.0]*50)
    result = bootstrap_ci(labels, scores)
    assert result["auroc"] == pytest.approx(1.0, abs=0.01)
    assert result["auroc_lo"] >= 0.95
    assert result["auprc"] == pytest.approx(1.0, abs=0.01)

def test_bootstrap_ci_random_classifier():
    rng = np.random.RandomState(42)
    labels = rng.randint(0, 2, 200)
    scores = rng.uniform(0, 1, 200)
    result = bootstrap_ci(labels, scores)
    assert 0.35 <= result["auroc"] <= 0.65
    assert result["auroc_lo"] < 0.55

def test_bootstrap_ci_returns_required_keys():
    labels = np.array([0]*50 + [1]*50)
    scores = np.array([0.0]*50 + [1.0]*50)
    result = bootstrap_ci(labels, scores)
    required = {"auroc", "auroc_lo", "auroc_hi", "auprc", "auprc_lo", "auprc_hi", "class_prior"}
    assert required.issubset(result.keys())

def test_bootstrap_ci_class_prior():
    labels = np.array([0]*80 + [1]*20)
    scores = np.random.uniform(0, 1, 100)
    result = bootstrap_ci(labels, scores)
    assert result["class_prior"] == pytest.approx(0.20, abs=0.01)

def test_bootstrap_ci_reproducible():
    labels = np.array([0]*50 + [1]*50)
    scores = np.random.uniform(0, 1, 100)
    r1 = bootstrap_ci(labels, scores, seed=42)
    r2 = bootstrap_ci(labels, scores, seed=42)
    assert r1["auroc"] == r2["auroc"]
    assert r1["auroc_lo"] == r2["auroc_lo"]

# ── CV pipeline tests ─────────────────────────────────────────────────────────

def test_run_cv_no_leakage():
    rng = np.random.RandomState(0)
    X = rng.randn(200, 12)
    y = np.array([0]*100 + [1]*100)
    result = run_cv_experiment(X, y)
    assert len(result["oof_scores"]) == 200
    assert "auroc" in result and "auprc" in result
    assert result["n_train"] + result["n_test"] == 200

def test_run_cv_gate_pass_for_perfect():
    rng = np.random.RandomState(0)
    X = np.column_stack([np.array([0.0]*100 + [1.0]*100)] + [rng.randn(200) for _ in range(11)])
    y = np.array([0]*100 + [1]*100)
    result = run_cv_experiment(X, y)
    assert result["gate_pass"] is True
    assert result["auroc"] >= 0.60
    assert result["auroc_lo"] > 0.50

def test_run_cv_gate_fail_for_random():
    """Random features should not pass gate."""
    rng = np.random.RandomState(42)
    X = rng.randn(400, 12)
    y = np.array([0]*200 + [1]*200)
    result = run_cv_experiment(X, y)
    assert result["gate_pass"] is False, (
        f"Random features should not pass gate — AUROC={result['auroc']:.3f}, "
        f"CI=[{result['auroc_lo']:.3f},{result['auroc_hi']:.3f}]"
    )

def test_run_cv_8020_path_oof_integrity():
    """80/20 path (N >= 800): OOF scores length matches test set size."""
    rng = np.random.RandomState(0)
    X = rng.randn(1000, 12)
    y = np.array([0]*500 + [1]*500)
    result = run_cv_experiment(X, y)
    assert result["n_test"] == 200  # 20% of 1000
    assert len(result["oof_scores"]) == 200
    assert len(result["oof_labels"]) == 200

def test_run_cv_5fold_path_oof_integrity():
    """5-fold CV path (N < 800): every sample scored exactly once."""
    rng = np.random.RandomState(0)
    X = rng.randn(400, 12)
    y = np.array([0]*200 + [1]*200)
    result = run_cv_experiment(X, y)
    assert len(result["oof_scores"]) == 400
    assert len(result["oof_labels"]) == 400
