# tests/test_density_matrix_scorer.py
import numpy as np
import pytest
from llm_guard.verification.density_matrix_scorer import (
    PatternGraph, build_pattern_graph,
    _encode_amplitude, _build_density_matrix, _compute_step_features,
    DensityMatrixScorer,
)

MINIMAL_STEPS = [
    {"thought": "search for answer", "action_type": "Search",
     "action_arg": "who invented telephone", "observation": "Alexander Graham Bell invented the telephone in 1876"},
    {"thought": "found the answer", "action_type": "Finish",
     "action_arg": "Alexander Graham Bell", "observation": ""},
]

def test_pattern_graph_has_nodes():
    g = build_pattern_graph(MINIMAL_STEPS, embedder=None)
    assert len(g.nodes) > 0

def test_pattern_graph_has_sequential_edges():
    g = build_pattern_graph(MINIMAL_STEPS, embedder=None)
    assert len(g.edges) > 0

def test_spectral_gap_non_negative():
    g = build_pattern_graph(MINIMAL_STEPS, embedder=None)
    assert g.spectral_gap() >= 0.0

def test_spectral_gap_single_node_is_zero():
    g = PatternGraph(nodes=[{"id": "t0", "text": "hi"}], edges=[])
    assert g.spectral_gap() == 0.0

def test_pagerank_sums_to_one():
    g = build_pattern_graph(MINIMAL_STEPS, embedder=None)
    pr = g.pagerank()
    assert abs(pr.sum() - 1.0) < 1e-6

def test_adjacency_symmetric():
    g = build_pattern_graph(MINIMAL_STEPS, embedder=None)
    A = g.adjacency()
    np.testing.assert_allclose(A, A.T, atol=1e-10)


def test_fit_raises_on_missing_correct_key():
    scorer = DensityMatrixScorer()
    chains = _make_labeled_chains(n_correct=10, n_wrong=10)
    # Insert one chain without "correct" key into the middle
    keyless = {"question": "q", "steps": MINIMAL_STEPS, "final_answer": "a"}
    chains.insert(5, keyless)  # insert at position 5 to hit the per-item guard
    with pytest.raises(ValueError, match="'correct'"):
        scorer.fit(chains)


# ── Task 2 tests: amplitude encoding helpers and _build_density_matrix public API ──

def test_amplitude_normalized():
    phi = np.array([1.0, 2.0, 3.0, 0.0])
    alpha = _encode_amplitude(phi)
    assert abs(np.dot(alpha, alpha) - 1.0) < 1e-10

def test_amplitude_zero_input():
    phi = np.zeros(4)
    alpha = _encode_amplitude(phi)
    assert alpha.shape == (4,)

def test_density_matrix_trace_one():
    alpha = np.array([0.5, 0.5, 0.5, 0.5])  # already normalized
    rho = _build_density_matrix([alpha], np.array([1.0]))
    assert abs(np.trace(rho) - 1.0) < 1e-10

def test_density_matrix_symmetric():
    alphas = [_encode_amplitude(np.random.rand(6)) for _ in range(3)]
    weights = np.array([0.4, 0.3, 0.3])
    rho = _build_density_matrix(alphas, weights)
    np.testing.assert_allclose(rho, rho.T, atol=1e-10)

def test_density_matrix_psd():
    alphas = [_encode_amplitude(np.random.rand(6)) for _ in range(3)]
    weights = np.array([0.4, 0.3, 0.3])
    rho = _build_density_matrix(alphas, weights)
    eigs = np.linalg.eigvalsh(rho)
    assert np.all(eigs >= -1e-10)

def test_step_features_shape():
    g = build_pattern_graph(MINIMAL_STEPS, embedder=None)
    phi = _compute_step_features(MINIMAL_STEPS[0], MINIMAL_STEPS, g, step_idx=0)
    assert phi.ndim == 1
    assert len(phi) >= 6
    assert np.all(np.isfinite(phi))


# ── Task 3 tests: fit() with min-chain validation ──

def _make_labeled_chains(n_correct=15, n_wrong=12):
    chains = []
    for _ in range(n_correct):
        chains.append({
            "question": "Who invented the telephone?",
            "steps": MINIMAL_STEPS,
            "final_answer": "Alexander Graham Bell",
            "correct": True,
        })
    for _ in range(n_wrong):
        chains.append({
            "question": "Who invented the telephone?",
            "steps": [
                {"thought": "not sure", "action_type": "Search",
                 "action_arg": "telephone inventor", "observation": "unclear"},
                {"thought": "guessing", "action_type": "Finish",
                 "action_arg": "Edison", "observation": ""},
            ],
            "final_answer": "Edison",
            "correct": False,
        })
    return chains

def test_fit_requires_min_correct():
    scorer = DensityMatrixScorer()
    chains = _make_labeled_chains(n_correct=9, n_wrong=15)
    with pytest.raises(ValueError, match="got 9 correct"):
        scorer.fit(chains)

def test_fit_requires_min_wrong():
    scorer = DensityMatrixScorer()
    chains = _make_labeled_chains(n_correct=15, n_wrong=9)
    with pytest.raises(ValueError, match="got 15 correct and 9 wrong"):
        scorer.fit(chains)

def test_fit_succeeds_with_min_chains():
    scorer = DensityMatrixScorer()
    chains = _make_labeled_chains(n_correct=10, n_wrong=10)
    scorer.fit(chains)
    assert scorer._m_correct is not None
    assert scorer._m_incorrect is not None

def test_measurement_operators_are_psd():
    scorer = DensityMatrixScorer()
    scorer.fit(_make_labeled_chains())
    for M in [scorer._m_correct, scorer._m_incorrect]:
        eigs = np.linalg.eigvalsh(M)
        assert np.all(eigs >= -1e-8)

def test_measurement_operators_same_shape():
    scorer = DensityMatrixScorer()
    scorer.fit(_make_labeled_chains())
    assert scorer._m_correct.shape == scorer._m_incorrect.shape


# ── Task 4 tests: score() with RuntimeError guard, mode, density_matrix ──

def test_score_before_fit_raises():
    scorer = DensityMatrixScorer()
    with pytest.raises(RuntimeError, match="must be fitted"):
        scorer.score("q", MINIMAL_STEPS, "a")

def test_risk_score_in_range():
    scorer = DensityMatrixScorer()
    scorer.fit(_make_labeled_chains())
    result = scorer.score("Who invented the telephone?", MINIMAL_STEPS, "Bell")
    assert 0.0 <= result.risk_score <= 1.0

def test_sc_old_risk_always_present():
    scorer = DensityMatrixScorer()
    scorer.fit(_make_labeled_chains())
    result = scorer.score("Who invented the telephone?", MINIMAL_STEPS, "Bell")
    assert hasattr(result, "sc_old_risk")
    assert 0.0 <= result.sc_old_risk <= 1.0

def test_density_matrix_none_by_default():
    scorer = DensityMatrixScorer(return_density_matrix=False)
    scorer.fit(_make_labeled_chains())
    result = scorer.score("q", MINIMAL_STEPS, "a")
    assert result.density_matrix is None

def test_density_matrix_returned_when_requested():
    scorer = DensityMatrixScorer(return_density_matrix=True)
    scorer.fit(_make_labeled_chains())
    result = scorer.score("q", MINIMAL_STEPS, "a")
    assert isinstance(result.density_matrix, np.ndarray)
    assert result.density_matrix.ndim == 2

def test_mode_is_jaccard_without_embedder():
    scorer = DensityMatrixScorer(embedder=None)
    scorer.fit(_make_labeled_chains())
    result = scorer.score("q", MINIMAL_STEPS, "a")
    assert result.mode == "jaccard"

def test_graph_spectral_gap_non_negative():
    scorer = DensityMatrixScorer()
    scorer.fit(_make_labeled_chains())
    result = scorer.score("q", MINIMAL_STEPS, "a")
    assert result.graph_spectral_gap >= 0.0

def test_result_fields_always_present():
    scorer = DensityMatrixScorer()
    scorer.fit(_make_labeled_chains())
    result = scorer.score("q", MINIMAL_STEPS, "a")
    for f in ["risk_score", "sc_old_risk", "graph_spectral_gap", "mode", "density_matrix"]:
        assert hasattr(result, f), f"Missing field: {f}"
