"""Independent eta-product fundamental-domain integrals for numberdb.org/T684.

Run: sage -python quadrature.py. No modular-form or Petersson library calls.
Three controls cover weights 3, 4 and 6: eta(z)^k eta(p*z)^k for
(p,k)=(7,3),(5,4),(3,6). Coefficients are exact product expansions.
For prime p, cosets are 1 and S*T^j, 0<=j<p. Eta inversion gives
|f|S(z)|^2=p^(-k)|f(z/p)|^2. Summing translations leaves only n=m mod p.
The y integral over y>=sqrt(1-x^2) is an integer incomplete-gamma finite sum;
the remaining x integral uses independently implemented Gauss-Legendre nodes.
Two coefficient cutoffs and quadrature orders test truncation. No certified
error enclosure is claimed. The polynomial moments are known controls.
"""
import math,json,time
import numberdb.sage as numberdb
from sage.rings.real_mpfr import RealField


def gauss(n,R):
    result=[]
    for j in range(1,n//2+1):
        z=(R.pi()*(R(j)-R(1)/4)/(R(n)+R(1)/2)).cos()
        for it in range(30):
            a,b=R(1),z
            for k in range(2,n+1):a,b=b,((2*k-1)*z*b-(k-1)*a)/k
            deriv=n*(z*b-a)/(z*z-1);new=z-b/deriv
            if abs(new-z)<R(2)**(-R.prec()+10):z=new;break
            z=new
        else:raise AssertionError('Legendre Newton did not converge')
        weight=2/((1-z*z)*deriv*deriv)
        # Transform [-1,1] to [0,1/2].
        result.extend([((1-z)/4,weight/4),((1+z)/4,weight/4)])
    for power in range(8):
        exact=(R(1)/2)**(power+1)/(power+1)
        assert abs(sum(w*x**power for x,w in result)-exact)<R(2)**(-R.prec()+20)
    return result


def coefficients(p,k,T):
    a=[0]*T;a[0]=1
    for scale in [1,p]:
        for n in range(1,(T-1)//scale+1):
            step=scale*n
            for _ in range(k):
                for j in range(T-1,step-1,-1):a[j]-=a[j-step]
    assert k*(1+p)==24
    return [0]+a[:-1]


def integrate(p,k,T,order,working, supplied_coefficients=None):
    R=RealField(numberdb.bits(working,losing=32));pi=R.pi()
    a=coefficients(p,k,T) if supplied_coefficients is None else supplied_coefficients
    assert a[1]==1
    if supplied_coefficients is None: assert a[2]==-k
    # Truncate by n+m<T, not by each index separately. Both tails are
    # varied independently of the Gauss quadrature order in the driver.
    terms=[]
    for scale in [1,p]:
        pairs=[]
        for t in range(2,T):
            row=[]
            for n in range(1,t//2+1):
                m=t-n;d=m-n
                if d%scale==0 and a[n] and a[m]:
                    row.append((d//scale,(1 if n==m else 2)*(a[n]*a[m].conjugate()).real() if supplied_coefficients is not None else (1 if n==m else 2)*a[n]*a[m]))
            pairs.append((t,row))
        terms.append((scale,pairs))
    total=R(0)
    for x,w in gauss(order,R):
        Y=(1-x*x).sqrt()
        c1=(2*pi*x).cos();cosines=[R(1),c1]
        for d in range(2,T):cosines.append(2*c1*cosines[-1]-cosines[-2])
        point=R(0)
        for scale,pairs in terms:
            b=2*pi/R(scale);ratio=(-b*Y).exp();ex=ratio
            subtotal=R(0)
            for t,row in pairs:
                ex*=ratio
                if not row:continue
                u=b*t*Y;term=R(1);poly=term
                for j in range(1,k-1):term*=u/j;poly+=term
                yg=math.factorial(k-2)*ex*poly/(b*t)**(k-1)
                trig=sum((c*cosines[d] for d,c in row),R(0))
                subtotal+=trig*yg
            point+=subtotal if scale==1 else subtotal/R(p)**(k-1)
        total+=2*w*point
    assert total>0
    return total


def main():
    for p,k in [(3,6),(5,4),(7,3)]:
        # About 5.44*T/p in the smallest exponential at the boundary.
        # These are measured truncation settings, not a formal tail bound.
        vals=[]
        for cutoff,order,working in [(70*p,140,150),(90*p,180,190)]:
            t=time.time();v=integrate(p,k,cutoff,order,working);vals.append(v)
            print('QUADRATURE',json.dumps({'p':p,'k':k,'cutoff':cutoff,'order':order,'working':working,'integral':str(v),'seconds':time.time()-t}),flush=True)
        rel=abs(vals[0]-vals[1])/vals[1]
        assert rel<RealField(800)('1e-120'),(p,k,str(rel))
        print('TRUNCATION',p,k,str(rel),flush=True)
if __name__=='__main__':main()
