"""Independent checks for theta-two CM constants -- numberdb.org/T507.

Run sage -python checks.py before filling, or checks.py --stored afterwards.
Use NUMBERDB_KEY_FROM_STDIN=1 with a key on stdin to read the private draft.
The generator uses Arb's theta routine. These checks use elementary real-ball
series and products with explicit tails at 512 bits; no call to that routine
occurs in these independent computations. The Gamma special value and elliptic
relations give checks outside the defining theta-two series. All stored
numbers are interpreted as intervals, including one last-place unit on either
side of a decimal. Controls reject nonfinite, wide and shifted enclosures.
"""
import json
import math
import os
import sys
from pathlib import Path
from fractions import Fraction
from decimal import Decimal
from importlib.metadata import version
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.version import version as sage_version
from numberdb._write import to_text
from generate import ThetaTwoCM
from check import exactness, measure, names_its_rings

R = RealBallField(512)
C = ComplexBallField(512)
N = 12
M = 80


def finite(x):
    return all(math.isfinite(float(t)) for t in (x.mid(), x.rad()))


def narrow(x):
    assert finite(x), 'nonfinite enclosure'
    assert R(x.rad()) < R(1) / 10**90, 'enclosure too wide'


def agree(x, y):
    x, y = R(x), R(y)
    narrow(x)
    narrow(y)
    assert x.overlaps(y), 'disjoint enclosures'


def rational(x):
    f = Fraction(x)
    return QQ(f.numerator) / QQ(f.denominator)


def decimal_bounds(text):
    if isinstance(text, dict):
        text = text['number']
    text = str(text)
    if '.' not in text and 'e' not in text.lower():
        q = rational(text)
        return q, q
    d = Decimal(text)
    q = rational(text)
    unit = QQ(10)**int(d.as_tuple().exponent)
    return q-unit, q+unit


def interval(text):
    low, high = decimal_bounds(text)
    return R((low+high)/2).add_error(R((high-low)/2))


def theta_series(r):
    q = (-R.pi()*R(r).sqrt()).exp()
    two = 2*sum((q**((QQ(n)+QQ(1)/2)**2) for n in range(N+1)), R(0))
    two_tail = 2*q**((QQ(N)+QQ(3)/2)**2)/(1-q**(2*N+4))
    tail = 2*q**((N+1)**2)/(1-q**(2*N+3))
    three = 1+2*sum((q**(n*n) for n in range(1,N+1)),R(0))
    four = 1+2*sum(((-1)**n*q**(n*n) for n in range(1,N+1)),R(0))
    # Symmetric errors enclose each one-sided or alternating tail.
    return q, two.add_error(two_tail), three.add_error(tail), four.add_error(tail)


def products(q):
    Q = q*q
    minus = R(1)
    plus = R(1)
    for n in range(1,M+1):
        minus *= 1-Q**n
        plus *= (1+Q**n)**2
    s = Q**(M+1)/(1-Q)
    assert 0 <= s < 1
    # The omitted minus product is in [1-s,1]; the omitted plus
    # product is in [1,exp(2s)]. Their product is in [1-s,exp(2s)].
    two_error = (2*s).expm1()+s
    theta_two = 2*q**(QQ(1)/4)*minus*plus*R(1).add_error(two_error)
    eta = q**(QQ(1)/12)*minus*R(1).add_error(s)
    return theta_two, eta


def controls():
    agree(R(1).gamma(),1)
    agree(R(4).sqrt(),2)
    agree(C(0).elliptic_k().real(),R.pi()/2)
    # q=0: theta3=theta4=1, theta2=0; their tails vanish.
    q=R(0)
    agree(1+2*sum((q**(n*n) for n in range(1,N+1)),R(0)),1)
    agree(2*sum((q**((QQ(n)+QQ(1)/2)**2) for n in range(N+1)),R(0)),0)
    rejected = 0
    for bad in [R('nan'), R(1).add_error(1), R(1)+R(1)/10**80]:
        try:
            agree(bad,1)
        except AssertionError:
            rejected += 1
    assert rejected == 3
    return {'known_value_controls':5,'required_failure_controls':rejected}


def load_table(tid):
    return numberdb.table(tid)


def run(stored=False):
    result = {'controls':controls(), 'mode':'stored' if stored else 'computed',
              'sage':sage_version, 'numberdb':version('numberdb'),
              'working_bits':numberdb.bits(100,losing=96),
              'independent_bits':512,'theta_series_N':N,'product_M':M}
    g=ThetaTwoCM()
    values={int(p['r']):g.value(p,100) for p in g.enumerate()}
    assert not names_its_rings(str(Path(__file__).with_name('generate.py')))
    assert not exactness(values)
    raw={r:v['number'] if isinstance(v,dict) else v for r,v in values.items()}
    result['minimum_accuracy_bits']=min(int(v.accuracy()) for v in raw.values())
    written={r:to_text(v,100) for r,v in raw.items()}
    for r,v in raw.items():
        low,high=decimal_bounds(written[r])
        mid,rad=QQ(v.mid()),QQ(v.rad())
        assert low <= mid-rad and mid+rad <= high, 'serialized interval is too narrow'
    result['measurement_balls']=measure(values)
    result['measurement_written']=measure(written)
    result['serialized_containment_count']=len(written)
    result['entries']=written
    documents={tid:load_table(tid) for tid in ['T295','T296','T297','T299','T504','T505','T506','T25']}
    if stored:
        own=load_table('T507')
        assert list(own['Numbers'])==[str(r) for r in range(1,13)]
        values_in={r:interval(own['Numbers'][str(r)]) for r in raw}
        for r in raw:
            low,high=decimal_bounds(own['Numbers'][str(r)])
            mid,rad=QQ(raw[r].mid()),QQ(raw[r].rad())
            assert low <= mid-rad and mid+rad <= high
    else:
        values_in={r:R(v) for r,v in raw.items()}
    def row(tid,r,normalisation=None):
        x=documents[tid]['Numbers'][str(r)]
        if normalisation is not None:x=x[normalisation]
        return interval(x)
    checked=[]
    for r,two in values_in.items():
        q,t2,t3,t4=theta_series(r)
        p2,eta=products(q)
        for x in [two,t2,t3,t4,p2,eta]:narrow(x)
        assert two > 0 and q > 0 and q < 1
        agree(two,t2)
        agree(two,p2)
        agree(t3,row('T504',r))
        agree(t4,row('T505',r))
        agree(eta,row('T506',r))
        agree(t3**4,two**4+t4**4)
        agree(two*t3*t4,2*eta**3)
        # Repeat shared identities directly against all three stored siblings.
        agree(row('T504',r)**4,two**4+row('T505',r)**4)
        agree(two*row('T504',r)*row('T505',r),2*row('T506',r)**3)
        m=row('T295',r,'parameter')
        k=row('T295',r,'modulus')
        K=row('T296',r)
        agree(m,k*k)
        agree(m,(two/t3)**4)
        agree(K,R.pi()*t3**2/2)
        agree(K,C(m).elliptic_k().real())
        agree(C(1-m).elliptic_k().real(),R(r).sqrt()*K)
        agree(two,m**(QQ(1)/4)*(2*K/R.pi()).sqrt())
        G=row('T297',r,'ramanujan')
        f=row('T297',r,'weber')
        agree(f,R(2)**(QQ(1)/4)*G)
        agree(two,m**(QQ(1)/4)*f**2*eta)
        agree(two,R(2).sqrt()*m**(QQ(1)/4)*G**2*eta)
        checked.append(r)
    gamma=R(QQ(1)/4).gamma()
    agree(values_in[1],gamma/(2*R.pi())**(QQ(3)/4))
    agree(values_in[1],row('T505',1))
    agree(row('T25','1/2'),row('T296',1))
    result['checked_indices']=checked
    result['identities_per_row']=['theta2 series','theta2 product','theta3 sibling','theta4 sibling','eta sibling','Jacobi quartic','theta eta identity','quartic stored siblings','eta stored siblings','k squared equals m','modular lambda ratio','K theta3 identity','K independent evaluation','complementary K','theta2 elliptic identity','Weber normalisation','theta2 Weber identity','theta2 Ramanujan identity']
    result['special_checks']=['square lattice Gamma identity','theta2 equals theta4 at r=1','T25 m=1/2 equals T296 r=1']
    result['passed']=True
    return result

if __name__=='__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    print(json.dumps(run('--stored' in sys.argv or os.environ.get('CHECK_STORED')=='1'),indent=2,ensure_ascii=False))