"""Exact newform roster and dimension checks for numberdb.org/T694.
Run: sage -python prepare.py > prepare.log; final ROSTER line is roster.json.
Needs companion-roster.json, downloaded from T683 and restricted to k=4,6.
"""
import numberdb.sage as numberdb
from sage.libs.pari import pari
from sage.rings.rational_field import QQ
from sage.rings.integer_ring import ZZ
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from sage.rings.qqbar import AA, _init_qqbar
from pathlib import Path
from functools import lru_cache
import json,math
from sage.arith.misc import euler_phi
_init_qqbar()
P=PolynomialRing(QQ,'x')
COMP=json.loads((Path(__file__).parent/'companion-roster.json').read_text())
def poly(c):return P(str(pari.lift(c)).replace('y','x'))
def coeffs(c):return [str(x) for x in c.list()]
def rational(s):
 from fractions import Fraction
 q=Fraction(s);return QQ(q.numerator)/q.denominator
@lru_cache(None)
def dim(N,k):
 n=ZZ(N);ps=n.prime_divisors();mu=QQ(N)
 for p in ps:mu*=1+QQ(1)/p
 e2=0 if N%4==0 else math.prod(1+int(ZZ(-4).kronecker(p)) for p in ps)
 e3=0 if N%2==0 or N%9==0 else math.prod(1+int(ZZ(-3).kronecker(p)) for p in ps)
 cusps=sum(euler_phi(math.gcd(d,N//d)) for d in n.divisors())
 g=1+mu/12-QQ(e2)/4-QQ(e3)/3-QQ(cusps)/2
 return int((k-1)*(g-1)+(k//2-1)*cusps+e2*(k//4)+e3*(k//3))
@lru_cache(None)
def newdim(N,k):
 return dim(N,k)-sum(len(ZZ(N//d).divisors())*newdim(int(d),k) for d in ZZ(N).divisors() if d<N)
def companion_forms(N,k,B):
 found=[]
 for row in COMP:
  if (row['level'],row['weight'])!=(N,k):continue
  F=P(row['field_poly']);roots=F.roots(AA,multiplicities=False);cs=[P([QQ(c) for c in v]) for v in row['an_power_basis']]
  assert len(cs)>B
  for e in row['embeddings']:
   low,hi=map(rational,e['root_rectangle'][0]);rs=[z for z in roots if low<z<hi];assert len(rs)==1
   found.append((e['label'],[c(rs[0]) for c in cs[:B+1]]))
 return found
spaces=[('A',1,k) for k in range(4,27,2)]+[('B',n,4) for n in range(2,31)]+[('C',n,6) for n in range(2,16)]+[('D',n,12) for n in range(2,11)]
roster=[];dims=[]
for group,N,k in spaces:
 mf=pari.mfinit([N,k],0);dd=int(pari.mfdim(mf));assert dd==newdim(N,k)
 assert int(pari.mfdim([N,k],1))==dim(N,k)
 B=max(2,int(pari.mfsturm([N,k])));embedded=[];companions=companion_forms(N,k,B) if group in ['B','C'] else []
 A=pari.mfatkininit(mf,N,precision=128) if dd else None
 for orbit,f in enumerate(pari.mfeigenbasis(mf),1):
  F=poly(pari.mfparams(f)[3]);cs=[poly(c) for c in pari.mfcoefs(f,B)];roots=F.roots(AA,multiplicities=False)
  assert len(roots)==F.degree();assert cs[0]==0 and cs[1]==1
  assert A[2]==1
  ac=list(pari.mfcoefs(pari.mfatkin(A,f),B));fc=list(pari.mfcoefs(f,B));sgn=next(s for s in [-1,1] if all(a==s*b for a,b in zip(ac,fc)))
  for z in roots:
   q=(z*2**30).floor();lo=QQ(q-1)/2**30;hi=QQ(q+2)/2**30
   assert sum(lo<r<hi for r in roots)==1
   vals=[c(z) for c in cs];label=None
   if companions:
    matches=[name for name,cc in companions if cc==vals];assert len(matches)==1;label=matches[0]
   row={'group':group,'N':N,'k':k,'sturm_bound':B,'field_poly':str(F),'field_degree':int(F.degree()),'root_interval':[str(lo),str(hi)],'coefficients':[coeffs(c) for c in cs],'fricke':sgn,'epsilon':(-1)**(k//2)*sgn,'companion_label':label}
   embedded.append((vals,row))
 embedded.sort(key=lambda r:r[0][2:])
 assert len(embedded)==dd and all(embedded[i-1][0]!=embedded[i][0] for i in range(1,dd))
 for i,(v,row) in enumerate(embedded,1):row['i']=i;roster.append(row)
 dims.append({'group':group,'N':N,'k':k,'cusp':dim(N,k),'new':dd})
 print('SPACE',group,N,k,dd,flush=True)
assert len(roster)==121
assert sum(r['k']-2 for r in roster)==626
print('DIMENSIONS '+json.dumps(dims))
print('ROSTER '+json.dumps(roster))