"""O’Nan moonshine coefficients -- numberdb.org/T672.

Run with SageMath (install numberdb once with sage -pip install numberdb):
    sage -python generate.py                 # verify stored entries
    sage -python generate.py --fill          # fill draft, API key on stdin

The module and normalization are DMO, arXiv:1702.03516v4, Theorems 3.1,
4.1. Positive coefficients use Table 5.1, Proposition 5.1, and Appendix D.
Every positive definite integral form is included, also imprimitive forms.
The trace uses PSL2 stabilizers. A full projective-line coset sum divided
by the level-one stabilizer equals the level-N orbit sum with its correct
stabilizers, by orbit-stabilizer. The integrality theorem and a finite narrow
Arb enclosure identify each output integer, without heuristic rounding.
See argument-selection.md for the deliberate class and grade bounds.
"""
import sys
from functools import lru_cache
from math import gcd,isqrt
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

CLASSES=('1A','2A','3A','4A','4B','5A','6A','7A','7B')
MAX_GRADE=200
WORKING_BITS=1024
# q coefficients of the normalized Hauptmoduln in DMO Table 5.1.
FIRST={1:196884,2:276,3:54,4:20,5:9,6:6,7:2}
CORRECTIONS={1:{},2:{1:12,2:-12},3:{1:12,3:-12},
             4:{2:12,4:-12},5:{1:6,5:-6},
             6:{1:-12,2:8,3:QQ(21)/2,6:-QQ(13)/2},7:{1:4,7:-4}}


def reduced_forms(D):
    """Unique reduced SL2 classes, including imprimitive ones."""
    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 c<a or ((abs(b)==a or a==c) and b<0):
                continue
            yield a,b,c


@lru_cache(None)
def cosets(N):
    """SL2(Z)/Gamma0(N), indexed by first column in P1(Z/N)."""
    if N==1:
        return ((1,0,0,1),)
    units=[u for u in range(N) if gcd(u,N)==1]
    seen=set();result=[]
    for r in range(N+1):
        for t in range(N+1):
            if gcd(r,t)!=1:
                continue
            key=min(((v*r)%N,(v*t)%N) for v in units)
            if key in seen:
                continue
            seen.add(key)
            if t==0:
                s,u=0,1
            else:
                u=pow(r,-1,t) if t>1 else 0
                s=(r*u-1)//t
            assert r*u-s*t==1
            result.append((r,s,t,u))
    expected=N
    for p in range(2,N+1):
        if N%p==0 and all(p%j for j in range(2,isqrt(p)+1)):
            expected=expected*(p+1)//p
    assert len(result)==expected
    return tuple(result)


def hauptmodul(N,z):
    if N==1:
        return z.modular_j()-744
    if N==6:
        return z.modular_eta()**5*(3*z).modular_eta()/((2*z).modular_eta()*(6*z).modular_eta()**5)+5
    e={2:24,3:12,4:8,5:6,7:4}[N]
    return (z.modular_eta()/(N*z).modular_eta())**e+e


def trace_data(D,bits=WORKING_BITS):
    C=ComplexBallField(bits)
    H={N:QQ(0) for N in range(1,8)}
    T1={N:C(0) for N in range(1,8)}
    T2={N:C(0) for N in range(1,8)}
    for a,b,c in reduced_forms(D):
        w=3 if a==b==c else 2 if b==0 and a==c else 1
        tau=(-C(b)+C(0,1)*C(D).sqrt())/(2*a)
        for N in range(1,8):
            for r,s,t,u in cosets(N):
                if (a*r*r+b*r*t+c*t*t)%N:
                    continue
                H[N]+=QQ(1)/w
                z=(u*tau-s)/(-t*tau+r)
                J=hauptmodul(N,z)
                T1[N]+=J/w
                T2[N]+=(J*J-2*FIRST[N])/w
    return H,T1,T2


def isolate_integer(ball):
    """Require a finite enclosure tighter than 100 absolute decimal places."""
    assert ball.real().is_finite() and ball.imag().is_finite() and ball.imag().contains_zero(), 'nonreal or nonfinite trace'
    bound=ball.real().parent()(10)**-110
    assert ball.real().parent()(ball.real().rad())<bound and ball.real().parent()(ball.imag().rad())<bound, 'insufficient precision'
    # ZZ(ball.real()) succeeds only if its interval contains a unique integer.
    z=ZZ(ball.real())
    assert z in ball.real()
    return z


@lru_cache(None)
def coefficients(D,bits=WORKING_BITS):
    H,T1,T2=trace_data(D,bits)
    out={}
    for N in range(1,8):
        v=(T2[N]-T1[N//gcd(N,2)])/2
        for level,c in CORRECTIONS[N].items():
            v+=c*H[level]
        out[N]=(isolate_integer(v),v)
    return out


class ONanCoefficients(numberdb.Generator):
    table='T672'
    parameters=('g','n')
    type='Z'
    rigour='exact'
    digits=100
    files=('generate.py','checks.py','source-controls.json','argument-selection.md',
           'numerical-evidence.md','skill-evidence.json','companion-controls.json','gap-ON-excerpt.txt')

    def enumerate(self):
        for g in CLASSES:
            for n in range(3,MAX_GRADE+1):
                if n%4 in (0,3):
                    yield {'g':g,'n':str(n)}

    def value(self,params,digits):
        return coefficients(int(params['n']))[int(params['g'][:-1])][0]


if __name__=='__main__':
    if '--fill' in sys.argv:
        numberdb.configure(api_key=sys.stdin.read().strip())
        print(ONanCoefficients().publish(message='Fill the draft with certified signed ONan traces.'))
    else:
        report=ONanCoefficients().verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)