"""Independent CM-orbit and Fourier checks for numberdb.org/T687.

    sage -python checks.py                   # checks locally generated values
    sage -python checks.py --stored          # checks values read from T687

Install numberdb first. Stored draft reads require NUMBERDB_API_KEY.
The numerical CM algorithm shares no Jacobi coefficient arithmetic with the
exact generator. It enumerates reduced forms, their projective level
structures and stabilizer orbits, then evaluates Dedekind eta with Arb.
"""
import json
import os
import sys
from math import gcd, isqrt
from functools import lru_cache
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 generate import FrickeTraces, LEVELS, MAX_D, basis, solve, singular
from check import exactness, measure, names_its_rings

BITS=640
F=ComplexBallField(BITS)
I=F(0,1)
PRIMES={2:(2,),3:(3,),5:(5,),6:(2,3),7:(7,),10:(2,5),13:(13,)}
FIRST={2:4372,3:783,5:134,6:79,7:51,10:22,13:12}


def narrow_contains(ball, exact):
    assert ball.real().is_finite() and ball.imag().is_finite(), 'nonfinite ball'
    assert ball.real().rad()<QQ(1)/10**100 and ball.imag().rad()<QQ(1)/10**100, 'wide ball'
    assert exact in ball.real() and ball.imag().contains_zero(), ('disagrees',exact,ball)


def star(N,tau):
    if N in (6,10):
        p=N//2; e=6 if N==6 else 4
        r=(tau.modular_eta()*(p*tau).modular_eta()/((2*tau).modular_eta()*(N*tau).modular_eta()))**e
        return r+e+ZZ(2)**e/r
    e=24//(N-1)
    r=(tau.modular_eta()/(N*tau).modular_eta())**e
    return r+e+ZZ(N)**(e//2)/r


def reduced_forms(d):
    for a in range(1,isqrt(d//3)+1):
        for b in range(-a,a+1):
            if (b*b+d)%(4*a): continue
            c=(b*b+d)//(4*a)
            if a>c or ((abs(b)==a or a==c) and b<0): continue
            yield (a,b,c)


def projective(N,u,v):
    return min(((k*u)%N,(k*v)%N) for k in range(N) if gcd(k,N)==1)


@lru_cache(None)
def level_structures(N):
    pairs=sorted({projective(N,u,v) for u in range(N) for v in range(N) if gcd(gcd(u,v),N)==1})
    expected=N
    for p in PRIMES[N]: expected=expected*(p+1)//p
    assert len(pairs)==expected
    return pairs


def egcd(a,b):
    if b==0: return a,1,0
    g,x,y=egcd(b,a%b)
    return g,y,x-(a//b)*y


def lift(N,pair):
    for h in range(N+1):
        for k in range(N+1):
            u,v=pair[0]+h*N,pair[1]+k*N
            if gcd(u,v)==1:
                g,t,minus_s=egcd(u,v)
                assert u*t+v*minus_s==1
                return (u,-minus_s,v,t)
    raise AssertionError('no primitive lift')


def gamma0_orbits(N,d):
    reps=[]
    for a,b,c in reduced_forms(d):
        allowed={p for p in level_structures(N) if (a*p[0]**2+b*p[0]*p[1]+c*p[1]**2)%N==0}
        w=3 if a==b==c else 2 if a==c and b==0 else 1
        while allowed:
            p=min(allowed); orbit={p}; u,v=p
            for _ in range(w-1):
                u,v=(-v,u+v) if w==3 else (-v,u)
                orbit.add(projective(N,u,v))
            assert orbit<=allowed
            allowed-=orbit
            u,s,v,t=lift(N,p)
            A=a*u*u+b*u*v+c*v*v
            B=2*a*u*s+b*(u*t+s*v)+2*c*v*t
            C=a*s*s+b*s*t+c*t*t
            assert B*B-4*A*C==-d and A%N==0
            w0=w//len(orbit)
            reps.append({'base':[a,b,c],'level':list(p),'form':[A,B,C],
                         'stabilizer_gamma0':w0,'star_sum_weight':str(QQ(1)/(w0*2**len(PRIMES[N])))})
    return reps


def cm_check(values):
    evidence=[]; maxrad=F(0).real().rad(); compared=0; residue_compared=0
    for N in LEVELS:
        for d in range(1,MAX_D+1):
            if (N,1,d) not in values: continue
            reps=gamma0_orbits(N,d)
            sums={1:F(0),2:F(0)}; by_residue={}
            for rep in reps:
                a,b,c=rep['form']; w0=rep['stabilizer_gamma0']
                tau=(-F(b)+I*F(d).sqrt())/(2*a)
                j=star(N,tau)
                fabers={1:j,2:j*j-2*FIRST[N]}
                h=b%(2*N)
                by_residue.setdefault(h,{1:F(0),2:F(0)})
                for m,val in fabers.items():
                    sums[m]+=val/QQ(w0*2**len(PRIMES[N]))
                    by_residue[h][m]+=val/QQ(w0)
            expected_h={h for h in range(2*N) if (h*h+d)%(4*N)==0}
            assert set(by_residue)==expected_h
            for m in (1,2):
                exact=values[N,m,d]
                narrow_contains(sums[m],exact); compared+=1
                maxrad=max(maxrad,sums[m].real().rad(),sums[m].imag().rad())
                factor=2**sum(d%p==0 for p in PRIMES[N])
                for h,part in by_residue.items():
                    narrow_contains(part[m],exact*factor); residue_compared+=1
            evidence.append({'N':N,'d':d,'gamma0_representatives':reps,
                             'trace1_ball':str(sums[1]),'trace2_ball':str(sums[2])})
        print('CM level',N,'passed',flush=True)
    return evidence,{'cm_entries_checked':compared,'fixed_residue_sums_checked':residue_compared,
                     'bits':BITS,'maximum_absolute_radius':str(maxrad)}


def controls():
    narrow_contains(F(0,1).modular_j(),1728)
    narrow_contains(star(2,F(0,1)),544)
    narrow_contains(star(2,(F(1)+I)/2),-104)
    narrow_contains(star(2,I/F(2).sqrt()),152)
    narrow_contains(star(3,(-F(3)+I*F(3).sqrt())/6),-42)
    rejected=0
    for ball,answer in [(F(1),2),(F('nan'),0),(F('0 +/- 1'),0)]:
        try: narrow_contains(ball,answer)
        except AssertionError: rejected+=1
    assert rejected==3
    assert list(reduced_forms(3))==[(1,1,1)]
    assert list(reduced_forms(4))==[(1,0,1)]
    assert list(reduced_forms(12))==[(1,0,3),(2,2,2)]
    assert gamma0_orbits(2,4)[0]['stabilizer_gamma0']==2
    print('Known CM values, enumeration and three rejection controls passed',flush=True)


def series_mul(a,b):
    return [sum(a[k]*b[n-k] for k in range(n+1)) for n in range(len(a))]


def eta_product(exponents,length):
    """Exact Euler product, independent of the Jacobi generator."""
    from math import comb
    out=[ZZ(1)]+[ZZ(0)]*(length-1)
    for delta,e in exponents.items():
        for step in range(delta,length,delta):
            f=[ZZ(0)]*length
            for j in range((length-1)//step+1):
                if e>=0:
                    if j<=e: f[j*step]=ZZ((-1)**j*comb(e,j))
                else: f[j*step]=ZZ(comb(-e+j-1,j))
            out=series_mul(out,f)
    return out


def fourier(N,limit=50):
    if N in (6,10):
        p=N//2; e=6 if N==6 else 4
        ex={1:e,p:e,2:-e,N:-e}; const=e; multiplier=2**e
        un_ex=({2:3,3:9,1:-3,6:-9} if N==6 else {2:1,5:5,1:-1,10:-5})
    else:
        e=24//(N-1); ex={1:e,N:-e}; const=e; multiplier=N**(e//2)
        un_ex=ex
    a=eta_product(ex,limit+2)
    b=eta_product({k:-v for k,v in ex.items()},limit+2)
    un=eta_product(un_ex,limit+2)
    jstar={n:a[n+1]+multiplier*b[n-1] for n in range(1,limit+1)}
    j0={n:un[n+1] for n in range(1,limit+1)}
    assert jstar[1]==FIRST[N]
    assert a[1]+const==0
    return jstar,j0


def sigma_excluding(N,n):
    if n.denominator()!=1 or n<=0: return QQ(0)
    n=int(n)
    return sum(ZZ(d) for d in range(1,n+1) if n%d==0 and d%N)


# Use the documented entries format rather than client-private parsers.
def records(block,keys,params=None):
    params={} if params is None else params
    if isinstance(block,list):
        for row in block: yield row
    elif isinstance(block,dict) and 'equals' in block:
        return
    elif not keys:
        yield {'params':params,'number':block['number'] if isinstance(block,dict) else block}
    else:
        for k,v in block.items():
            if k in ('comment','param-latex'): continue
            yield from records(v,keys[1:],dict(params,**{keys[0]:str(k)}))

def identities(values):
    checks=0; companion=0
    # Optional frozen published companion values, not inputs to the generator.
    path=Path(__file__).with_name('T652.json')
    comp=json.loads(path.read_text()) if path.exists() else numberdb.table('T652')
    compvalues={(r['params']['g'],int(r['params']['n'])):QQ(r['number']) for r in records(comp['Numbers'],list(comp['Parameters'])) if 'number' in r}
    for N in LEVELS:
        star_coeff,un_coeff=fourier(N)
        for n in range(1,31):
            for label,calc in [(str(N)+'A',star_coeff),(str(N)+('E' if N in (6,10) else 'B'),un_coeff)]:
                if (label,n) in compvalues:
                    assert calc[n]==compvalues[label,n],('companion',label,n); companion+=1
        def t(d):
            if d==0: return QQ(5 if N==2 else 5 if N in (6,10) else 3)/(2 if N in (6,10) else 1)
            if d in (-1,-4): return QQ(-1 if d==-1 else -2)
            if d<0: return QQ(0)
            return values.get((N,2,d),QQ(0))
        def trace_sum(k):
            assert 4*k<=MAX_D
            return sum(t(4*k-r*r) for r in range(-isqrt(4*k+4),isqrt(4*k+4)+1))
        for n in range(1,51):
            if N in (6,10):
                if N==6:
                    correction=7*sigma_excluding(6,QQ(n))+26*sigma_excluding(3,QQ(n)/2)-3*sigma_excluding(2,QQ(n)/3)
                else: correction=4*sigma_excluding(10,QQ(n))+12*sigma_excluding(5,QQ(n)/2)
            else:
                correction=QQ(24*(3-(N if N==2 else 0)))/(N-1)*sigma_excluding(N,QQ(n))
            assert 2*n*un_coeff[n]==trace_sum(n)+correction,('unstarred identity',N,n)
            checks+=1
        for n in range(1,MAX_D//(4*N)+1):
            rhs=0
            for e in range(1,N+1):
                if N%e==0:
                    sign=(-1)**sum(e%p==0 for p in PRIMES[N])
                    rhs+=sign*trace_sum(e*n)
            assert 2*n*star_coeff[n]==rhs,('starred identity',N,n)
            checks+=1
    return {'coefficient_identities':checks,'published_coefficient_controls':companion}


def main():
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    controls()
    g=FrickeTraces()
    if '--stored' in sys.argv or os.environ.get('CHECK_STORED')=='1':
        doc=numberdb.table('T687')
        values={tuple(int(r['params'][k]) for k in g.parameters):QQ(r['number'])
                for r in records(doc['Numbers'],list(doc['Parameters']))}
        expected={tuple(int(p[k]) for k in g.parameters) for p in g.enumerate()}
        assert set(values)==expected, 'Stored enumeration differs from generator'
    else:
        values={tuple(int(p[k]) for k in g.parameters):g.value(p,100) for p in g.enumerate()}
    assert not exactness(values)
    assert not names_its_rings(str(Path(__file__).with_name('generate.py')))
    assert values[2,1,4]==-26 and values[2,1,7]==-23 and values[2,1,8]==76 and values[3,1,3]==-7
    print('Exact rational controls passed',flush=True)
    orbit,evidence=cm_check(values)
    evidence.update(identities(values))
    evidence.update({'table':'T687','entries':len(values),'measurement':measure(values),
                     'skill_manifest':'agents/runs/20261109T100832Z-build-skills/manifest.json',
                     'skill_sha256':'886390324dd1971e381a43cf3b46b822ab0bc67d62f31d8481e29ec79556746b',
                     'known_controls':5,'negative_controls':3})
    print('NUMERICAL_EVIDENCE '+json.dumps(evidence),flush=True)
    print('ORBIT_EVIDENCE '+json.dumps(orbit),flush=True)
    print('VALUES '+json.dumps([{'params':dict(zip(g.parameters,map(str,k))),'number':str(v)} for k,v in values.items()]),flush=True)

if __name__=='__main__': main()