"""
Statistical functions for benchmark evaluation.

Provides confidence intervals, sample size calculations, and
proportion statistics for interpreting benchmark results.
All implementations are pure Python.
"""

import math


def normal_quantile(p: float) -> float:
    """
    Inverse normal CDF (percent-point function).

    Uses the Beasley-Springer-Moro rational approximation,
    accurate to approximately 1e-8.

    Args:
        p: Probability in (0, 1).

    Returns:
        z such that P(Z <= z) = p for standard normal Z.

    Raises:
        ValueError: If p is not in (0, 1).
    """
    if p <= 0 or p >= 1:
        raise ValueError(f"p must be in (0, 1), got {p}")

    # Rational approximation coefficients (Beasley-Springer-Moro)
    a = [
        -3.969683028665376e01,
        2.209460984245205e02,
        -2.759285104469687e02,
        1.383577518672690e02,
        -3.066479806614716e01,
        2.506628277459239e00,
    ]
    b = [
        -5.447609879822406e01,
        1.615858368580409e02,
        -1.556989798598866e02,
        6.680131188771972e01,
        -1.328068155288572e01,
    ]
    c = [
        -7.784894002430293e-03,
        -3.223964580411365e-01,
        -2.400758277161838e00,
        -2.549732539343734e00,
        4.374664141464968e00,
        2.938163982698783e00,
    ]
    d = [
        7.784695709041462e-03,
        3.224671290700398e-01,
        2.445134137142996e00,
        3.754408661907416e00,
    ]

    p_low = 0.02425
    p_high = 1 - p_low

    if p < p_low:
        # Lower tail
        q = math.sqrt(-2 * math.log(p))
        return (
            ((((c[0] * q + c[1]) * q + c[2]) * q + c[3]) * q + c[4]) * q + c[5]
        ) / ((((d[0] * q + d[1]) * q + d[2]) * q + d[3]) * q + 1)
    if p <= p_high:
        # Central region
        q = p - 0.5
        r = q * q
        return (
            (
                ((((a[0] * r + a[1]) * r + a[2]) * r + a[3]) * r + a[4]) * r
                + a[5]
            )
            * q
        ) / (((((b[0] * r + b[1]) * r + b[2]) * r + b[3]) * r + b[4]) * r + 1)
    # Upper tail
    q = math.sqrt(-2 * math.log(1 - p))
    return -(
        ((((c[0] * q + c[1]) * q + c[2]) * q + c[3]) * q + c[4]) * q + c[5]
    ) / ((((d[0] * q + d[1]) * q + d[2]) * q + d[3]) * q + 1)


def wilson_score_interval(
    successes: int,
    total: int,
    confidence: float = 0.95,
) -> dict:
    """
    Wilson score confidence interval for a binomial proportion.

    Unlike the normal (Wald) approximation, the Wilson interval:
    - Never produces bounds outside [0, 1]
    - Behaves correctly at 0% and 100% accuracy
    - Has better coverage at small sample sizes

    Args:
        successes: Number of successes (correct answers).
        total: Total number of trials (graded examples).
        confidence: Confidence level (default 0.95 for 95% CI).

    Returns:
        Dictionary with keys:
        - lower: Lower bound of the interval
        - upper: Upper bound of the interval
        - center: Center of the Wilson interval (not the raw proportion)
        - margin_of_error: Half-width of the interval
        - sample_size: The total n used
    """
    if total <= 0:
        return {
            "lower": 0.0,
            "upper": 0.0,
            "center": 0.0,
            "margin_of_error": 0.0,
            "sample_size": 0,
        }

    if successes < 0 or successes > total:
        raise ValueError(
            f"successes must be in [0, total], got successes={successes}, total={total}"
        )

    p_hat = successes / total
    z = normal_quantile(1 - (1 - confidence) / 2)
    z2 = z * z

    denominator = 1 + z2 / total
    center = (p_hat + z2 / (2 * total)) / denominator
    half_width = (
        z * math.sqrt(p_hat * (1 - p_hat) / total + z2 / (4 * total * total))
    ) / denominator

    lower = max(0.0, center - half_width)
    upper = min(1.0, center + half_width)

    return {
        "lower": lower,
        "upper": upper,
        "center": center,
        "margin_of_error": half_width,
        "sample_size": total,
    }


def proportion_std_error(p: float, n: int) -> float:
    """
    Standard error of a sample proportion.

    Args:
        p: Observed proportion (0 to 1).
        n: Sample size.

    Returns:
        Standard error, or 0.0 if n <= 0.
    """
    if n <= 0:
        return 0.0
    if p < 0 or p > 1:
        raise ValueError(f"p must be in [0, 1], got {p}")
    return math.sqrt(p * (1 - p) / n)


def sample_size_for_difference(
    p1: float,
    p2: float,
    power: float = 0.8,
    alpha: float = 0.05,
) -> int:
    """
    Required sample size per group to detect a difference between
    two independent proportions.

    Uses the formula for a two-sided two-proportion z-test:
        n = (z_alpha/2 + z_beta)^2 * (p1(1-p1) + p2(1-p2)) / (p1 - p2)^2

    Args:
        p1: Expected proportion for group 1.
        p2: Expected proportion for group 2.
        power: Statistical power (default 0.8 = 80%).
        alpha: Significance level (default 0.05 for two-sided test).

    Returns:
        Required sample size per group (rounded up).

    Raises:
        ValueError: If p1 == p2 (no difference to detect).
    """
    if p1 == p2:
        raise ValueError(
            "p1 and p2 must differ to compute required sample size"
        )
    for name, val in [("p1", p1), ("p2", p2)]:
        if val < 0 or val > 1:
            raise ValueError(f"{name} must be in [0, 1], got {val}")
    for name, val in [("power", power), ("alpha", alpha)]:
        if val <= 0 or val >= 1:
            raise ValueError(f"{name} must be in (0, 1), got {val}")

    z_alpha = normal_quantile(1 - alpha / 2)
    z_beta = normal_quantile(power)

    numerator = (z_alpha + z_beta) ** 2 * (p1 * (1 - p1) + p2 * (1 - p2))
    denominator = (p1 - p2) ** 2

    return math.ceil(numerator / denominator)
