# SCRIPT: BH-SHADOW-07
# Ramanujan's radial limits. f(q) = sum_n q^{n^2}/(-q;q)_n^2 is his third-order mock theta function and
# b(q) = (q;q)_inf/(-q;q)_inf^2 = (1-q)(1-q^3)(1-q^5)...(1-2q+2q^4-...) its modular partner. At a
# root of unity zeta of even order 2k both blow up along the radius q = zeta(1-eps), eps -> 0, but
# f - (-1)^k b stays bounded (Ramanujan, last letter, 1920; proved by Folsom, Ono and Rhoades 2013).
# The bounded value is a finite sum at the root of unity,
#     L(zeta) = -4 sum_{n>=0} (-zeta;zeta)_n^2 zeta^{n+1},
# which terminates because (-zeta;zeta)_n = 0 once n reaches the order of -zeta... (it vanishes as soon
# as a factor 1 + zeta^j is zero).
# Kill conditions, written before the run:
#  S1  |f| and |b| each grow without bound along the radius (zeta = -1, i, e^{i pi/3}).
#  S2  the difference converges; its Richardson extrapolation (fit a + b eps + c eps^2 through
#      eps = 3e-3, 1e-3, 3e-4) agrees with L(zeta) to 1e-3, and the raw error shrinks roughly in
#      proportion to eps. Kill otherwise.
# Precision: the working precision is set from the sizes of the products involved, so that the
# cancellation of two numbers as large as 1e600 is exact enough. b is summed by Gauss's theta series
# and Euler's pentagonal series, b = theta(q) (q;q)_inf/(q^2;q^2)_inf, which converge fast.
import mpmath as mp
ok=True
import json
DUMP={}
def report(name,cond,val=''):
    global ok
    print(('PASS ' if cond else 'FAIL ')+name, val); ok=ok and cond
import numpy as np
def logmag(zeta,eps):
    # double-precision estimate of the sizes involved, to set the working precision
    q=complex(zeta)*(1-float(eps)); j=np.arange(1,int(80/float(eps))+50)
    lq=np.sum(np.log(np.abs(1-q**j))); lmq=np.sum(np.log(np.abs(1+q**j)))
    lq2=np.sum(np.log(np.abs(1-q**(2*j))))
    return max(abs(lq-2*lmq),abs(lq),abs(lq2),abs(lmq))/np.log(10)
def euler(q,M):   # (q;q)_inf by the pentagonal number theorem
    return sum((-1)**n*q**(n*(3*n-1)//2) for n in range(-M,M+1))
def theta(q,M):   # (q;q)_inf/(-q;q)_inf by Gauss's identity
    return sum((-1)**n*q**(n*n) for n in range(-M,M+1))
def run(zeta,k,eps):
    digits=logmag(zeta,eps)
    mp.mp.dps=int(60+2*digits)
    q=zeta*(1-eps)
    M=int(mp.sqrt(2*mp.mp.dps*2.31/eps))+10
    b=theta(q,M)*euler(q,M)/euler(q*q,M)   # = (q;q)/(-q;q)^2, since 1/(-q;q) = (q;q)/(q^2;q^2)
    s=mp.mpf(0); den=mp.mpf(1); n=0; small=0
    while True:
        if n>0: den*=(1+q**n)**2
        t=q**(n*n)/den; s+=t
        if abs(t)<abs(s)*mp.mpf(10)**(-mp.mp.dps+5) and n>10: small+=1
        else: small=0
        if small>5 or n>400000: break
        n+=1
    return s,b,s-(-1)**k*b
def L(zeta):
    mp.mp.dps=30; s=0; p=mp.mpf(1); n=0
    while True:
        s+=p**2*zeta**(n+1); n+=1; p*=(1+zeta**n)
        if abs(p)<1e-25 or n>50: break
    return -4*s
for zeta,k,lab in [(mp.mpc(-1,0),1,'zeta = -1'),(mp.mpc(0,1),2,'zeta = i'),(mp.exp(1j*mp.pi/3),3,'zeta = e^(i pi/3)')]:
    lim=L(zeta); rows=[]
    print(f'{lab}:  L(zeta) = {mp.nstr(lim,10)}')
    for eps in ['1e-2','3e-3','1e-3','3e-4']:
        fv,bv,d=run(zeta,k,mp.mpf(eps)); rows.append((mp.mpf(eps),fv,bv,d))
        print(f'   1-|q| = {eps:>5}:  |f| = {mp.nstr(abs(fv),6):>12}   |b| = {mp.nstr(abs(bv),6):>12}   f-(-1)^k b = {mp.nstr(d,10)}')
    DUMP[lab]={'L':[float(lim.real),float(lim.imag)],'rows':[[float(e),float(mp.log10(abs(fv))),float(mp.log10(abs(bv))),float(d.real),float(d.imag)] for e,fv,bv,d in rows]}
    report(f'S1 {lab}: f and b blow up', abs(rows[-1][1])>1e10*abs(rows[0][1]) or abs(rows[-1][1])>1e20, f'|f| {mp.nstr(abs(rows[0][1]),3)} -> {mp.nstr(abs(rows[-1][1]),3)}')
    (x1,_,_,y1),(x2,_,_,y2),(x3,_,_,y3)=rows[1],rows[2],rows[3]
    rich=y1*x2*x3/((x1-x2)*(x1-x3))+y2*x1*x3/((x2-x1)*(x2-x3))+y3*x1*x2/((x3-x1)*(x3-x2))
    e1=abs(rows[2][3]-lim); e2=abs(rows[3][3]-lim)
    report(f'S2 {lab}: f-(-1)^k b -> L(zeta)', abs(rich-lim)<1e-3 and 2<e1/e2<4.5, f'Richardson {mp.nstr(rich,8)}; raw error {mp.nstr(e1,3)} -> {mp.nstr(e2,3)}')
json.dump(DUMP,open('bh_shadow_07.json','w'),indent=1)
print('ALL PASS' if ok else 'SOME FAILED')