"""Independent checks for planar-walk moments -- numberdb.org/T522.

Run: sage -python checks.py (offline), or --stored (NUMBERDB_API_KEY needed
for this private draft). Bessel-integral checks share neither seeds nor
recurrence code with generate.py. They are numerical checks at 50 digits,
not certified enclosures. The main generator supplies the certification.
"""
import os, sys, json, time
from functools import lru_cache
from decimal import Decimal
from fractions import Fraction
from pathlib import Path
import numberdb.sage as numberdb
from sage.rings.real_arb import RealBallField
from sage.rings.complex_arb import ComplexBallField
from sage.rings.rational_field import QQ
import mpmath as mp
from generate import PlanarWalkMoments, moments, log_integrals
from check import exactness, measure, agrees_with, names_its_rings
from numberdb._write import to_text


def bessel_moment(n, s):
    """BSWZ (2012), eqs. (75),(76), valid for Re(s)>-2."""
    half = mp.mpf(s)/2
    if n == 3:
        factor = mp.power(3, s+mp.mpf(3)/2)/(mp.pi*mp.power(2,s)*mp.gamma(half+1)**2)
    else:
        factor = mp.power(4,s+2)/(mp.pi**2*mp.gamma(half+1)**2)
    return factor * mp.quad(lambda t: t**(s+1)*bessel_product(t,n), [0,1,4,16,64,mp.inf])


@lru_cache(maxsize=24000)
def bessel_values(t):
    return mp.besseli(0,t), mp.besselk(0,t)


def bessel_product(t,n):
    i,k = bessel_values(t)
    return i*k**(n-1)


def check_recurrences(values, exact=False):
    tests=[]
    for n in (3,4):
        for s in (-1,1,3):
            if n==3:
                residual=(s+4)**2*values[n,s+4]-2*(5*s*s+30*s+46)*values[n,s+2]+9*(s+2)**2*values[n,s]
            else:
                residual=(s+4)**3*values[n,s+4]-4*(s+3)*(5*s*s+30*s+48)*values[n,s+2]+64*(s+2)**3*values[n,s]
            assert residual.contains_zero(), (n,s,residual)
            tests.append({'n':n,'s':s,'residual':str(residual)})
    return tests


def main():
    start=time.monotonic()
    mp.mp.dps=50
    report={'table':'T522','instruction_manifest':'../20261026T172724Z-build-skills/manifest.json',
            'instruction_manifest_sha256':'d03b756eaa1cd22ce39c8cee78fb3a44b052146786d6371925ac7bc9dec564f3',
            'instruction_sha256':'b4d51ded0d52cfe3dcc0859de0845c3879b0c2363273ba7e809a7e06ded40b18',
            'versions':{'numberdb':numberdb.__version__,'mpmath':mp.__version__,'python':sys.version},
            'controls':[]}
    try:
        from sage.env import SAGE_VERSION
        report['versions']['sage']=SAGE_VERSION
    except ImportError: pass
    # Known controls before trusting either quadrature.
    C=ComplexBallField(500)
    control=C.integral(lambda z,analytic:z*z,QQ(1)/2,1)
    assert (QQ(7)/24 in control.real()) and control.real().rad()<QQ(1)/10**120
    assert control.imag().is_finite() and control.imag().contains_zero()
    report['controls'].append({'method':'Arb integral k^2, 1/2..1','known':'7/24','result':str(control)})
    for n in (3,4):
        for s,known in ((0,1),(2,n)):
            t=time.monotonic(); value=bessel_moment(n,s)
            assert abs(value-known)<mp.mpf('1e-42'),(n,s,value)
            row={'n':n,'s':s,'known':known,'value':mp.nstr(value,50),'absolute_error':mp.nstr(abs(value-known),6),'seconds':time.monotonic()-t}
            report['controls'].append(row); print('CONTROL',json.dumps(row),flush=True)
    g=PlanarWalkMoments(); values={}; evidence=[]
    for params in g.enumerate():
        n,s=params['n'],params['s']; lo=moments(n,150)[s]; hi=moments(n,200)[s]
        assert lo.overlaps(hi)
        value=g.value(params,100); values[n,s]=value
        t=time.monotonic(); independent=bessel_moment(n,s)
        midpoint=mp.mpf(str(hi.mid()))
        rel=abs(independent-midpoint)/abs(midpoint)
        assert rel<mp.mpf('1e-42'),(n,s,rel)
        row={'n':n,'s':s,'certified_100_digits':to_text(hi,digits=100),'radius_150':str(lo.rad()),'radius_200':str(hi.rad()),'precision_overlap':True,'bessel_50_digits':mp.nstr(independent,50),'relative_difference':mp.nstr(rel,8),'seconds':time.monotonic()-t}
        evidence.append(row); print('INDEPENDENT',json.dumps(row),flush=True)
    assert not exactness(values),exactness(values)
    assert not names_its_rings(__file__.replace('checks.py','generate.py'))
    report['rows']=evidence; report['measure']=measure(values)
    report['recurrences']=check_recurrences({(n,s):moments(n,150)[s] for n,s in values})
    # Independent gamma-to-elliptic seed identity (BSW eqs. (4),(7)).
    R=RealBallField(550); CC=ComplexBallField(550)
    m=(2-R(3).sqrt())/4
    k=CC(m).elliptic_k().real()
    a=2*R(3).sqrt()*k*k/R.pi()**2
    b=a+R(3).sqrt()/(k*k)
    assert a.overlaps(moments(3,150)[-1]); assert b.overlaps(moments(3,150)[1])
    report['elliptic_three_step_seed_check']=True
    # Formula integral k^q L^2 is independently integrated at q=0,2,10.
    for q in (0,2,10):
        lhs=mp.quad(lambda k:k**q*mp.log(4/k)**2,[0,mp.mpf(1)/2])
        rhs=log_integrals(R,q)[2]
        assert abs(lhs-mp.mpf(str(rhs.mid())))<mp.mpf('1e-45')
    report['log_integral_controls']=True
    reference=json.loads(Path(__file__).with_name('reference-values.json').read_text())
    fcc=reference['T393']['values']['number']
    fcc_ball=R(QQ(Fraction(Decimal(fcc))),QQ(10)**Decimal(fcc).as_tuple().exponent)
    assert (R(2)/3*fcc_ball).overlaps(moments(3,150)[-1])
    report['stored_watson_identity']=True
    ktext=reference['T296']['values']['number']
    kball=R(QQ(Fraction(Decimal(ktext))),QQ(10)**Decimal(ktext).as_tuple().exponent)
    assert k.overlaps(kball)
    report['stored_singular_elliptic_identity']=True
    report['mahler_derivatives']=[]
    for n in (3,4):
        factor=3*mp.sqrt(3)/mp.pi if n==3 else 16/mp.pi**2
        logfactor=mp.log(mp.mpf(3)/2)+mp.euler if n==3 else mp.log(4)+mp.euler
        derivative=logfactor+factor*mp.quad(lambda t:t*mp.log(t)*bessel_product(t,n),[0,1,4,16,64,mp.inf])
        val=reference['T283']['values'][str(n)]
        if isinstance(val,dict): val=val['number']
        difference=abs(derivative-mp.mpf(val))
        assert difference<mp.mpf('1e-42'), (n,difference)
        report['mahler_derivatives'].append({'n':n,'value':mp.nstr(derivative,50),'absolute_difference':mp.nstr(difference,8)})

    if '--stored' in sys.argv or os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
        if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
            key=sys.stdin.read().strip()
            if '=' in key and key.split('=',1)[0].isidentifier(): key=key.split('=',1)[1].strip().strip('\"\'')
            numberdb.configure(api_key=key)
        result=g.verify(sample=None)
        assert result.ok,str(result)
        print('VERIFY',result,flush=True)
        document=numberdb.table('T522')
        records=document['Numbers']
        sv={}
        if isinstance(records,dict):
            records=[{'params':{'n':n,'s':s}, **(v if isinstance(v,dict) else {'number':v})}
                     for n,rows in records.items() for s,v in rows.items()]
        for record in records:
            params=record['params']
            n,s=int(params['n']),int(params['s'])
            text=record['number']
            center=QQ(Fraction(Decimal(text))); ulp=QQ(10)**Decimal(text).as_tuple().exponent
            sv[n,s]=RealBallField(550)(center,ulp)
        assert set(sv)==set(values)
        report['stored_recurrences']=check_recurrences(sv)
        for row in evidence:
            v=sv[row['n'],row['s']]
            certified=moments(row['n'],200)[row['s']]
            assert v.lower() <= certified.lower() and certified.upper() <= v.upper()
            independent=mp.mpf(row['bessel_50_digits'])
            assert abs(independent-mp.mpf(str(v.mid())))/abs(independent)<mp.mpf('1e-42')
        report['stored_independent_rows']=len(sv)
        report['stored_intervals_contain_certified_balls']=True
        report['verify']=str(result)
    report['seconds']=time.monotonic()-start
    print('EVIDENCE_JSON',json.dumps(report,sort_keys=True),flush=True)

if __name__=='__main__': main()