"""Independent checks for Gi derivative zeros -- numberdb.org/T517.

    sage -python checks.py             # local certificates, no key needed
    sage -python checks.py --stored    # additionally read and check the draft

Keep generate.py and check.py beside this file. NUMBERDB_API_KEY may be
needed for draft reads; NUMBERDB_KEY_FROM_STDIN=1 reads it from stdin.
All evidence is printed as JSON for saving by the caller.

The independent integral route differentiates DLMF 9.12.4. Its cross terms
cancel, giving Gi'(x) = Bi'(x)(1/3 - integral_0^x Ai) +
Ai'(x) integral_0^x Bi. The Airy functions and certified quadrature are Arb
implementations independent of the generator's Gi coefficient recurrence.
Controls integrate t^2 and Ai' and require the known nonzero answers.
Every compared ball is finite and bounded in absolute radius.
"""
import os
import sys
import json
from decimal import Decimal
from collections import Counter
import numberdb.sage as numberdb
from sage.rings.complex_arb import ComplexBallField
from sage.rings.rational_field import QQ
import mpmath as mp
import generate as g
from check import exactness, measure, names_its_rings, agrees_with

C=ComplexBallField(g.PRECISION_BITS)
R=g.R
TOL=R(2)**(-520)


def small(z, bound='1e-145'):
    assert z.is_finite()
    assert R(z.rad()) < R(bound), 'enclosure too wide to support this check'
    return z


def real(z):
    assert z.real().is_finite() and z.imag().is_finite() and z.imag().contains_zero()
    return small(z.real())


def integral(fun, a, b):
    return real(C.integral(fun,C(a),C(b),rel_tol=TOL,abs_tol=TOL))


def controls():
    p=integral(lambda z,_:z*z,0,1)
    assert QQ(1)/3 in p and not p.contains_zero()
    a=integral(lambda z,_:z.airy_ai_prime(),0,1)
    answer=real(C(1).airy_ai()-C(0).airy_ai())
    assert a.overlaps(answer) and not a.contains_zero() and not answer.contains_zero()
    return {'polynomial':str(p),'airy_derivative':str(a),'expected_airy_difference':str(answer)}


def airy_derivative(x):
    A=integral(lambda z,_:z.airy_ai(),0,x)
    B=integral(lambda z,_:z.airy_bi(),0,x)
    z=C(x)
    return small(real(z.airy_bi_prime())*(QQ(1)/3-A)+real(z.airy_ai_prime())*B)


def exact_decimal(s):
    a,b=Decimal(str(s)).as_integer_ratio()
    return QQ(a)/b


def decimal_interval(s):
    s=str(s)
    dec=Decimal(s)
    center=exact_decimal(s)
    place=QQ(10)**int(dec.as_tuple().exponent)
    return center-place,center+place


def check_interval(n,lo,hi):
    sl=g.sign(g.evaluate(lo,1));sr=g.sign(g.evaluate(hi,1))
    assert sl*sr==-1
    curvature=g.local_enclosures(lo,hi)[1]
    assert g.sign(curvature)==(-1 if n%2==0 else 1)
    il=airy_derivative(lo);ir=airy_derivative(hi)
    assert g.sign(il)==sl and g.sign(ir)==sr
    return {'n':n,'lo':str(lo),'hi':str(hi),'left_sign':sl,'right_sign':sr,
            'curvature_sign':g.sign(curvature),'independent_left':str(il),
            'independent_right':str(ir)}


def run(stored=False):
    out={'table':'T517','controls':controls(),'environment':g.ScorerGiPrimeZeros().environment()}
    out['environment']['mpmath']=mp.__version__
    assert not names_its_rings(g.__file__)
    gen=g.ScorerGiPrimeZeros()
    values={p['n']:gen.value(p,100) for p in gen.enumerate()}
    assert not exactness(values)
    out['measure']=measure(values)
    cells,records=g.root_cells()
    out['count_certificate']={'cells':len(records),'types':dict(Counter(r['test'] for r in records)),
                              'roots':sum(bool(r.get('root')) for r in records),
                              'negative_left_endpoint':str(cells[-1][0]),'records':records}
    out['tail_bounds']=[str(g.tail(d)) for d in range(3)]
    assert sum(bool(r.get('root')) for r in records)==13
    out['root_certificates']=[]
    out['mpmath_checks']=[]
    for n in range(13):
        ball,lo,hi=g.certified_root(n)
        out['root_certificates'].append(check_interval(n,lo,hi))
        roots=[]
        for dps in (150,200):
            with mp.workdps(dps):
                a,b=cells[n]
                root=mp.findroot(lambda x:mp.diff(mp.scorergi,x),(g.mp_rational(a),g.mp_rational(b)))
                assert g.mp_rational(lo)<root<g.mp_rational(hi)
                roots.append(mp.nstr(root,140))
                # ODE, initial values and differentiated sum identity, with
                # mpmath computations that do not call the generator evaluator.
                y=mp.scorergi(root)
                assert abs(mp.diff(mp.scorergi,root,2)-root*y+1/mp.pi)<mp.mpf('1e-140')
                assert abs(mp.diff(mp.scorergi,root)+mp.diff(mp.scorerhi,root)-mp.airybi(root,1))<mp.mpf('1e-140')
                assert abs(mp.scorergi(0)-1/(mp.power(3,mp.mpf(7)/6)*mp.gamma(mp.mpf(2)/3)))<mp.mpf('1e-145')
                assert abs(mp.diff(mp.scorergi,0)-1/(mp.power(3,mp.mpf(5)/6)*mp.gamma(mp.mpf(1)/3)))<mp.mpf('1e-145')
        assert roots[0]==roots[1]
        out['mpmath_checks'].append({'n':n,'root_140_digits':roots[1],'precisions':[150,200]})
        midpoint=(lo+hi)/2
        ode=g.evaluate(midpoint,2)-R(midpoint)*g.evaluate(midpoint)+1/R.pi()
        small(ode)
        assert ode.contains_zero()
        print('checked local root',n,flush=True)
    # Use the shared independent-comparison helper as well as exact brackets.
    independent={r['n']:R(r['root_140_digits']) for r in out['mpmath_checks']}
    for n in range(13):
        assert independent[n] in g.certified_root(n)[0]
    endpoint_signs = {}
    independent_signs = {}
    for cert in out['root_certificates']:
        for side in ('left', 'right'):
            key = (cert['n'], side)
            endpoint_signs[key] = cert[side+'_sign']
            independent_signs[key] = g.sign(R(cert['independent_'+side]))
    complaints = agrees_with(endpoint_signs, lambda key: independent_signs[key])
    assert not complaints, complaints
    if stored:
        document=numberdb.table('T517')
        numbers=document['Numbers']
        assert len(numbers)==13
        out['stored_certificates']=[]
        for key,entry in numbers.items():
            n=int(key)
            written=entry['number'] if isinstance(entry,dict) else entry
            lo,hi=decimal_interval(written)
            _,a,b=g.certified_root(n)
            assert lo<a<b<hi
            out['stored_certificates'].append(check_interval(n,lo,hi))
            print('checked stored root',n,flush=True)
        report=gen.verify(sample=None)
        assert report.ok,str(report)
        out['verify']=str(report)
    print('EVIDENCE_JSON_BEGIN')
    print(json.dumps(out,indent=2))
    print('EVIDENCE_JSON_END')
    return out


if __name__=='__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
        os.environ['NUMBERDB_API_KEY']=sys.stdin.read().strip()
    run('--stored' in sys.argv or os.environ.get('NUMBERDB_PUBLISH')=='1')
