"""wholelim.py -- does the whole-norm tail's loss stay bounded in r?

QUESTION. vartail.py's tail at order s = r - 3 (d = 2) pays, over the
exact certificate, Lam = 2V mean|R|/||P||_1 times sqrt(E g^2)/E|g|
(P = Q_s c*, R = P/(1 - x), V the sign changes of P). It read 2.93 to
3.84 at ranks 5..10. Bounded in r, a tail from O(J_c), and with it a
uniform bound J_ch <= C r^(3/2), is in reach along this route; growing,
it is not.

BY HAND (the Gaussian limit at fixed r, J -> oo). sigma = sqrt(J)/2;
P_k -> 2^J sigma^(-1-s) f(y)/sqrt(2 pi), f = (He_s + b He_(s+2)) phi
least in L2 (tailsplit.py's limit, b1 = 0 by parity). Since
int He_n phi e^(i w y) dy/sqrt(2 pi) = (i w)^n e^(-w^2/2), at w =
sigma theta |P(e^(i theta))| -> 2^J sigma^(-s) |w|^s e^(-w^2/2)
|1 - b w^2|, and |1 - x| -> |theta|, so
  mean|R| -> 2^J sigma^(-s) 2^(s/2) Gamma(s/2) E|1 - bW|/(2 pi),
  W ~ chi^2_s (U = sin^2(theta/2) -> W/J), ||P||_1 -> 2^J sigma^(-s)
  int|f|/sqrt(2 pi), and the Cauchy-Schwarz factor tends to
  sqrt(E(1 - bW)^2)/E|1 - bW|, E(1 - bW)^2 = 1 - 2bs + b^2 s(s + 2):
  Lam -> 2V 2^(s/2) Gamma(s/2) sqrt(E(1 - bW)^2)/(sqrt(2 pi) int|f|).
Split Lam = A_var A_max A_cs: A_var = 2V max|F|/int|f|, F = (He_(s-1)
+ b He_(s+1)) phi, R's limit shape ((He_(n-1) phi)' = -He_n phi);
A_max = 2^(s/2) Gamma(s/2) E|1 - bW|/(sqrt(2 pi) max|F|) (max|R_k| <=
mean|R|); A_cs the Cauchy-Schwarz factor. By hand int|He_n phi| ~
sqrt(n!) n^(-1/4) and max|He_(s-1) phi|, 2^(s/2) Gamma(s/2) ~ sqrt(s!)
s^(-3/4): A_max, A_cs = O(1), A_var ~ sqrt(r): V counts every sign
change of f, only about sqrt(s) of them in phi's bulk.

PREDICTIONS (frozen before the engine). PL1 A_max and A_cs move under
20% over r = 20..40. PL2 A_var's and Lam's local slopes in log r over
r = 20..40 in [0.4, 0.6]. PL3 V_true = r - 1 at even r.
KILL (printed): Lam with V = V_true (the real roots of He_s + b
He_(s+2)) has local slope log(Lam(40)/Lam(20))/log 2 >= 0.3.
CONTROL (positive, read first): chi_abs against quadrature at s = 5,
37, 57 (added after review: E|1 - bW| cancels in Lam, so the band
below never tests it; it feeds A_max and A_cs). Lam with V = r - 1
within 15% of
vartail.exact_read's 2(r - 1) mean/||P||_1 sqrt(E g^2)/E|g| at
dualtail.py's J_t, ranks 8..10. PLANT: b = 0 in the circle mean and in
E(1 - bW)^2, read against the same band.

FINDINGS (copied from the printed output).
  (1) Control: chi_abs equals its quadrature to 1e-15 at s = 5, 37, 57;
      the limit over the exact loss is 0.999, 1.000, 1.000 at
      ranks 8..10 (exact 3.5524, 3.4780, 3.8357); the plant reads 1.856,
      1.986, 2.109, caught at every rank.
  (2) PL1, PL2, PL3 hold; the KILL fires. Lam(V_true) = 5.1468, 6.2180,
      7.1357, 7.9501, 8.6894 at r = 20, 30, 40, 50, 60; its local slope
      0.471 over r = 20..40 and 0.486 over 40..60. A_var carries it
      (slope 0.503; 3.0742 at r = 20 to 5.3248 at 60); A_max stays
      within 1.305 .. 1.361 at r >= 15 and A_cs rises to 1.2462 (slopes
      -0.045, 0.014).
  (3) V_true is r - 1 at even r and r - 3 at odd r at every rank read
      (vartail's var(P) 7, 6, 9 at ranks 8..10). b reads 2/(2s + 3) to
      the printed digits at every rank (0.28571 at s = 2, 0.01709 at
      s = 57), b1 below 1e-16. By Parseval int He_m phi He_n phi dy =
      i^(m - n) Gamma((m + n + 1)/2) at m + n even, 0 at odd, so b1 = 0
      and b = Gamma(s + 3/2)/Gamma(s + 5/2) = 2/(2s + 3) exactly; the
      grid's b is within 4.4e-16 of it, relative.
  (4) So the whole-norm tail's loss grows as sqrt(r), all of it in the
      variation step: P changes sign about r times, but only about
      sqrt(r) of its runs sit in the Gaussian bulk, and 2V max|R| pays
      for every one at the bulk's height.

RUN RECORD. 4.6 s, numpy grid of step 2e-4 over |y| <= 40 (tailsplit's),
E|1 - bW| through the regularised gamma at 30 digits; the control reads
vartail.exact_read in rationals.
"""
import math
import os
import sys
import time

os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
import numpy as np  # noqa: E402
import mpmath as mp  # noqa: E402
from numpy.polynomial import hermite_e as He  # noqa: E402

mp.mp.dps = 30
Y = np.linspace(-40, 40, 400001)
DY = Y[1] - Y[0]
PHI = np.exp(-Y ** 2 / 2)


def h(n):
    c = np.zeros(n + 1)
    c[n] = 1
    return He.hermeval(Y, c) * PHI


def chi_abs(s, b):
    """E|1 - bW|, W ~ chi^2_s, exactly through the regularised gamma."""
    s, b = mp.mpf(s), mp.mpf(b)
    m0 = 1 - b * s                       # E(1 - bW)
    if b == 0:
        return abs(m0)
    w0 = 1 / b
    if w0 <= 0:                          # 1 - bW keeps one sign
        return abs(m0)
    # E[(1 - bW) 1(W > w0)]: P(W > w0) and E[W 1(W > w0)] = s P(chi^2_(s+2) > w0)
    tail = mp.gammainc(s / 2, w0 / 2, mp.inf, regularized=True)
    tailw = s * mp.gammainc(s / 2 + 1, w0 / 2, mp.inf, regularized=True)
    neg = tail - b * tailw              # E[(1 - bW) 1(W > w0)] < 0 if b > 0
    return float(m0 - 2 * neg)


def limit(r, plant=False):
    s = r - 3
    fs = [h(s + m) for m in range(3)]
    G = np.array([[np.sum(fs[i] * fs[j]) * DY for j in range(3)]
                  for i in range(3)])
    beta = np.linalg.solve(G[1:, 1:], -G[1:, 0])
    b1, b = beta
    f = fs[0] + b1 * fs[1] + b * fs[2]
    F = h(s - 1) + b1 * h(s) + b * h(s + 1)
    l1f = np.sum(np.abs(f)) * DY
    maxF = np.max(np.abs(F))
    # V_true: real roots of He_s + b1 He_(s+1) + b He_(s+2)
    c = np.zeros(s + 3)
    c[s], c[s + 1], c[s + 2] = 1, b1, b
    rts = He.hermeroots(c)
    vtrue = int(np.sum(np.abs(rts.imag) < 1e-7 * (1 + np.abs(rts.real))))
    bb = 0.0 if plant else b
    gam = float(mp.mpf(2) ** (mp.mpf(s) / 2) * mp.gamma(mp.mpf(s) / 2))
    Eabs = chi_abs(s, bb)
    E2 = 1 - 2 * bb * s + bb * bb * s * (s + 2)
    Acs = math.sqrt(E2) / Eabs
    Amax = gam * Eabs / (math.sqrt(2 * math.pi) * maxF)
    unit = gam * math.sqrt(E2) / (math.sqrt(2 * math.pi) * l1f)
    return dict(b=b, b1=b1, vtrue=vtrue, Acs=Acs, Amax=Amax,
                Avar1=2 * (r - 1) * maxF / l1f, Avar=2 * vtrue * maxF / l1f,
                lam1=2 * (r - 1) * unit, lam=2 * vtrue * unit)


def chi_quad(s, b):
    """E|1 - bW| by quadrature of the chi^2_s density, split at 1/b."""
    s, b = mp.mpf(s), mp.mpf(b)
    lc = (s / 2) * mp.log(2) + mp.loggamma(s / 2)
    def dens(w):
        return abs(1 - b * w) * mp.exp((s / 2 - 1) * mp.log(w) - w / 2 - lc)
    return float(mp.quad(dens, [0, 1 / b, mp.inf]))


def control():
    from vartail import exact_read, JT
    ok, caught = True, False
    for s, b in ((5, 2 / 13), (37, 2 / 77), (57, 2 / 117)):
        cq, cf = chi_quad(s, b), chi_abs(s, b)
        ok &= abs(cq - cf) <= 1e-12
        print("CONTROL E|1 - bW| at s = %d: closed %.15f, quadrature %.15f"
              % (s, cf, cq))
    for r in (8, 9, 10):
        J = JT[r]
        e = exact_read(r, r - 3, J)
        l1P = mp.mpf(e["l1P"].numerator) / e["l1P"].denominator
        ex = float(2 * (r - 1) * e["mean"] / l1P * mp.sqrt(e["Eg2"]) / e["Eg"])
        lm = limit(r)["lam1"]
        lp = limit(r, plant=True)["lam1"]
        within = abs(lm / ex - 1) <= 0.15
        pin = abs(lp / ex - 1) <= 0.15
        ok &= within
        caught |= not pin
        print("CONTROL r = %2d J = %4d: exact %.4f, limit %.4f (%.3f), "
              "PLANT %.4f (%.3f) %s; var(P) %d" % (
                  r, J, ex, lm, lm / ex, lp, lp / ex,
                  "in band" if pin else "out", e["V"]), flush=True)
    print("control: %s; plant %s" % ("ok" if ok else "OFF",
                                      "caught" if caught else "MISSED"))
    return ok, caught


def main():
    t0 = time.time()
    ok, caught = control()
    rows = {}
    for r in list(range(5, 11)) + [15, 20, 25, 30, 35, 40, 50, 60]:
        L = limit(r)
        rows[r] = L
        print("r = %2d  b %+.5f  b1 %.1e  V_true %2d  A_var %.4f  A_max %.4f"
              "  A_cs %.4f  Lam(V_true) %.4f  Lam(r - 1) %.4f" % (
                  r, L["b"], L["b1"], L["vtrue"], L["Avar"], L["Amax"],
                  L["Acs"], L["lam"], L["lam1"]), flush=True)

    def slope(key, a=20, z=40):
        return math.log(rows[z][key] / rows[a][key]) / math.log(z / a)
    print("local slopes over r = 20..40: Lam(V_true) %.3f, Lam(r - 1) %.3f,"
          " A_var %.3f, A_max %.3f, A_cs %.3f; over 40..60: Lam %.3f" % (
              slope("lam"), slope("lam1"), slope("Avar"), slope("Amax"),
              slope("Acs"), slope("lam", 40, 60)))
    rb = max(abs(L["b"] * (2 * r - 3) / 2 - 1) for r, L in rows.items())
    print("largest |b (2s + 3)/2 - 1| over the ranks read: %.1e" % rb)
    pl1 = all(max(rows[r][k] for r in (20, 25, 30, 35, 40))
              / min(rows[r][k] for r in (20, 25, 30, 35, 40)) < 1.2
              for k in ("Amax", "Acs"))
    pl2 = all(0.4 <= slope(k) <= 0.6 for k in ("Avar", "lam"))
    pl3 = all(rows[r]["vtrue"] == r - 1 for r in rows if r % 2 == 0)
    print("PL1 %s; PL2 %s; PL3 %s" % tuple("holds" if v else "fails"
                                           for v in (pl1, pl2, pl3)))
    kill = slope("lam") >= 0.3
    print("KILL %s (%.1f s)" % ("fires" if kill else "misses",
                                time.time() - t0))
    return not ok or not caught


if __name__ == "__main__":
    sys.exit(1 if main() else 0)