"""Quantiles of the studentized range distribution -- numberdb.org/T457.

Run with SageMath (or passagemath):
    sage -pip install numberdb mpmath scipy  # once
    sage -python generate.py                  # recompute and verify every row
    sage -python generate.py --fill-draft     # fill the existing draft
Set NUMBERDB_API_KEY for a private draft, or pipe the key on stdin with
NUMBERDB_KEY_FROM_STDIN=1. Filling entries does not publish a table.

The selection is k=2,3, nu=1,...,10,20, and fourteen conventional
probabilities. Values are enclosed to 100 significant digits. mpmath
quadrature finds candidates only; Arb signs at exact rational endpoints
certify the root brackets. See math-notes.md for the three-mean angular
integral, its coefficient recurrence and its explicit Cauchy remainder bound.
No numerical integration error estimate is used as a proof.
"""
import os
import sys
from functools import lru_cache

import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from sage.rings.real_arb import RealBallField
from sage.rings.complex_arb import ComplexBallField
import mpmath as mp
from scipy.stats import studentized_range, t

PROBABILITIES = ('1/200','1/100','1/40','1/20','1/10','1/4','1/2',
                 '3/4','9/10','19/20','39/40','99/100','199/200','999/1000')
DEGREES = tuple(range(1,11)) + (20,)
WORKING_GUARD = 192


def mp_rational(s):
    x = QQ(s)
    return mp.mpf(int(x.numerator())) / int(x.denominator())


def cdf_candidate(q, k, nu):
    """Fixed precision, used only to propose a bracket, never certify it."""
    if k == 2:
        return mp.betainc(mp.mpf(1)/2, mp.mpf(nu)/2, 0,
                         q*q/(2*nu+q*q), regularized=True)
    return 1 - 6/mp.pi * mp.quad(
        lambda theta: (1+q*q/(2*nu*mp.cos(theta)**2))**(-mp.mpf(nu)/2),
        [0,mp.pi/6])


def beta_regularized(x, a, b, B):
    C = ComplexBallField(B.precision())
    def lower(z, aa, bb):
        hyp = C(z).hypergeometric([C(aa), C(1)-C(bb)], [C(aa)+1])
        assert hyp.imag().contains_zero()
        return z**aa / aa * hyp.real() * (aa+bb).gamma() / (aa.gamma()*bb.gamma())
    if x > B(1)/2:
        return 1-lower(1-x,b,a)
    return lower(x,a,b)


def cdf_ball(q, k, nu, bits):
    """Certified CDF at a positive exact rational q."""
    B = RealBallField(bits)
    x = B(q)
    if k == 2:
        z = x*x/(2*nu+x*x)
        return beta_regularized(z,B(1)/2,B(nu)/2,B)
    if k != 3:
        raise ValueError('This generator implements only k=2 and k=3.')
    a = x*x/(2*nu)
    c = 1/(1+a)
    b = B(nu-1)/2
    d = B(nu)/2
    # On |z|=1/2, |(1-z)^b (1-cz)^(-d)| <= 3^nu.
    # Integrating the series at sin(theta)<=1/2 gives tail <3^nu*2^(-N).
    # N uses 2*nu rather than nu*log2(3), a convenient exact upper bound.
    N = bits + 2*nu + 32
    previous, current = B(0), B(1)
    total = B(0)
    for n in range(N):
        total += current/(2*n+1)
        following = (((1+c)*n-b+d*c)*current/4
                     + c*(b-d-n+1)*previous/16)/(n+1)
        previous, current = current, following
    integral = total/2
    integral = integral.add_error(B(QQ(3)**nu / QQ(2)**N))
    return 1 - (6/B.pi()) * (1+a)**(-d) * integral


@lru_cache(maxsize=None)
def enclose(k, nu, probability, digits=100):
    p = QQ(probability)
    with mp.workdps(digits+35):
        if k == 2:
            guess = mp.mpf(str(float(t.ppf((1+float(p))/2,nu))))*mp.sqrt(2)
        else:
            guess = mp.mpf(str(float(studentized_range.ppf(float(p),k,nu))))
        root = mp.findroot(lambda q:cdf_candidate(q,k,nu)-mp_rational(probability),
                           (guess,guess*mp.mpf('1.001')),
                           tol=mp.mpf(10)**(-digits-25),maxsteps=40)
        # The candidate becomes exact rational endpoints before certification.
        scale = QQ(10)**(digits+15)
        center = QQ(int(mp.floor(root*int(scale))))/scale
    radius = QQ(2)/scale
    low, high = center-radius, center+radius
    assert low>0
    bits = numberdb.bits(digits,losing=WORKING_GUARD)
    for extra in (0,128,256,512,1024):
        working = bits+extra
        f_low = cdf_ball(low,k,nu,working)
        f_high = cdf_ball(high,k,nu,working)
        if not f_low.is_finite() or not f_high.is_finite():
            continue
        if f_low < p and f_high > p:
            B=RealBallField(working)
            return B((low+high)/2).add_error(B((high-low)/2))
    raise ArithmeticError('The proposed root bracket was not certified: %s' % ((k,nu,probability),))


class StudentizedRangeQuantiles(numberdb.Generator):
    table = 'T457'
    parameters = ('nu','k','p')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py','checks.py','math-notes.md','reference-values.json')

    def enumerate(self):
        for nu in DEGREES:
            for k in (2,3):
                for p in PROBABILITIES:
                    yield {'nu':str(nu),'k':str(k),'p':p}

    def value(self, params, digits):
        return enclose(int(params['k']),int(params['nu']),params['p'],digits)


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    generator = StudentizedRangeQuantiles()
    if '--fill-draft' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(generator.publish(message='Fill the draft with certified 100-digit two- and three-mean quantiles'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)