"""Independent numerical checks for numberdb.org/T490.

Run: sage -python checks.py
Read the saved output as JSON; the `rows` field contains the certified decimal
entries. Run with --stored to repeat on the table fetched from the API (a
private draft needs NUMBERDB_API_KEY). All numerical assertions run before a
successful JSON report is emitted.

The independent DLMF 10.65.1 series uses exact rational coefficients and the
exact cosine cycle for orders 0, 1, 2. After N terms its absolute remainder is
at most a_N/(1-r), where a_k=(x/2)^n (x^2/4)^k/[k!(n+k)!] and
r=(x^2/4)/[(N+1)(n+N+1)]<1 bounds every subsequent term ratio.
The rational sum is converted to a 600-bit ball only at the end; odd orders
include 1/sqrt(2). The explicit tail is added as an error bound.

A separate Schlaefli integral (DLMF 10.64.1), using mpmath quadrature at 125
decimal digits, checks every row to 110 relative decimal places. This is an
independent numerical check, not the source of the table's proven rigour.
"""
import json
import math
import os
import sys
from decimal import Decimal
from fractions import Fraction
from pathlib import Path
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
from sage.env import SAGE_VERSION
from numberdb._generate import _written, _check_precision, _check_rigour
import mpmath as mp
from generate import KelvinBer
from check import exactness, measure, names_its_rings

R = RealBallField(600)
C = ComplexBallField(600)


def series(n, x, terms=200):
    """A rigorous independent series enclosure for integer n=0,1,2."""
    x = QQ(x)
    a = (x/QQ(2))**n / QQ(math.factorial(n))
    total = QQ(0)
    cycles = ((1, 0, -1, 0), (-1, -1, 1, 1), (0, 1, 0, -1))
    for k in range(terms):
        total += cycles[n][k % 4]*a
        a *= x*x / QQ(4*(k+1)*(n+k+1))
    ratio = x*x / QQ(4*(terms+1)*(n+terms+1))
    assert ratio < 1
    tail = a / (1-ratio)
    assert tail < QQ(1)/10**150
    value = R(total)
    if n == 1:
        value /= R(2).sqrt()
    return value.add_error(R(tail))


def rational_decimal(text):
    return QQ(str(Fraction(Decimal(text))))


def written_interval(text):
    d = Decimal(text)
    center = QQ(str(Fraction(d)))
    unit = QQ(10)**int(d.as_tuple().exponent)
    return center-unit, center+unit, unit


def tight(v):
    assert v.is_finite(), 'nonfinite ball'
    assert not v.contains_zero(), 'comparison lacks relative accuracy'
    assert v.rad() < abs(v).lower()*QQ(1)/10**110, 'ball too wide'


def integral(n, x):
    x = mp.mpf(str(x.numerator()))/mp.mpf(str(x.denominator()))
    a = x/mp.sqrt(2)
    return (-1)**n/mp.pi * mp.quad(
        lambda t: mp.cos(a*mp.sin(t)-n*t)*mp.cosh(a*mp.sin(t)),
        [0, mp.pi/2, mp.pi])


def check_rows(rows, source):
    mp.mp.dps = 125
    # Known nonzero and zero controls, before interpreting quadrature checks.
    assert abs(integral(0, QQ(0))-1) < mp.mpf('1e-120')
    assert abs(integral(1, QQ(0))) < mp.mpf('1e-120')
    assert series(0, QQ(0)) == 1
    assert series(1, QQ(0)) == 0
    assert series(2, QQ(0)) == 0
    mp.mp.dps = 160
    by_key = {(int(p['nu']), QQ(p['x'])): text for p,text in rows}
    expected = {(n, QQ(a)/b) for n in range(3)
                for b in range(1,5) for a in range(1,20*b+1)}
    assert set(by_key) == expected and len(rows) == 360
    worst_mp = mp.mpf(0)
    worst_integral = mp.mpf(0)
    for j, (params, text) in enumerate(rows):
        n, x = int(params['nu']), QQ(params['x'])
        low, high, unit = written_interval(text)
        s = series(n, x)
        tight(s)
        assert s.lower() > low and s.upper() < high, (params,'series not contained')
        w = C(x)*C(1,1)/C(2).sqrt()
        alt = ((C(0,1)**n)*w.bessel_I(n)).real()
        tight(alt)
        assert alt.lower() > low and alt.upper() < high, (params,'rotated I')
        z = C(x)*C(-1,1)/C(2).sqrt()
        reflected = z.bessel_J(-n).real()*((-1)**n)
        parity = (-z).bessel_J(n).real()*((-1)**n)
        for v in [reflected, parity]:
            tight(v)
            assert v.lower() > low and v.upper() < high
        mpx = mp.mpf(str(x.numerator()))/mp.mpf(str(x.denominator()))
        outside = mp.ber(n, mpx)
        mid = mp.mpf(text)
        mpunit = mp.mpf(str(unit.numerator()))/mp.mpf(str(unit.denominator()))
        error = abs(outside-mid)/mpunit
        assert error < 1, (params,'mpmath outside printed interval')
        worst_mp = max(worst_mp,error)
        with mp.workdps(125):
            q = integral(n,x)
            error = abs(q-outside)/abs(outside)
            assert error < mp.mpf('1e-110'), (params,'integral disagreement')
            worst_integral = max(worst_integral,error)
        if (j+1)%60 == 0:
            print(f'{source}: independently checked {j+1}/360', file=sys.stderr,flush=True)
    # Recurrence at nu=1 uses three stored ber values and independently
    # evaluated bei_1, so it tests the data read back rather than the generator.
    for x in sorted({key[1] for key in by_key}):
        balls = {}
        for n in range(3):
            lo, hi, unit = written_interval(by_key[n,x])
            balls[n] = R((lo+hi)/2).add_error(R(unit))
        bei = (C(x)*C(-1,1)/C(2).sqrt()).bessel_J(1).imag()
        residual = balls[0]+balls[2]+R(2).sqrt()/R(x)*(balls[1]-bei)
        assert residual.is_finite() and residual.contains_zero()
        assert residual.rad() < QQ(1)/10**90
    return {'checked':len(rows), 'recurrences':120,
            'mpmath_max_error_in_last_place':str(worst_mp),
            'integral_max_relative_error':str(worst_integral),
            'controls':'series at x=0: 1,0,0; integral at x=0: 1,0',
            'series_terms':200,'series_precision_bits':600,
            'mpmath_dps':160,'quadrature_dps':125}


def stored_rows(client):
    doc = numberdb.table('T490',client=client)
    # The public client API uses nested parameter keys.
    return [(dict(nu=str(n),x=str(x)),entry['number'] if isinstance(entry,dict) else entry)
            for n,group in doc['Numbers'].items() for x,entry in group.items()]


def main(stored=False):
    g = KelvinBer()
    if stored:
        import urllib.request
        key = (sys.stdin.read().strip() if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1'
               else os.environ.get('NUMBERDB_API_KEY'))
        c = numberdb.Client(api_key=key, opener=urllib.request.urlopen)
        rows = stored_rows(c)
        verification = g.verify(sample=None,client=c)
        assert verification.ok, str(verification)
        extra = {'verify':str(verification)}
    else:
        values = []
        rows = []
        max_rad = 0.0
        min_relative_digits = 1000.0
        for p in g.enumerate():
            v = g.value(p,g.digits)
            tight(v)
            identity = p['nu']+','+p['x']
            text = _written(v,g.digits,g.format)
            _check_rigour(g,g.table,identity,v,bounded=True)
            _check_precision(g.table,identity,text,g.digits,False)
            lo,hi,unit = written_interval(text)
            assert v.lower() > lo and v.upper() < hi
            max_rad = max(max_rad,float(v.rad()))
            min_relative_digits = min(min_relative_digits,
                                      -math.log10(float(v.rad())/abs(float(v))))
            values.append(v)
            rows.append((p,text))
        assert not exactness(values)
        assert not names_its_rings(str(Path(__file__).with_name('generate.py')))
        extra = {'measurement':measure({p['nu']+','+p['x']:t for p,t in rows}),
                 'maximum_ball_radius_approx':max_rad,
                 'minimum_relative_decimal_accuracy_approx':min_relative_digits,
                 'working_precision_bits':numberdb.bits(100,losing=64)}
    result = check_rows(rows, 'stored' if stored else 'computed')
    result.update(extra)
    result['versions']={'sage':SAGE_VERSION,'python':sys.version.split()[0],
                        'mpmath':mp.__version__}
    result['rows']=rows
    print(json.dumps(result,indent=2))


if __name__ == '__main__':
    main('--stored' in sys.argv)