"""Independent checks for Stationary points of the Hardy Z-function, T585.

Run: sage -python checks.py. Reads the live table unless values.json is beside
this file. Set NUMBERDB_KEY_FROM_STDIN=1 and pipe a key to read a private draft.
The direct gamma/zeta derivative expression shares no root-finding or Hardy
jet code with generate.py. mpmath supplies separate 125- and 155-digit root
computations. These agreements supplement the generator's certified interval
cover. The tests must pass on the decimal intervals read back from the site.
"""
import os,sys,json,time
from pathlib import Path
from decimal import Decimal
from fractions import Fraction
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
from sage.misc.cython import cython_import
import mpmath as mp

PREC=560
R=RealBallField(PREC)
C=ComplexBallField(PREC)

def rational_decimal(text):
    f=Fraction(Decimal(text))
    return QQ(f.numerator)/QQ(f.denominator)

def written(text):
    d=Decimal(text)
    c=rational_decimal(text)
    exponent=d.as_tuple().exponent
    u=QQ(10)**exponent
    return c-u,c+u,u

def direct(t):
    """Return Z(t), Z'(t) from log-Gamma, psi, zeta and zeta' directly."""
    t=R(t)
    g=C(QQ(1)/4,t/2)
    theta=g.log_gamma().imag()-t*R.pi().log()/2
    theta_prime=(g.psi().real()-R.pi().log())/2
    phase=C(0,theta).exp()
    s=C(QQ(1)/2,t)
    z=s.zeta()
    zp=s.zetaderiv(1)
    values=[phase*z,C(0,1)*phase*(theta_prime*z+zp)]
    for v in values:
        assert v.real().is_finite() and v.imag().is_finite()
        assert v.imag().contains_zero()
    return values[0].real(),values[1].real()

def strict_sign(x):
    assert x.is_finite()
    return 1 if x>0 else -1 if x<0 else 0

def endpoint_check(left,right):
    l=direct(left)[1]; r=direct(right)[1]
    assert R(l.rad())<R(2)**(-440) and R(r.rad())<R(2)**(-440)
    return strict_sign(l)*strict_sign(r)==-1

def mp_root(text,dps):
    with mp.workdps(dps):
        x=mp.mpf(text)
        for _ in range(3):
            f=mp.siegelz(x,derivative=1)
            df=mp.siegelz(x,derivative=2)
            assert mp.isfinite(f) and mp.isfinite(df) and df != 0
            x-=f/df
        assert abs(mp.siegelz(x,derivative=1)) < mp.mpf(10)**(-dps+12)
        return mp.nstr(x,dps-5)

def records(document):
    nums=document['Numbers']
    if isinstance(nums,dict):
        return {str(k):v['number'] if isinstance(v,dict) else v for k,v in nums.items()}
    return {str(r['params']['j']):r['number'] for r in nums}

def main(values=None):
    start=time.time()
    # Known-answer controls before trusting derivatives or interval tests.
    known=C(2).zeta().real()-R.pi()**2/6
    assert known.contains_zero() and R(known.rad())<R(2)**(-500)
    z0,d0=direct(0)
    assert z0 < -1 and z0.overlaps(C(QQ(1)/2).zeta().real())
    assert d0.contains_zero() and R(d0.rad())<R(2)**(-500)
    with mp.workdps(170):
        assert mp.siegelz(0)==mp.zeta(mp.mpf(1)/2)
        assert mp.siegelz(0,derivative=1)==0
        assert abs(mp.zeta(2)-mp.pi**2/6)<mp.mpf('1e-165')
    if values is None:
        local=Path(__file__).with_name('values.json')
        if local.exists(): values=json.loads(local.read_text())
        else: values=records(numberdb.table('T585'))
    assert list(values)==[str(j) for j in range(1,101)]
    roots125={};roots155={}; rejected=0
    for j,text in values.items():
        left,right,ulp=written(text)
        assert endpoint_check(left,right),j
        # The same checker must reject intervals shifted ten last places.
        for shift in [-10,10]:
            assert not endpoint_check(left+shift*ulp,right+shift*ulp),j
            rejected+=1
        centre=(left+right)/2
        z,d=direct(centre); zn,dn=direct(-centre)
        assert (z-zn).contains_zero() and R((z-zn).rad())<R(2)**(-430)
        assert (d+dn).contains_zero() and R((d+dn).rad())<R(2)**(-430)
        for digits,store in [(125,roots125),(155,roots155)]:
            independent=mp_root(text,digits)
            q=rational_decimal(independent)
            assert left<q<right,(j,digits)
            store[j]=independent
        if int(j)%10==0: print('checked',j,flush=True)
    # External landmarks: certify the first 99 zeta zeros and their count.
    bridge=cython_import(str(Path(__file__).with_name('hardy_arb.pyx')))
    zeros=bridge.zero_prefix(99,128)
    assert bridge.count(233,128)==98
    assert bridge.count(234,128)==99
    assert zeros[-1].imag()<234
    for n,rho in enumerate(zeros,1):
        assert rho.real()==QQ(1)/2
        assert rho.imag().is_finite() and rho.imag().accuracy()>110
        if n>1: assert zeros[n-2].imag()<rho.imag()
    assert R(written(values['2'])[1])<zeros[0].imag()
    for j in range(3,101):
        left,right,_=written(values[str(j)])
        assert zeros[j-3].imag()<R(left)<R(right)<zeros[j-2].imag()
    # Hall p.125 quotes only coarse initial locations and signed heights.
    assert 2<R(written(values['1'])[0])<3
    assert 10<R(written(values['2'])[0])<11
    for j,lo,hi in [('1',QQ(-53)/100,QQ(-51)/100),('2',QQ(-157)/100,QQ(-154)/100)]:
        left,right,_=written(values[j]); z,_=direct((left+right)/2)
        assert R(lo)<z<R(hi)
    result={'table':'T585','entries':100,'direct_endpoint_certificates':100,
        'shifted_intervals_rejected':rejected,'parity_checks':100,
        'mpmath_root_precisions':[125,155],'mpmath_roots_per_precision':100,
        'nontrivial_zeta_zeros_below_233':98,'nontrivial_zeta_zeros_below_234':99,'verified_zeta_gap_relations':98,
        'controls':'zeta(2)=pi^2/6; Z(0)=zeta(1/2)<-1; Zprime(0)=0; 200 shifted intervals rejected',
        'hall_initial_heights':'both signed critical values agree with the coarse quoted values',
        'mpmath_version':mp.__version__,'flint_version':bridge.flint_version(),
        'python_version':sys.version,'seconds':time.time()-start}
    print('CHECKS_JSON',json.dumps(result))
    print('MP125_JSON',json.dumps(roots125))
    print('MP155_JSON',json.dumps(roots155))
    return result

if __name__=='__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
        key=sys.stdin.read().strip()
        if '=' in key: key=key.split('=',1)[1].strip().strip('\"\'')
        os.environ['NUMBERDB_API_KEY']=key
    main()