"""Fricke traces of singular moduli -- numberdb.org/T687.

Run with SageMath:
    sage -pip install numberdb              # once
    sage -python generate.py                 # verify the stored draft/table
    sage -python generate.py --publish       # fill entries with your API key

All entries are rational Fourier coefficients of weight-2 weak Jacobi forms,
converted from fixed-residue Gamma_0(N) traces to full Atkin-Lehner traces.
Matsusaka (2017), Theorems 3.2, 3.4 and Lemma 4.2, identify these coefficients
with the CM orbit sums. No rational is reconstructed by rounding.

The ring theorem gives the basis E4^u E6^v A^a B^(N-a), 2u+3v=a+1,
where A=phi_{-2,1}, B=phi_{0,1}. Exact full-rank elimination of all singular
coefficients proves uniqueness within this complete weight/index basis.
See argument-selection.md and mathematical-method.md for selection and proof.
"""
import os
import sys
from functools import lru_cache
from math import gcd, isqrt
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.polynomial.laurent_polynomial_ring import LaurentPolynomialRing

LEVELS = (2, 3, 5, 6, 7, 10, 13)
ORDERS = (1, 2)
MAX_D = 200
L = LaurentPolynomialRing(QQ, 'z')
z = L.gen()


def add(f, g):
    return [a+b for a,b in zip(f,g)]


def scale(f, c):
    return [c*a for a in f]


def mul(f, g):
    length = len(f)
    out = [L.zero() for _ in f]
    for i,a in enumerate(f):
        if a:
            for j in range(length-i):
                if g[j]:
                    out[i+j] += a*g[j]
    return out


def one(length):
    return [L.one()] + [L.zero()]*(length-1)


def powers(f, bound):
    out = [one(len(f))]
    for _ in range(bound):
        out.append(mul(out[-1], f))
    return out


def sigma(n, k=1):
    return sum(ZZ(d)**k for d in range(1, n+1) if n % d == 0)


def prime_divisors(n):
    return [p for p in range(2,n+1) if n%p==0 and all(p%j for j in range(2,isqrt(p)+1))]


def ramification_factor(N,d):
    return ZZ(2)**sum(d%p==0 for p in prime_divisors(N))


@lru_cache(None)
def jacobi_generators(length):
    A = one(length)
    for k in range(1,length):
        factor = [L.zero() for _ in A]
        # (1-z*q^k)^2 (1-z^-1*q^k)^2
        numerator = [L(1),-2*(z+z**-1),z*z+z**-2+4,-2*(z+z**-1),L(1)]
        # (1-q^k)^(-4), exact binomial coefficients.
        for j in range((length-1)//k+1):
            factor[j*k] = sum(numerator[h]*ZZ((j-h+1)*(j-h+2)*(j-h+3)//6)
                              for h in range(min(4,j)+1))
        A=mul(A,factor)
    A=scale(A,z-2+z**-1)
    E2=[L(1)]+[L(-24*sigma(n)) for n in range(1,length)]
    E4=[L(1)]+[L(240*sigma(n,3)) for n in range(1,length)]
    E6=[L(1)]+[L(-504*sigma(n,5)) for n in range(1,length)]
    B=add([L({r: (6*r*r-24*n)*c for r,c in a.dict().items()})
           for n,a in enumerate(A)],scale(mul(E2,A),-5))
    assert A[0]==z-2+z**-1 and B[0]==z+10+z**-1
    assert A[1]==-2*z**2+8*z-12+8*z**-1-2*z**-2
    assert B[1]==10*z**2-64*z+108-64*z**-1+10*z**-2
    return A,B,E4,E6


def solve(rows, rhs, columns):
    """Exact rectangular elimination; require full column rank and consistency."""
    a=[[QQ(x) for x in row]+[QQ(y)] for row,y in zip(rows,rhs)]
    pivot=0
    for col in range(columns):
        k=next((i for i in range(pivot,len(a)) if a[i][col]),None)
        if k is None:
            raise ArithmeticError('Jacobi singular coefficients do not determine the form')
        a[pivot],a[k]=a[k],a[pivot]
        div=a[pivot][col]
        a[pivot]=[x/div for x in a[pivot]]
        for i in range(len(a)):
            if i != pivot and a[i][col]:
                c=a[i][col]
                a[i]=[x-c*y for x,y in zip(a[i],a[pivot])]
        pivot+=1
    assert all(not any(row[:-1]) and not row[-1] for row in a[pivot:])
    return [a[i][-1] for i in range(columns)]


def singular(N,m,d):
    if d==0:
        return 2*sum(e*sigma(m//e) for e in range(1,m+1) if N%e==0 and m%e==0)
    k=isqrt(-d)
    if k*k==-d and m%k==0:
        return -k*ramification_factor(N,k)
    return 0


@lru_cache(None)
def basis(N,bound=MAX_D):
    length=(bound+N*N)//(4*N)+2
    A,B,E4,E6=jacobi_generators(length)
    ap,bp=powers(A,N),powers(B,N)
    ep,fp=powers(E4,(N+1)//2),powers(E6,(N+1)//3)
    monomials=[]; labels=[]
    for a in range(N+1):
        ab=mul(ap[a],bp[N-a])
        for u in range((a+1)//2+1):
            for v in range((a+1)//3+1):
                if 2*u+3*v==a+1:
                    labels.append((a,u,v))
                    monomials.append(mul(ab,mul(ep[u],fp[v])))
    polar=[(n,r) for r in range(N+1) for n in range(r*r//(4*N)+1)]
    rows=[[f[n][r] for f in monomials] for n,r in polar]
    return monomials,labels,polar,rows


@lru_cache(None)
def trace_coefficients(N,m,bound=MAX_D):
    monomials,labels,polar,rows=basis(N,bound)
    solution=solve(rows,[singular(N,m,4*N*n-r*r) for n,r in polar],len(labels))
    f=[L.zero() for _ in monomials[0]]
    for c,g in zip(solution,monomials):
        f=add(f,scale(g,c))
    out={}
    # Check all computed elliptic translates, not just the extraction residue.
    for n,a in enumerate(f):
        for r in range(-N,N+1):
            d=4*N*n-r*r
            v=QQ(a[r])
            if d<=0:
                assert v==singular(N,m,d),(N,m,d,v)
            if 0<d<=bound:
                if d in out:
                    assert out[d]==v,(N,m,d,'residue-dependent coefficient')
                out[d]=v
    wanted={d for d in range(1,bound+1) if any((d+r*r)%(4*N)==0 for r in range(N+1))}
    assert set(out)==wanted
    return {d:v/ramification_factor(N,d) for d,v in out.items()}


class FrickeTraces(numberdb.Generator):
    table='T687'
    parameters=('N','m','d')
    type='Q'
    rigour='exact'
    digits=100
    files=('generate.py','argument-selection.md','mathematical-method.md','checks.py','check.py',
           'skill-manifest.json','numerical-evidence.json','orbit-evidence.json',
           'T652.json')

    def enumerate(self):
        for N in LEVELS:
            admissible={(-r*r)%(4*N) for r in range(2*N)}
            for m in ORDERS:
                for d in range(1,MAX_D+1):
                    if d%(4*N) in admissible:
                        yield {'N':str(N),'m':str(m),'d':str(d)}

    def value(self,params,digits):
        N,m,d=(int(params[k]) for k in self.parameters)
        return trace_coefficients(N,m)[d]


if __name__=='__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    g=FrickeTraces()
    if '--publish' in sys.argv or os.environ.get('NUMBERDB_PUBLISH')=='1':
        print(g.publish(message='Fill exact Fricke traces after Jacobi and CM-orbit checks.'))
    else:
        report=g.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)