"""Confidence bounds used to certify a routing policy on the held-out split.

Everything here is written from scratch (regularised incomplete beta by continued fraction,
Lanczos log-gamma) so that web/router.js can carry the same code line for line and the two
implementations give identical digits. tests/test_bounds.py checks it against scipy.
"""
from __future__ import annotations

import math

_LANCZOS = (0.99999999999980993, 676.5203681218851, -1259.1392167224028, 771.32342877765313,
            -176.61502916214059, 12.507343278686905, -0.13857109526572012,
            9.9843695780195716e-6, 1.5056327351493116e-7)


def log_gamma(x: float) -> float:
    """Lanczos approximation (g = 7, 9 terms), good to about 1e-15 for x > 0."""
    if x < 0.5:
        return math.log(math.pi / math.sin(math.pi * x)) - log_gamma(1.0 - x)
    x -= 1.0
    a = _LANCZOS[0]
    t = x + 7.5
    for i in range(1, 9):
        a += _LANCZOS[i] / (x + i)
    return 0.5 * math.log(2.0 * math.pi) + (x + 0.5) * math.log(t) - t + math.log(a)


def _betacf(a: float, b: float, x: float) -> float:
    """Continued fraction for the incomplete beta function (modified Lentz)."""
    tiny = 1e-300
    qab, qap, qam = a + b, a + 1.0, a - 1.0
    c, d = 1.0, 1.0 - qab * x / qap
    if abs(d) < tiny:
        d = tiny
    d = 1.0 / d
    h = d
    for m in range(1, 300):
        m2 = 2 * m
        aa = m * (b - m) * x / ((qam + m2) * (a + m2))
        d = 1.0 + aa * d
        if abs(d) < tiny:
            d = tiny
        c = 1.0 + aa / c
        if abs(c) < tiny:
            c = tiny
        d = 1.0 / d
        h *= d * c
        aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2))
        d = 1.0 + aa * d
        if abs(d) < tiny:
            d = tiny
        c = 1.0 + aa / c
        if abs(c) < tiny:
            c = tiny
        d = 1.0 / d
        de = d * c
        h *= de
        if abs(de - 1.0) < 3e-16:
            break
    return h


def beta_cdf(x: float, a: float, b: float) -> float:
    """Regularised incomplete beta I_x(a, b)."""
    if x <= 0.0:
        return 0.0
    if x >= 1.0:
        return 1.0
    lbeta = log_gamma(a + b) - log_gamma(a) - log_gamma(b) + a * math.log(x) + b * math.log(1.0 - x)
    front = math.exp(lbeta)
    if x < (a + 1.0) / (a + b + 2.0):
        return front * _betacf(a, b, x) / a
    return 1.0 - front * _betacf(b, a, 1.0 - x) / b


def beta_ppf(p: float, a: float, b: float) -> float:
    """Inverse of beta_cdf in x, by bisection (80 steps: exact to double precision)."""
    lo, hi = 0.0, 1.0
    for _ in range(80):
        mid = 0.5 * (lo + hi)
        if beta_cdf(mid, a, b) < p:
            lo = mid
        else:
            hi = mid
    return 0.5 * (lo + hi)


def clopper_pearson_lower(k: float, n: int, delta: float) -> float:
    """One-sided lower bound on a Bernoulli mean from k successes in n trials.

    With probability at least 1 - delta (over the random held-out sample) the true accuracy
    is above this number. k = 0 gives 0; k = n gives delta**(1/n).
    """
    if n <= 0:
        return 0.0
    k = int(math.floor(k + 0.5))
    if k <= 0:
        return 0.0
    if k >= n:
        return delta ** (1.0 / n)
    return beta_ppf(delta, k, n - k + 1)


def hoeffding_radius(n: int, delta: float, span: float = 1.0) -> float:
    """Half-width of a one-sided Hoeffding bound for a mean of n values in an interval of width span."""
    if n <= 0:
        return float("inf")
    return span * math.sqrt(math.log(1.0 / delta) / (2.0 * n))


def quality_lower_bound(total: float, n: int, delta: float, binary: bool) -> float:
    """Lower bound on mean quality: Clopper-Pearson when scores are 0/1, Hoeffding otherwise."""
    if n <= 0:
        return 0.0
    if binary:
        return clopper_pearson_lower(total, n, delta)
    return max(0.0, total / n - hoeffding_radius(n, delta, 1.0))


def rows_to_certify(p_hat: float, target: float, delta: float, binary: bool, cap: int = 10_000_000):
    """Smallest held-out size at which this observed rate would certify the target.

    Returns None when p_hat < target (no amount of data would help) and cap+1 when more than cap
    rows would be needed.
    """
    if p_hat < target:
        return None
    if not binary:
        gap = p_hat - target
        if gap <= 0:
            return cap + 1
        n = math.ceil(math.log(1.0 / delta) / (2.0 * gap * gap))
        return min(n, cap + 1)
    lo, hi = 1, 1
    while quality_lower_bound(int(math.floor(p_hat * hi + 0.5)), hi, delta, True) < target:
        hi *= 2
        if hi > cap:
            return cap + 1
    while hi - lo > 1:
        mid = (lo + hi) // 2
        if quality_lower_bound(int(math.floor(p_hat * mid + 0.5)), mid, delta, True) >= target:
            hi = mid
        else:
            lo = mid
    return hi


def required_rate(target: float, n: int, delta: float, binary: bool):
    """Smallest observed rate on n held-out rows whose lower bound still clears the target.

    This is the margin a policy must be tuned to before the held-out rows are looked at. None
    means that even a perfect score on n rows could not certify the target (too few rows).
    """
    if n <= 0:
        return None
    if not binary:
        need = target + hoeffding_radius(n, delta, 1.0)
        return need if need <= 1.0 else None
    if quality_lower_bound(n, n, delta, True) < target:
        return None
    lo, hi = 0, n          # smallest k with LB(k, n) >= target; LB is increasing in k
    while hi - lo > 1:
        mid = (lo + hi) // 2
        if quality_lower_bound(mid, n, delta, True) >= target:
            hi = mid
        else:
            lo = mid
    if quality_lower_bound(lo, n, delta, True) >= target:
        hi = lo
    return hi / n