"""Split conformal: empirical coverage matches the finite-sample guarantee."""

import numpy as np
import pytest
from scipy.special import softmax

from clfkit.conformal import SplitConformal, conformal_quantile, evaluate_sets

K = 5


def _draw(n, rng):
    """y ~ softmax(true logits); the 'model' sees a distorted, misspecified version."""
    logits = rng.normal(0, 2.0, size=(n, K)) + np.array([0.8, 0.4, 0.0, -0.4, -0.8])
    y = np.array([rng.choice(K, p=p) for p in softmax(logits, axis=1)])
    model_p = softmax(1.7 * logits + rng.normal(0, 1.0, size=(n, K)), axis=1)  # overconfident + noisy
    return model_p, y


@pytest.mark.parametrize("score", ["lac", "aps"])
@pytest.mark.parametrize("mondrian", [False, True])
def test_marginal_coverage_across_seeds(score, mondrian):
    alpha, n_cal, n_test = 0.1, 500, 2000
    covs = []
    for seed in range(12):
        rng = np.random.default_rng(seed)
        Pc, yc = _draw(n_cal, rng)
        Pt, yt = _draw(n_test, rng)
        cp = SplitConformal(alpha, score, mondrian, seed=seed).fit(Pc, yc)
        covs.append(evaluate_sets(cp.predict_sets(Pt), yt)["coverage"])
    covs = np.array(covs)
    # Guarantee: E[coverage] in [1-alpha, 1-alpha+1/(n+1)] (Mondrian: per class, so
    # marginal is >= 1-alpha too). Per-seed coverage fluctuates (Beta-distributed
    # around the target with sd ~ sqrt(a(1-a)/n_cal) ~ 0.013 + test noise).
    assert covs.mean() >= 1 - alpha - 0.01
    assert covs.mean() <= 1 - alpha + 0.04
    assert covs.min() >= 1 - alpha - 0.05


def test_mondrian_fixes_class_conditional_gap():
    """Marginal LAC under-covers the hard (rare / ambiguous) classes; Mondrian does not."""
    alpha = 0.1
    per_marg, per_mond = [], []
    for seed in range(8):
        rng = np.random.default_rng(100 + seed)
        Pc, yc = _draw(1500, rng)
        Pt, yt = _draw(3000, rng)
        m = SplitConformal(alpha, "lac", False, seed=seed).fit(Pc, yc)
        c = SplitConformal(alpha, "lac", True, seed=seed).fit(Pc, yc)
        per_marg.append(list(evaluate_sets(m.predict_sets(Pt), yt)["per_class_coverage"].values()))
        per_mond.append(list(evaluate_sets(c.predict_sets(Pt), yt)["per_class_coverage"].values()))
    worst_marg = np.min(np.mean(per_marg, axis=0))
    worst_mond = np.min(np.mean(per_mond, axis=0))
    assert worst_marg < 1 - alpha - 0.02  # the gap exists
    assert worst_mond >= 1 - alpha - 0.015  # and Mondrian closes it (on average)


def test_quantile_edge_cases():
    s = np.arange(1, 10) / 10.0  # n = 9
    assert conformal_quantile(s, 0.05) == float("inf")  # ceil(10*0.95)=10 > 9
    assert conformal_quantile(s, 0.2) == 0.8  # ceil(10*0.8)=8 -> 8th smallest