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

    sage -python checks.py                    # compute and check offline
    sage -python checks.py --stored           # check every stored entry

Requires generate.py and the attached campaign helper check.py. All writes are
stdout only. No publication, review or table-edit request is made.

The principal independent check implements Beukers–Cohen–Mellit (BCM),
Finite hypergeometric functions, Theorem 1.3 (arXiv:1505.02900v2): explicit
finite-field Gauss sums in 384-bit complex balls. It shares no HGM or p-adic
code with the generator. Integrality of the trace plus a unique integer in the
ball determines the exact answer. PARI is another implementation check, with
t=1/z. A quadratic-character convolution and literal elliptic point counts
supply controls outside the Gross–Koblitz implementation.
"""
import json
import os
import sys
from functools import lru_cache
from math import gcd
from pathlib import Path
import numberdb.sage as numberdb
from sage.rings.integer_ring import ZZ
from sage.rings.rational_field import QQ
from sage.rings.complex_arb import ComplexBallField
from sage.rings.finite_rings.finite_field_constructor import GF
from sage.libs.pari import pari
from sage.env import SAGE_VERSION
from generate import (CalabiYauFrobenius, TEMPLATES, FIBRES, R, T, motive,
                      traces, trace_precision)
from check import exactness, measure, agrees_with, names_its_rings

CB = ComplexBallField(384)

# Derive gamma from cyclotomic indices, independently of Sage's HGM object.
def mobius(n):
    result, p = 1, 2
    while p*p <= n:
        if n % p == 0:
            n //= p
            result = -result
            if n % p == 0:
                return 0
        p += 1
    return -result if n > 1 else result

@lru_cache(None)
def gamma_vectors(template):
    gamma = {1:-4}
    for d in map(int,template.split('.')):
        for e in range(1,d+1):
            if d % e == 0:
                gamma[e] = gamma.get(e,0) + mobius(d//e)
    plus = [n for n,c in gamma.items() for _ in range(max(c,0))]
    minus = [n for n,c in gamma.items() for _ in range(max(-c,0))]
    M = QQ(1)
    for n in plus: M *= ZZ(n)**n
    for n in minus: M /= ZZ(n)**n
    assert sum(plus) == sum(minus)
    return plus, minus, M

@lru_cache(None)
def gauss_data(p,f):
    q = p**f
    field = GF(q,'a')
    g = field.multiplicative_generator()
    xs = [field.one()]
    for _ in range(q-2): xs.append(xs[-1]*g)
    assert len(set(xs)) == q-1 and xs[-1]*g == 1
    logs = {x:k for k,x in enumerate(xs)}
    omega = [CB(QQ(2*k)/(q-1)).exppii()
             for k in range(q-1)]
    psi = [CB(QQ(2*int(x.trace()))/p).exppii() for x in xs]
    gauss = [sum((omega[(m*k)%(q-1)]*psi[k] for k in range(q-1)),CB(0))
             for m in range(q-1)]
    # Nonzero controls: trivial Gauss sum=-1; every other Gauss norm=q.
    assert (-1 in gauss[0].real()) and (0 in gauss[0].imag())
    for m,h in enumerate(gauss):
        assert (h.real().is_finite() and h.imag().is_finite()) and h.real().rad() < QQ(1)/10**90
        norm = h*h.conjugate()
        assert ((1 if m == 0 else q) in norm.real())
        assert (0 in norm.imag()) and norm.real().rad() < QQ(1)/10**85
    return field, logs, omega, gauss

@lru_cache(None)
def bcm_trace(template,z_string,p,f):
    plus, minus, M = gamma_vectors(template)
    q, n = p**f, p**f-1
    F, logs, omega, G = gauss_data(p,f)
    arg = F(((-1)**sum(minus))*QQ(z_string)/M)
    l = logs[arg]
    s0 = min(len(plus),len(minus))
    total = CB(0)
    for m in range(n):
        sm = min(sum((k*m)%n == 0 for k in plus),
                 sum((k*m)%n == 0 for k in minus))
        term = CB(QQ(q)**(sm-s0)) * omega[(l*m)%n]
        for k in plus: term *= G[(k*m)%n]
        for k in minus: term *= G[(-k*m)%n]
        total += term
    total *= QQ((-1)**(len(plus)+len(minus)))/QQ(1-q)
    assert (total.real().is_finite() and total.imag().is_finite()) and (0 in total.imag())
    assert total.real().rad() < QQ(1)/10**75
    assert total.imag().rad() < QQ(1)/10**75
    value = ZZ(total.real())
    assert value*value <= 16*ZZ(p)**(3*f)
    return ZZ(value)

def bcm_polynomial(key):
    template,z,p = key
    a = bcm_trace(template,z,p,1)
    c = bcm_trace(template,z,p,2)
    assert (a*a-c)%2 == 0
    return R([1,-a,(a*a-c)//2,-ZZ(p)**3*a,ZZ(p)**6])

@lru_cache(None)
def pari_template(key):
    return pari('hgminit([%s])'%key.replace('.',','))

def pari_polynomial(key):
    template,z,p = key
    poly = pari.hgmeulerfactor(pari_template(template),pari(QQ(1)/QQ(z)),p)
    return R([ZZ(poly.polcoef(i)) for i in range(5)])

def quadratic_controls():
    results = []
    for p in [3,5,7]:
        for f in [1,2]:
            F = GF(p**f,'b')
            chi = lambda x: 0 if not x else (1 if x.is_square() else -1)
            assert sum(chi(x*x-F(1)) for x in F) == -1
            # Finite Euler convolution, normalized by J(phi,phi)^3.
            # J(phi,phi)=-phi(-1), from the elementary quadratic Jacobi sum.
            jacobi = sum(chi(x*(1-x)) for x in F)
            assert jacobi == -chi(F(-1))
            for z_string in FIBRES:
                z = F(QQ(z_string))
                nonzero = [x for x in F if x and x != 1]
                pairs = {}
                for x in nonzero:
                    for y in nonzero:
                        k=x*y
                        pairs[k]=pairs.get(k,0)+chi(x*(1-x)*y*(1-y))
                S=sum(v*chi(w*(1-w)*(1-z*k*w)) for k,v in pairs.items()
                      for w in nonzero)
                a=ZZ(jacobi**3*S)
                assert a == bcm_trace('2.2.2.2',z_string,p,f)
                # Outside rank four: BCM Theorem 1.2, literal point count.
                count=1+sum(1+chi(x*(x-1)*(x-z)) for x in F)
                h2=pari.hgmeulerfactor(pari('hgminit([2,2])'),
                                       pari(QQ(1)/QQ(z_string)),p)
                ap=-ZZ(h2.polcoef(1))
                predicted=ap if f==1 else ap*ap-2*p
                assert chi(F(-1))*(p**f+1-count) == predicted
                if p==5 and f==1 and z_string=='2': assert count==8
                results.append({'p':p,'f':f,'z':z_string,'trace4':int(a),
                                'legendre_points':int(count)})
    return results

def load_values(stored=False):
    gen=CalabiYauFrobenius()
    parameters=list(gen.enumerate())
    if stored:
        from numberdb._generate import _stored_entries
        doc=numberdb.table('T674')
        entries=_stored_entries(doc)
        expected={','.join(str(p[k]) for k in gen.parameters) for p in parameters}
        assert set(entries)==expected, (len(entries),len(expected))
        return {(p['template'],p['z'],int(p['p'])):
                R(entries[','.join(str(p[k]) for k in gen.parameters)])
                for p in parameters}
    return {(p['template'],p['z'],int(p['p'])):gen.value(p,100) for p in parameters}

def main(stored=False):
    print('VERSIONS',json.dumps({'sage':SAGE_VERSION,'pari':str(pari('version()')),
                                'numberdb':numberdb.__version__ if hasattr(numberdb,'__version__') else 'client-source'}),flush=True)
    print('MODE','stored' if stored else 'offline',flush=True)
    # Published, nonzero full-polynomial controls precede all comparisons.
    assert bcm_polynomial(('5','-1',3)) == R([1,5,45,135,729])
    assert bcm_polynomial(('2.2.2.2','-1',5)) == R([1,4,-130,500,15625])
    assert pari_polynomial(('5','-1',3)) == R([1,5,45,135,729])
    assert pari_polynomial(('5','1/2',7)) == R([1,-35,805,-12005,117649])
    assert pari_polynomial(('5','2',7)) == R([1,10,420,3430,117649])
    print('CONTROLS known polynomials and asymmetric coordinate passed',flush=True)
    values=load_values(stored)
    assert len(values)==390
    assert not exactness(values)
    assert not names_its_rings(str(Path(__file__).with_name('generate.py')))
    print('MEASURE',json.dumps(measure(values),default=str),flush=True)
    assert not agrees_with(values,pari_polynomial)
    print('PARI all 390 complete polynomials agree at reciprocal coordinate',flush=True)
    for template in TEMPLATES:
        subset={k:v for k,v in values.items() if k[0]==template}
        failures=agrees_with(subset,bcm_polynomial)
        assert not failures, failures[:3]
        print('BCM',template,len(subset),'polynomials pass',flush=True)
    for key,P in values.items():
        template,z,p=key
        a=-P[1]
        assert P.degree()==4 and P[0]==1 and P[4]==p**6
        assert P[3]==p**3*P[1] and P[2]%p==0
        assert 2*P[2]==a*a-bcm_trace(template,z,p,2)
        assert a*a<=16*p**3 and abs(P[2])<=6*p**3
        # Exact weight-three Weil test: X^-2 Q(X)=U²-aU+(b-2p³),
        # U=X+p³/X. Both real U roots must lie in [-2sqrt(p³),2sqrt(p³)].
        disc=a*a-4*(P[2]-2*p**3)
        assert disc>=0
        gap=16*p**3-a*a-disc
        assert gap>=0 and gap*gap>=4*a*a*disc
        assert P==motive(template).euler_factor(QQ(z),ZZ(p))
        for f in (1,2):
            n=trace_precision(ZZ(p),f)
            assert bcm_trace(template,z,p,f)==motive(template).padic_H_value(
                ZZ(p),f,QQ(z),prec=n+3,cache_p=False)
    print('IDENTITIES shape, Newton reconstruction, exact Weil test, Hodge bound, Sage standard factors and extra-precision traces passed',flush=True)
    controls=quadratic_controls()
    print('QUADRATIC_CONTROLS',json.dumps(controls),flush=True)
    print('VALUES_JSON',json.dumps([{'params':{'template':k[0],'z':k[1],'p':str(k[2])},
                                    'number':str(v),'coefficients':list(map(int,v.list()))}
                                   for k,v in values.items()]),flush=True)
    print('PASS',len(values),'polynomials; 780 explicit Gauss-sum trace enclosures; 18 quadratic convolution and elliptic point-count controls',flush=True)

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