import numpy as np
import pytest
import warnings

import promptstats as ps
from promptstats.compare import CompareReport, EntityStats


def _rng(seed: int = 0) -> np.random.Generator:
    return np.random.default_rng(seed)


def test_compare_models_requires_dict():
    with pytest.raises(TypeError, match="dict"):
        ps.compare_models([np.ones((2, 10)), np.ones((2, 10))])


def test_compare_models_requires_at_least_two_models():
    with pytest.raises(ValueError, match="at least 2"):
        ps.compare_models({"only_one": np.ones((2, 10))})


def test_compare_models_flat_3d_raises():
    with pytest.raises(ValueError, match=r"Expected 1-D \(M inputs\) or 2-D \(M inputs, R runs\)"):
        ps.compare_models(
            {
                "m1": np.ones((2, 10, 3)),
                "m2": np.ones((2, 10, 3)),
            },
            rng=_rng(),
        )


def test_compare_models_mismatched_template_counts_raises():
    with pytest.raises(ValueError, match="same template keys"):
        ps.compare_models(
            {
                "m1": {"t1": np.ones(10), "t2": np.ones(10)},
                "m2": {"t1": np.ones(10), "t2": np.ones(10), "t3": np.ones(10)},
            },
            rng=_rng(),
        )


def test_compare_models_mismatched_run_counts_raises():
    with pytest.raises(ValueError, match="same number of runs"):
        ps.compare_models(
            {
                "m1": {
                    "t1": np.ones((10, 3)),
                    "t2": np.ones((10, 3)),
                },
                "m2": {
                    "t1": np.ones((10, 2)),
                    "t2": np.ones((10, 2)),
                },
            },
            rng=_rng(),
        )


def test_compare_models_template_labels_length_check():
    with pytest.raises(ValueError, match="template_labels length"):
        ps.compare_models(
            {
                "m1": {"t1": np.ones(10), "t2": np.ones(10)},
                "m2": {"t1": np.ones(10), "t2": np.ones(10)},
            },
            template_labels=["only_one"],
            rng=_rng(),
        )


@pytest.mark.parametrize("method", ["wilson", "newcombe", "fisher_exact"])
def test_compare_models_accepts_explicit_binary_methods(method: str):
    report = ps.compare_models(
        {
            "m1": np.array([1, 0, 1, 1, 0, 1, 0, 1], dtype=float),
            "m2": np.array([0, 0, 1, 0, 0, 1, 0, 0], dtype=float),
            "m3": np.array([1, 1, 1, 1, 0, 1, 1, 1], dtype=float),
        },
        method=method,
        n_bootstrap=300,
        rng=_rng(31),
    )
    pair = report.pairwise.get("m1", "m2")
    if method == "fisher_exact":
        assert "fisher exact" in pair.test_method.lower()
    else:
        assert "newcombe" in pair.test_method


@pytest.mark.parametrize("method", ["wilson", "newcombe", "fisher_exact"])
def test_compare_models_explicit_binary_methods_reject_non_binary(method: str):
    with pytest.raises(ValueError, match=r"requires binary \(0/1\) data"):
        ps.compare_models(
            {
                "m1": np.array([0.1, 0.4, 0.8, 0.6, 0.2], dtype=float),
                "m2": np.array([0.2, 0.5, 0.7, 0.4, 0.3], dtype=float),
            },
            method=method,
            rng=_rng(32),
        )


def test_compare_models_accepts_sign_test_for_non_binary_data():
    report = ps.compare_models(
        {
            "m1": np.array([0.70, 0.72, 0.69, 0.71, 0.70, 0.73, 0.72, 0.71], dtype=float),
            "m2": np.array([0.64, 0.66, 0.65, 0.67, 0.64, 0.68, 0.66, 0.65], dtype=float),
            "m3": np.array([0.76, 0.77, 0.75, 0.78, 0.76, 0.79, 0.77, 0.78], dtype=float),
        },
        method="sign_test",
        n_bootstrap=500,
        rng=_rng(33),
    )

    pair = report.pairwise.get("m1", "m2")
    assert "sign test" in pair.test_method.lower()


def test_compare_models_returns_report_and_fields():
    report = ps.compare_models(
        {
            "model_a": {
                "t1": [0.80, 0.82, 0.78, 0.81, 0.79],
                "t2": [0.77, 0.79, 0.76, 0.78, 0.75],
            },
            "model_b": {
                "t1": [0.86, 0.88, 0.84, 0.87, 0.85],
                "t2": [0.83, 0.84, 0.82, 0.85, 0.81],
            },
        },
        n_bootstrap=400,
        rng=_rng(5),
    )

    assert isinstance(report, CompareReport)
    assert set(report.labels) == {"model_a", "model_b"}
    assert set(report.model_stats.keys()) == {"model_a", "model_b"}
    assert isinstance(report.model_stats["model_a"], EntityStats)
    p = report.pairwise.get("model_a", "model_b").p_value
    assert 0.0 <= p <= 1.0
    assert isinstance(report.full_analysis, ps.MultiModelBundle)
    assert isinstance(report.quick_summary(), str)
    assert report.winner in (None, "model_a", "model_b")


def test_compare_models_detects_clear_winner():
    rng = _rng(11)
    m1 = rng.normal(loc=0.55, scale=0.05, size=120)
    m2 = rng.normal(loc=0.72, scale=0.05, size=120)

    report = ps.compare_models(
        {"weaker": m1, "stronger": m2},
        n_bootstrap=1200,
        rng=_rng(12),
    )

    assert report.unbeaten == ["stronger"]
    assert report.significant is True
    assert report.pairwise.get("weaker", "stronger").p_value < 0.05


def test_compare_models_pairwise_lookup_direction_agnostic():
    report = ps.compare_models(
        {
            "m1": {
                "t1": [0.7, 0.72, 0.71, 0.69, 0.70],
                "t2": [0.65, 0.66, 0.64, 0.67, 0.66],
            },
            "m2": {
                "t1": [0.75, 0.76, 0.74, 0.77, 0.75],
                "t2": [0.70, 0.71, 0.72, 0.69, 0.70],
            },
        },
        n_bootstrap=400,
        rng=_rng(2),
    )

    p12 = report.pairwise.get("m1", "m2").p_value
    p21 = report.pairwise.get("m2", "m1").p_value
    assert p12 == pytest.approx(p21, abs=1e-12)


def test_compare_models_accepts_flat_1d_per_model_scores():
    report = ps.compare_models(
        {
            "m1": np.array([0.70, 0.72, 0.69, 0.71, 0.70]),
            "m2": np.array([0.76, 0.77, 0.75, 0.78, 0.76]),
        },
        n_bootstrap=300,
        rng=_rng(21),
    )

    assert isinstance(report, CompareReport)
    assert set(report.labels) == {"m1", "m2"}


def test_compare_models_accepts_flat_nested_arrays_for_runs():
    report = ps.compare_models(
        {
            "m1": np.array([
                [0.70, 0.71, 0.69],
                [0.72, 0.73, 0.71],
                [0.68, 0.69, 0.67],
                [0.71, 0.72, 0.70],
            ]),
            "m2": np.array([
                [0.76, 0.77, 0.75],
                [0.78, 0.79, 0.77],
                [0.74, 0.75, 0.73],
                [0.77, 0.78, 0.76],
            ]),
        },
        template_labels=["single_template"],
        n_bootstrap=300,
        rng=_rng(22),
    )

    assert isinstance(report, CompareReport)
    assert report.full_analysis.benchmark.template_labels == ["single_template"]


def test_compare_models_multirun_uses_nested_bootstrap_for_pairwise():
    rng = _rng(123)
    n_inputs = 40
    n_runs = 5

    m1 = rng.normal(loc=0.70, scale=0.10, size=(n_inputs, n_runs))
    m2 = rng.normal(loc=0.62, scale=0.10, size=(n_inputs, n_runs))

    report = ps.compare_models(
        {
            "m1": m1,
            "m2": m2,
        },
        method="auto",
        n_bootstrap=500,
        rng=_rng(124),
    )

    pair = report.pairwise.get("m1", "m2")
    assert pair.n_runs == n_runs
    assert "nested" in pair.test_method.lower()


def test_compare_models_accepts_nested_template_dicts():
    report = ps.compare_models(
        {
            "m1": {
                "t_a": [0.70, 0.71, 0.69, 0.72],
                "t_b": [0.68, 0.69, 0.67, 0.70],
            },
            "m2": {
                "t_a": [0.76, 0.77, 0.75, 0.78],
                "t_b": [0.74, 0.75, 0.73, 0.76],
            },
        },
        n_bootstrap=300,
        rng=_rng(23),
    )

    assert isinstance(report, CompareReport)
    assert report.full_analysis.benchmark.template_labels == ["t_a", "t_b"]


def test_compare_models_accepts_nested_template_dicts_with_runs():
    report = ps.compare_models(
        {
            "m1": {
                "t_a": [[0.70, 0.71], [0.72, 0.73], [0.69, 0.70]],
                "t_b": [[0.68, 0.69], [0.70, 0.71], [0.67, 0.68]],
            },
            "m2": {
                "t_a": [[0.76, 0.77], [0.78, 0.79], [0.75, 0.76]],
                "t_b": [[0.74, 0.75], [0.76, 0.77], [0.73, 0.74]],
            },
        },
        n_bootstrap=300,
        rng=_rng(24),
    )

    assert isinstance(report, CompareReport)
    assert report.full_analysis.benchmark.template_labels == ["t_a", "t_b"]


def test_compare_models_nested_template_dict_keys_must_match():
    with pytest.raises(ValueError, match="same template keys"):
        ps.compare_models(
            {
                "m1": {
                    "t_a": [0.70, 0.71, 0.69, 0.72],
                    "t_b": [0.68, 0.69, 0.67, 0.70],
                },
                "m2": {
                    "t_a": [0.76, 0.77, 0.75, 0.78],
                    "t_c": [0.74, 0.75, 0.73, 0.76],
                },
            },
            rng=_rng(25),
        )


def test_compare_models_auto_collapse_avoids_pseudo_run_warning_single_template():
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        report = ps.compare_models(
            {
                "m1": np.array([0.70, 0.72, 0.69, 0.71, 0.70]),
                "m2": np.array([0.76, 0.77, 0.75, 0.78, 0.76]),
            },
            n_bootstrap=300,
            rng=_rng(26),
        )

    assert isinstance(report, CompareReport)
    two_run_msgs = [
        str(w.message) for w in caught if "only 2 runs detected" in str(w.message)
    ]
    assert not two_run_msgs
    assert report.full_analysis.template_level.benchmark.n_runs == 1


def test_compare_models_explicit_as_runs_keeps_run_like_model_axis():
    report = ps.compare_models(
        {
            "m1": np.array([0.70, 0.72, 0.69, 0.71, 0.70]),
            "m2": np.array([0.76, 0.77, 0.75, 0.78, 0.76]),
        },
        template_model_collapse="as_runs",
        n_bootstrap=300,
        rng=_rng(27),
    )

    assert isinstance(report, CompareReport)
    assert report.full_analysis.template_level.benchmark.n_runs == 2


def test_compare_models_rejects_unknown_template_model_collapse():
    with pytest.raises(ValueError, match="Unknown template_model_collapse"):
        ps.compare_models(
            {
                "m1": np.array([0.70, 0.72, 0.69, 0.71, 0.70]),
                "m2": np.array([0.76, 0.77, 0.75, 0.78, 0.76]),
            },
            template_model_collapse="median",  # type: ignore[arg-type]
            rng=_rng(28),
        )
