"""Independent checks for numberdb.org/T701.

Run: sage -python checks.py
Requires generate.py beside this file and mpmath (included with SageMath).
Outputs JSON evidence; no API writes. With --stored, authenticates through
NUMBERDB_API_KEY or NUMBERDB_KEY_FROM_STDIN=1 and checks the live draft too.

The independent Fourier calculation expands log|1-z| and sums the q-powers
analytically. It shares no evaluation code with the paired-log generator.
It is an agreement check, not the precision certificate. The latter is the
Arb computation with explicit tails in generate.py. The Mahler integral is
independent of the logarithmic J definition; square-lattice rotation follows
from reindexing the Eisenstein-Kronecker lattice. A bounded lattice truncation
fixes its sign against this nonzero integral.
"""
import json
import math
from decimal import Decimal
import os
import sys
from collections import Counter
import numberdb.sage as numberdb
from sage.env import SAGE_VERSION
from sage.rings.real_arb import RealBallField
from sage.rings.complex_arb import ComplexBallField
from sage.rings.rational_field import QQ
from numberdb._write import to_text
from numberdb._generate import _stored_entries
import mpmath as mp
import generate


def fourier(delta, N, r, s, dps=170):
    with mp.workdps(dps):
        b=delta % 2
        L=mp.pi*mp.sqrt(-delta)
        a=mp.mpf(r)/N
        rho=mp.exp(-L)
        q=(-1)**b*rho
        z=mp.exp(-L*a+2j*mp.pi*(mp.mpf(s)/N-mp.mpf(b*r)/(2*N)))
        w=q/z
        K=int(mp.ceil((dps+20)*mp.log(10)/(L*min(a,1-a))))
        zk,wk,qk=z,w,q
        total=mp.mpf(0)
        for k in range(1,K+1):
            t=qk/(1-qk)**2
            total += mp.re((zk*(a/(1-qk)+t)-wk*((1-a)/(1-qk)+t))/k)
            zk*=z; wk*=w; qk*=q
        return +(L*total+L**2*(a**3-3*a*a/2+a/2)/3)


def bloch_square(N, r, s):
    with mp.workdps(170):
        a=mp.mpf(r % N)/N
        q=mp.exp(-2*mp.pi)
        x=mp.exp(-2*mp.pi*a+2j*mp.pi*(mp.mpf(s % N)/N))
        def bw(z):
            return mp.im(mp.polylog(2,z)+mp.log(abs(z))*mp.log(1-z))
        total=bw(x)
        for k in range(1,70):
            total+=bw(x*q**k)-bw(q**k/x)
        return +total


def finite_narrow(v, target_bits=350):
    assert v.is_finite(), 'Nonfinite check result'
    assert v.accuracy() >= target_bits, ('Too little relative accuracy',v.accuracy())


def mahler_and_lattice():
    R=RealBallField(470); C=ComplexBallField(470)
    control=C.integral(lambda z,analytic:z*z,0,1)
    assert (control.real().is_finite() and control.imag().is_finite()) and (control.real()-QQ(1)/3).contains_zero()
    finite_narrow(control.real())
    rho=(-2*R.pi()).exp()
    # phi(q)=1+2 sum q^(n^2); n>=M+1 is bounded by a geometric series.
    M=12
    tail=2*rho**((M+1)**2)/(1-rho**(2*M+3))
    plus=(1+sum(2*rho**(n*n) for n in range(1,M+1))).add_error(tail)
    minus=(1+sum(2*(-rho)**(n*n) for n in range(1,M+1))).add_error(tail)
    kappa=4*plus**2/minus**2
    assert kappa>4
    def integrand(t,analytic):
        z=(C(kappa)+2*(2*C.pi()*t).cos())/2
        # Sage's installed arccosh analytic guard checks -1, not +1.
        # Restrict the whole image ball to Re(z)>1 to certify holomorphy.
        if analytic and not z.real()>1:
            return C('nan')
        return z.arccosh()
    integral=C.integral(integrand,0,1)
    finite_narrow(integral.real())
    assert integral.imag().contains_zero()
    # Independent constant-term expansion of the Mahler integral for kappa>4.
    # c_n=binom(2n,n)^2/(2n*kappa^(2n)); c_(n+1)/c_n < 16/kappa^2.
    ratio=16/kappa**2
    term=2/kappa**2
    total=R(0)
    n=1
    while term/(1-ratio)>R(2)**-450:
        total+=term
        term*=R(4*n*(2*n+1)**2)/R((n+1)**3)/kappa**2
        n+=1
    series_terms=n-1
    expansion=(kappa.log()-total).add_error(term/(1-ratio))
    assert integral.real().overlaps(expansion)
    finite_narrow(expansion)
    expected=R.pi()*integral.real()/4
    actual=generate.evaluate(-4,4,1,0)
    assert actual.overlaps(expected), ('Mahler disagreement',str(actual),str(expected),str(actual-expected),str(kappa),str(integral))
    finite_narrow(expected)
    # On tau=i, |m*i+n| >= max(|m|,|n|). Shell k has 8k points,
    # so the omitted K tail is <= 8 sum_{k>M}1/k^2 <= 8/M.
    radius=80
    Csmall=ComplexBallField(90); Rsmall=RealBallField(90)
    lattice={}
    for name,xi,eta in [('D',QQ(0),QQ(1)/4),('J',QQ(1)/4,QQ(0))]:
        total=Csmall(0)
        for m in range(-radius,radius+1):
            for n in range(-radius,radius+1):
                if m==n==0:continue
                lam=Csmall(n,m)
                phase=(2*Csmall.gen(0)*Csmall.pi()*(n*xi-m*eta)).exp()
                total+=phase/(lam**2*lam.conjugate())
        component=(-total.real() if name=='D' else total.imag())/Rsmall.pi()
        enclosed=component.add_error(Rsmall(8)/(radius*Rsmall.pi()))
        assert enclosed.is_finite() and enclosed>0
        assert enclosed.overlaps(expected)
        assert not (-enclosed).overlaps(expected)
        lattice[name]={'center':str(component),'with_tail':str(enclosed),'tail_bound':str(Rsmall(8)/(radius*Rsmall.pi()))}
    return {'control':str(control),'kappa':str(kappa),'mahler_integral':str(integral),'mahler_constant_term_series':str(expansion),'mahler_series_terms':series_terms,'integration_branch_guard':'Reject every analytic query unless Re(z)>1; installed Sage arccosh checks -1 instead of +1.','pi_times_integral_over_4':str(expected),'J_i_over_4':str(actual),'lattice_square_radius':radius,'lattice':lattice}


def main(stored=False):
    mp.mp.dps=170
    rows=list(generate.parameters())
    assert len(rows)==362 and len({tuple(p.values()) for p in rows})==362
    counts=Counter(p['Delta'] for p in rows)
    assert all(counts[d]==(26 if d%2==0 else 29) for d in generate.DISCRIMINANTS)
    # Independent enumeration: form group elements, their exact additive orders,
    # and inversion orbits without using the generator's gcd test.
    independent=[]
    for d in generate.DISCRIMINANTS:
        for n in range(3,7):
            for r in range(n):
                for s in range(n):
                    order=next(k for k in range(1,n+1) if (k*r)%n==0 and (k*s)%n==0)
                    if order!=n or (r,s)!=min((r,s),((-r)%n,(-s)%n)):continue
                    if (2*r)%n==0 and (d%2*r)%n==0:continue
                    independent.append((d,n,r,s))
    assert independent==[tuple(p.values()) for p in rows]
    retained=10000; terms=[]; worst=mp.mpf(0); rotation=mp.mpf(0)
    numbers=[]; values={}; references={}
    for index,p in enumerate(rows):
        args=tuple(p.values()); v,m,tail=generate.evaluate(*args,metadata=True)
        zero = args == (-3,3,1,2)
        if zero:
            assert v.is_finite() and v.contains_zero() and RealBallField(470)(v.rad())<RealBallField(470)(10)**-120
        else:
            finite_narrow(v); retained=min(retained,int(v.accuracy()))
        terms.append(m)
        high=generate.evaluate(*args,digits=130)
        if not zero:finite_narrow(high,450)
        assert v.overlaps(high)
        ref=fourier(*args);references[args]=ref
        approx=mp.mpf(str(v.mid()))
        difference=abs(approx-ref) if zero else abs((approx-ref)/ref)
        if not zero:worst=max(worst,difference)
        assert difference<mp.mpf('1e-120'),(args,difference)
        d,n,r,s=args
        inverted=generate.evaluate(d,n,-r,-s)
        conjugated=generate.evaluate(d,n,r,(d%2)*r-s)
        assert v.overlaps(-inverted) and v.overlaps(conjugated)
        if d==-4:
            other=bloch_square(n,-s,r)
            diff=abs((ref-other)/ref);rotation=max(rotation,diff)
            assert diff<mp.mpf('1e-150')
        text=to_text(0 if zero else v,100)
        assert len(text)<150
        if not zero:
            ulp=mp.mpf(10)**Decimal(text).as_tuple().exponent
            assert abs(mp.mpf(text)-ref)<=ulp
        numbers.append({'params':{k:str(t) for k,t in p.items()},'number':text})
        values[args]=v
        if index%80==0:print('Checked',index+1,'of',len(rows),flush=True)
    zero_count=0
    for d in generate.DISCRIMINANTS:
        for n in range(1,7):
            for r in range(n):
                for s in range(n):
                    if math.gcd(math.gcd(n,r),s)!=1:continue
                    if (r,s)>((-r)%n,(-s)%n):continue
                    if n<=2 or ((2*r)%n==0 and (d%2*r)%n==0):
                        v=generate.evaluate(d,n,r,s)
                        assert v.is_finite() and v.contains_zero()
                        assert RealBallField(470)(v.rad())<RealBallField(470)(10)**-120
                        zero_count+=1
    print('Running external Mahler and signed lattice controls',flush=True)
    external=mahler_and_lattice()
    stored_evidence=None
    if stored:
        if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':numberdb.configure(api_key=sys.stdin.read().strip())
        doc=numberdb.table('T701')
        # The API can normalize submitted flat records to a nested tree.
        records=[{'params':dict(zip(generate.Table.parameters,key.split(','))), 'number':value}
                 for key,value in _stored_entries(doc).items()]
        assert len(records)==len(rows)
        maxerr=mp.mpf(0)
        for record in records:
            key=tuple(int(record['params'][k]) for k in generate.Table.parameters)
            ref=references[key]; text=record['number'];num=mp.mpf(text)
            if key==(-3,3,1,2):
                assert text=='0' and abs(ref)<mp.mpf('1e-160')
                rel=mp.mpf(0)
            else:
                rel=abs((num-ref)/ref)
                ulp=mp.mpf(10)**Decimal(text).as_tuple().exponent
                assert abs(num-ref)<=ulp
            maxerr=max(maxerr,rel);assert rel<mp.mpf('1e-98')
            if key!=(-3,3,1,2):
                d,n,r,s=key
                neg=generate.evaluate(d,n,-r,-s,digits=130)
                conj=generate.evaluate(d,n,r,(d%2)*r-s,digits=130)
                finite_narrow(neg,450);finite_narrow(conj,450)
                assert abs(num+mp.mpf(str(neg.mid())))<=ulp
                assert abs(num-mp.mpf(str(conj.mid())))<=ulp
        # Check square rotation again on the actual stored decimal values.
        for record in records:
            p={k:int(v) for k,v in record['params'].items()}
            if p['Delta']==-4:
                ref=bloch_square(p['N'],-p['s'],p['r'])
                assert abs((mp.mpf(record['number'])-ref)/ref)<mp.mpf('1e-98')
        stored_evidence={'entries_independently_checked':len(records),'stored_inversion_and_conjugation_rows':len(records),'square_rotation_rows':26,'max_relative_decimal_error':mp.nstr(maxerr,12)}
    result={'table':'T701','skill_manifest':'agents/runs/20261110T005551Z-build-skills/manifest.json','skill_sha256':'886390324dd1971e381a43cf3b46b822ab0bc67d62f31d8481e29ec79556746b','numberdb_version':numberdb.__version__,'sage_version':SAGE_VERSION,'python_version':sys.version.split()[0],'mpmath_version':mp.__version__,'entries':len(rows),'exact_zero_rows':[['-3','3','1','2']],'counts_by_discriminant':dict(counts),'minimum_relative_accuracy_bits':retained,'term_count_range':[min(terms),max(terms)],'max_relative_fourier_difference':mp.nstr(worst,12),'max_relative_square_rotation_difference':mp.nstr(rotation,12),'omitted_zero_representatives_checked':zero_count,'inversion_and_conjugation_rows':len(rows),'external_controls':external,'stored':stored_evidence}
    print('EVIDENCE_JSON='+json.dumps(result,sort_keys=True))
    if not stored:print('NUMBERS_JSON='+json.dumps(numbers))

if __name__=='__main__':main('--stored' in sys.argv or os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1')