"""Critical values of Rankin-Selberg L-functions of level-one cusp forms -- numberdb.org/T691.

Run with SageMath (install numberdb once with sage -pip install numberdb):
    sage -python generate.py                 # verify every stored entry
    sage -python generate.py --compute       # print reproducible numerical evidence
    sage -python generate.py --publish       # fill this draft, using NUMBERDB_API_KEY

Reading a private draft also requires NUMBERDB_API_KEY with access to it.
The --publish switch submits entries to the existing draft; it does not make
the table public or offer it for review.

The default selection is 66 critical values for f_(k,i) tensor Delta with
12<k<=26. See argument-selection.md and exact-roster.json. Values are rounded
to 100 significant digits only after fresh evaluations at 160 and 220 working
decimal digits agree to 120 relative digits. This is heuristic agreement,
not a certified error bound. PARI's required internal coefficient precision
can exceed its output precision: the quadratic embedding is therefore computed
inside a PARI closure at each requested precision, never supplied as a fixed
precision decimal. Internally h=(k+10)/2 shifts L(s) to L_an(s-h), with
coefficients c_n/n^h and gamma parameters h,h+1,h-11,h-10. This exact change
of coordinates improved the measured 160/220-digit agreement for the first
pair from about 105 to about 164 digits. Arithmetic values are still stored.
check_values.py gives independent numerical checks.
"""
import os
import sys
import json
import time
from functools import lru_cache
from math import isqrt
import numberdb.sage as numberdb
import sage.symbolic.expression
from sage.rings.integer_ring import ZZ
from sage.rings.real_mpfr import RealField
from sage.rings.rational_field import QQ
from sage.rings.power_series_ring import PowerSeriesRing
from sage.libs.pari import pari
from sage.version import version as sage_version

WEIGHTS = (16, 18, 20, 22, 24, 26)
WORKING_DIGITS = (160, 220)
OUTPUT_DIGITS = 100
AGREEMENT_DIGITS = 120
GAMMA = (0, 1, -11, -10)


def add(a, b):
    return (a[0] + b[0], a[1] + b[1])


def scale(a, n):
    return (a[0] * n, a[1] * n)


def mul(a, b):
    # t^2 = 1080*t + 20468736, for t=a_2 of the weight-24 form.
    return (a[0]*b[0] + 20468736*a[1]*b[1],
            a[0]*b[1] + a[1]*b[0] + 1080*a[1]*b[1])


def eisenstein_coefficients(n, weight):
    factor = {4:240, 6:-504}[weight]
    c = [ZZ(0)] * (n+1)
    c[0] = ZZ(1)
    for d in range(1, n+1):
        v = ZZ(factor) * ZZ(d)**(weight-1)
        for m in range(d, n+1, d):
            c[m] += v
    return c


@lru_cache(None)
def exact_forms(n):
    R = PowerSeriesRing(ZZ, 'q', default_prec=n+1)
    e4 = R(eisenstein_coefficients(n,4)).add_bigoh(n+1)
    e6 = R(eisenstein_coefficients(n,6)).add_bigoh(n+1)
    num = e4**3-e6**2
    assert all(num[m] % 1728 == 0 for m in range(n+1))
    delta = R([num[m]//1728 for m in range(n+1)]).add_bigoh(n+1)
    assert [delta[m] for m in range(6)] == [0,1,-24,252,-1472,4830]
    polys = {12:delta,16:delta*e4,18:delta*e6,20:delta*e4**2,
             22:delta*e4*e6,24:delta*e4**3-696*delta**2,
             26:delta*e4**2*e6}
    delta2 = delta**2
    return {k:[(p[m],delta2[m] if k==24 else ZZ(0)) for m in range(n+1)]
            for k,p in polys.items()}


def prime_sieve(n):
    spf=list(range(n+1))
    for p in range(2,isqrt(n)+1):
        if spf[p]==p:
            for m in range(p*p,n+1,p):
                if spf[m]==m: spf[m]=p
    return spf


@lru_cache(None)
def tensor_coefficients(k, n):
    forms=exact_forms(n); f=forms[k]; g=forms[12]; spf=prime_sieve(n)
    local={}
    for p in range(2,n+1):
        if spf[p]!=p: continue
        a=f[p]; b=g[p]; A=ZZ(p)**(k-1); B=ZZ(p)**11
        ab=mul(a,b)
        middle=add(add(scale(mul(a,a),B),scale(mul(b,b),A)),(-2*A*B,ZZ(0)))
        den=[(ZZ(1),ZZ(0)),scale(ab,-1),middle,scale(ab,-A*B),((A*B)**2,ZZ(0))]
        h=[(ZZ(1),ZZ(0))]; pp=p
        while pp<=n:
            m=len(h); v=(ZZ(0),ZZ(0))
            for r in range(1,min(4,m)+1): v=add(v,mul(den[r],h[m-r]))
            h.append(scale(v,-1)); pp*=p
        local[p]=h
    c=[(ZZ(0),ZZ(0))]*(n+1); c[1]=(ZZ(1),ZZ(0))
    for m in range(2,n+1):
        p=spf[m]; r=m; e=0
        while r%p==0: r//=p; e+=1
        c[m]=mul(c[r],local[p][e])
    return c


def coefficient_closure(k, i, n):
    c=tensor_coefficients(k,n)
    name='rankin%d_%d_%d' % (k,i,n)
    pari('%sU=%s' % (name, '['+','.join(str(a[0]) for a in c[1:])+']'))
    if k!=24:
        return pari('n->vector(n,j,%sU[j]/j^%d)' % (name,(k+10)//2))
    pari('%sV=%s' % (name, '['+','.join(str(a[1]) for a in c[1:])+']'))
    sign='-' if i==1 else '+'
    return pari('n->vector(n,j,(%sU[j]+(540%s12*sqrt(144169))*%sV[j])/j^%d)' % (name,sign,name,(k+10)//2))


def objects():
    for k in WEIGHTS:
        for i in range(1,3 if k==24 else 2): yield k,i


@lru_cache(None)
def compute_object(k, i, working):
    bits=numberdb.bits(working)
    pari.set_real_precision_bits(bits)
    shift=(k+10)//2
    gamma=[x+shift for x in GAMMA]
    domain=[QQ(1)/QQ(2),QQ(k-13)/QQ(2),0]
    template=pari.lfuncreate([pari('n->vector(n,j,1)'),0,gamma,1,1,1])
    cost=pari.lfuncost(template,domain,precision=bits)
    n=2*int(cost[0])+100
    data=pari.lfuncreate([coefficient_closure(k,i,n),0,gamma,1,1,1])
    init=pari.lfuninit(data,domain,precision=bits)
    R=RealField(bits)
    values={s:R(pari.lfun(init,s-shift,precision=bits)) for s in range(12,k)}
    assert all((not v.is_NaN() and not v.is_infinity()) and v!=0 for v in values.values())
    feq=int(pari.lfuncheckfeq(data,precision=bits))
    assert feq <= -bits+64, (k,i,working,feq)
    return values, {'k':k,'i':i,'working_digits':working,'bits':bits,
                    'internal_argument_shift':shift,'gamma_parameters':gamma,
                    'coefficients':n,'internal_coefficient_bits':int(cost[1]),
                    'lfuncheckfeq_log2_error':feq}


def decimal_value(v, digits):
    return v.str(base=10, digits=digits, no_sci=False)


class RankinSelbergCritical(numberdb.Generator):
    table='T691'
    parameters=('k','i','ell','j','s')
    type='R'
    digits=OUTPUT_DIGITS
    rigour='heuristic (agreement-checked)'
    files=('generate.py','check_values.py','argument-selection.md','exact-roster.json',
           'numerical-evidence.json','skill-evidence.json','sources.json',
           'preflight-evidence.txt','readback_checks.py')

    def enumerate(self):
        for k,i in objects():
            for s in range(12,k):
                yield dict(k=str(k),i=str(i),ell='12',j='1',s=str(s))

    def value(self, params, digits):
        k,i,ell,j,s=(int(params[t]) for t in self.parameters)
        if (k,i) not in set(objects()) or (ell,j)!=(12,1) or not 12<=s<k:
            raise ValueError('outside the explicitly selected anchored critical set')
        lo=compute_object(k,i,WORKING_DIGITS[0])[0][s]
        hi=compute_object(k,i,WORKING_DIGITS[1])[0][s]
        R=hi.parent(); rel=abs(R(lo)-hi)/abs(hi)
        if rel>=R(10)**(-AGREEMENT_DIGITS):
            raise ArithmeticError('working-precision agreement failed at '+str(params))
        return decimal_value(hi,digits)


def compute_evidence():
    start=time.time(); records=[]; calculations=[]
    for k,i in objects():
        low,meta1=compute_object(k,i,WORKING_DIGITS[0])
        high,meta2=compute_object(k,i,WORKING_DIGITS[1]); calculations.extend([meta1,meta2])
        worst=high[12].parent()(0)
        for s in range(12,k):
            R=high[s].parent(); rel=abs(R(low[s])-high[s])/abs(high[s]); worst=max(worst,rel)
            assert rel<R(10)**(-AGREEMENT_DIGITS)
            records.append({'params':dict(k=str(k),i=str(i),ell='12',j='1',s=str(s)),
                            'number':decimal_value(high[s],OUTPUT_DIGITS),
                            'low':str(low[s]),'high':str(high[s]),'relative_difference':str(rel)})
        print('COMPUTED',k,i,'relative difference',worst,'elapsed',time.time()-start,flush=True)
    print('NUMERICAL_JSON '+json.dumps({'table':'T691','sage':sage_version,'pari':str(pari.version()),
        'numberdb':numberdb.__version__,'calculations':calculations,'records':records}),flush=True)


if __name__=='__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
        key=sys.stdin.read().strip()
        if '=' in key and key.split('=',1)[0].isidentifier(): key=key.split('=',1)[1].strip().strip('\"\'')
        os.environ['NUMBERDB_API_KEY']=key
    if '--compute' in sys.argv or os.environ.get('RANKIN_COMPUTE')=='1':
        compute_evidence()
    elif '--publish' in sys.argv or os.environ.get('NUMBERDB_PUBLISH')=='1':
        print(RankinSelbergCritical().publish(message='Fill all 66 critical values with reproducible checks'))
    else:
        report=RankinSelbergCritical().verify(sample=None); print(report)
        sys.exit(0 if report.ok else 1)