"""Numerical checks of constants of Sections 10 and 12.

- the phase average of the top coefficient of the frozen source (proof of l:res-permanent);
- Upsilon(q) = 2 sqrt(q) K_1(2 sqrt(q)) and its square integral;
- the spectral weight w(k) of e:coarse-weight, the Kontorovich--Lebedev projection of Upsilon,
  and the tail pi^{7/2}/(32 r^{3/2}) e^{-r/4} of the mean profile (c:mean-shape);
- the second moment of the mean profile, compared with 8 zeta(3) (r:coarse-mean-physics;
  proved in l:coarse-second-moment; this script supplies a numerical check only).
"""
import sympy as sp, mpmath as mp
# frozen source top coefficient and its theta average
th = sp.symbols('theta', real=True)
nR, nA = sp.symbols('nR nA', commutative=True)  # n_R^2 = n_A^2 = 0, commuting even nilpotents
t = 1 + sp.cos(th)
s0 = t*(1 + (nR + nA)/2) + sp.cos(th)/4*nR*nA
for k in range(2, 9):
    e = sp.expand(s0**k)
    e = sp.Poly(e, nR, nA)
    top = sum(c for (a, b), c in e.terms() if a == 1 and b == 1)
    # discard terms with nR^2 or nA^2 (they vanish)
    claim = sp.Rational(k, 4)*t**(k-1)*(k*t - 1)
    avg = sp.integrate(sp.expand(top), (th, 0, 2*sp.pi))/(2*sp.pi)
    print(k, sp.simplify(top - claim) == 0, sp.simplify(avg - k*sp.factorial2(2*k-3)/(2*sp.factorial(k-2))) == 0)
mp.mp.dps = 30
Ups = lambda q: 2*mp.sqrt(q)*mp.besselk(1, 2*mp.sqrt(q))
for q in [0.1, 1, 7]:
    print("Upsilon", q, Ups(q) - mp.quad(lambda x: mp.e**(-q/x - x), [0, 1, mp.inf]))
print("int Upsilon^2 =", mp.quad(lambda q: Ups(q)**2, [0, 1, 10, mp.inf]))
w = lambda k: 2*mp.pi**2*k*(k**2 + mp.mpf(1)/4)**2*mp.tanh(mp.pi*k)/mp.cosh(mp.pi*k)**2
print("int w =", mp.quad(w, [0, 1, 5, mp.inf]), " int w/(k^2+1/4) =", mp.quad(lambda k: w(k)/(k**2 + 0.25), [0, 1, 5, mp.inf]))
# projection <psi_k, U Upsilon> with psi_k(y) = (2/pi) sqrt(k sinh(2 pi k)) K_{2ik}(2 e^{y/2}), U f(y) = e^{y/2} f(e^y)
for k in [0.3, 1.1]:
    k = mp.mpf(k)
    psi = lambda y: 2/mp.pi*mp.sqrt(k*mp.sinh(2*mp.pi*k))*mp.besselk(2j*k, 2*mp.e**(y/2)).real
    Uu = lambda y: mp.e**(y/2)*Ups(mp.e**y)
    num = mp.quad(lambda y: psi(y)*Uu(y), [-40, -10, 0, 3, 6])
    claim = mp.pi*(k**2 + 0.25)*mp.sqrt(k*mp.sinh(2*mp.pi*k))/mp.cosh(mp.pi*k)**2
    print("projection k=", k, num, claim)
# eigenfunction normalization check: int psi_k psi_k' dy ~ delta(k-k'); check via the KL Parseval on a test function
def F_of(r):
    # For large r the integrand e^{-r k^2} w(k) is concentrated at k ~ r^{-1/2}; substitute
    # k = kappa/sqrt(r) so that the quadrature resolves the peak.
    r = mp.mpf(r)
    if r == 0:
        return mp.quad(w, [0, 0.5, 2, 6, mp.inf])
    I = mp.quad(lambda kap: mp.e**(-kap**2)*w(kap/mp.sqrt(r))/mp.sqrt(r), [0, 1, 3, 6, 12, mp.inf])
    return mp.e**(-r/4)*I
for r in [0, 50, 200, 800, 3200, 12800]:
    val = F_of(r)
    if r == 0:
        print("F(0) =", val)
    else:
        ratio = val/(mp.pi**3.5/(32*mp.mpf(r)**1.5)*mp.e**(-mp.mpf(r)/4))
        print("tail ratio r=", r, ratio, " r*(ratio - 1) =", mp.mpf(r)*(ratio - 1))

# second moment of the mean profile (r:coarse-mean-physics)
m2 = 8*mp.pi**2*mp.quad(lambda k: k*mp.tanh(mp.pi*k)/(mp.cosh(mp.pi*k)**2*(k**2 + mp.mpf(1)/4)), [0, 1, 5, mp.inf])
print("second moment =", m2, " 8 zeta(3) =", 8*mp.zeta(3))