"""partfive.py -- at fixed e = 3, which of the part-5 shape's three factors
breaks the concentration estimate of its X over F_e's mean?

QUESTION. Hold e = 3, J = 11 l - 1, l = r - 2, r = 21..51 both parities.
The part-5 shape P5 = k=0 2^(r-5) 5 binds the e = 3 arm. Its exact
centred X over F_e's circle mean splits into three factors at the test p:
  X/mean = B1(p) |C5(p)| / A,
  B1(p) = 2^(p+q-2)/|w|_1 B(al, be)/B(a, b), q = (r - 5) mod 2,
    al = (r - 4 + q)/2, be = (J + p)/2, a = l/2, b = (J + 1)/2,
    |w|_1 = 2^p (q = 0) or 2 C(p, p//2) (q = 1);
  C5(p) = E cos 5u, u drawn as cos^(2be-1) sin^(2al-1) on [0, pi/2],
    E cos^n u = B(al, be + n/2)/B(al, be), cos 5u = 16c^5 - 20c^3 + 5c;
  A = E|1 - 4eU|, U ~ Beta(a, b).
The concentration TRANSPLANT reads B1 -> e/c0, C5 -> cos 5u0 = c0/9,
A -> sigma sqrt(2/pi), c0^2 = 1 - 1/(4e), sin^2 u0 = 1/(4e): 1.35 at
r = 21 and growing as sqrt(l), where the run reads 2.67, flat.

BY HAND (re-derived before the code). 5u0 = 1.4642, pi/2 - 5u0 = 0.1066:
cos 5u sits near its zero with slope about -1, so an O(1/l) move of u is
O(1) against cos 5u0. Two O(1/l) moves, both down: P5's own mean,
al/(al + be) - 1/(4e) = -(8e - 3)/(16e^2 l) + O(l^-2), and Jensen on
u = arcsin sqrt U (u'' = -(1 - 2U)/(4 (U(1 - U))^(3/2))); together, at
odd r and p = 0, du = -(5e - 2)/(4e sqrt(4e - 1) l), -0.327/l at e = 3
(a larger p moves it further down). The corrected hand
form C5h = cos(5 ubar) exp(-25 Var u/2), ubar = arcsin sqrt(Ubar) +
u''(Ubar) Var U/2, Var u = Var U/(4 Ubar(1 - Ubar)), Ubar and Var U those
of P5's own Beta(al, be). At r = 21: C5h = 0.190 against c0/9 = 0.106,
B1's lead (a + b)^(3/2)/(4 (a - 1) sqrt b) = 3.502 against e/c0 = 3.133,
A ~ 0.247, product 2.69. At fixed e the limit grows as sqrt(l). Even r
(q = 1): B1's lead 2^(p-2)/C(p, p//2) sqrt((a + b)/a) (b/(a + b))^((p-1)/2)
peaks near p ~ 4e, and a larger p moves Ubar down, so p* is well above 1.

PREDICTIONS (frozen before the code). P1 odd r: p* = 0; B1/B1lead in
[0.98, 1.02]; A/(sigma sqrt(2/pi)) in [0.95, 1.05], sigma = 4e sqrt(Var U).
P2 C5/(c0/9) >= 1.4 at r = 21 and falling in r. P3 C5/C5h in [0.95, 1.05]
at every r = 21..51, both parities, at p*. P4 even r: p* odd in 7..15;
B1/B1lead in [0.94, 1.06].
KILL. Printed per r: C5/C5h. Outside [0.95, 1.05] at any r fires: the
O(1/l) moves are not the whole correction. Read beside A's and B1's
ratios: either outside its band at odd r and the break is not the cos
factor alone.
CONTROLS, read first. C1 the product B1 |C5|/A against exacttest's
exp(log_x - log_mean) at the same (shape, p) at every r run, and at r =
21, 22 against the exact integer pairing over |w|_1 and the mean:
relative difference < 1e-9. PLANT (sharpened after the run, the frozen cos 3u
for cos 5u being caught by construction): be + 1/2 for be in B1 and C5
must move the product off exacttest by more than 1e-9 at every r. C1
tests the split of X into B1 C5; A cancels
from both sides (log_mean runs circlemean.mad, the same A), so A rests on
circlemean's own controls. p* is exacttest.best_x's p.

FINDINGS (copied from the printed output).
  (1) C1 holds at every r = 21..51: product/exacttest - 1 at most
      3.5e-13, the exact integer pairing at r = 21, 22 to 5.7e-31 and
      6.0e-31, the be + 1/2 plant caught at every r (2.3e-02 .. 3.3e-02
      off). The odd rows 21..33 reproduce exacttest's per-rank least
      (2.6738 at r = 21 .. 2.7100 at r = 33), the even rows 38..50
      evenarm's (2.1747 .. 2.2927).
  (2) The KILL misses: C5/C5h 1.0019 (r = 21) down to 1.0004 (r = 51),
      both parities. The transplant's break is the cos factor: C5/(c0/9)
      1.7861 at r = 21 falling to 1.3024 at r = 51 (odd), 1.9126 to
      1.3256 (even), while A/(sigma sqrt(2/pi)) reads 0.9955 .. 0.9983
      and B1/(e/c0) 1.1033 .. 1.0374 at odd r (its lead's 1/(a - 1),
      B1/B1lead 0.9871 .. 0.9950). P5's mean at p = 0
      sits below 1/(4e) by 0.00778 (r = 21) .. 0.00299 (r = 51).
  (3) X/mean at odd r: 2.6738, 2.6630, 2.6617 (r = 21, 23, 25), then
      rising to 2.9281 at r = 51: the 1/l excess of C5 decays while
      1/A grows as sqrt(l). At even r 2.1043 (r = 22), 2.1000 (24), then
      rising to 2.2927 (50).
  (4) P1 holds: p* = 0 at every odd r, B1/B1lead 0.9871 .. 0.9950,
      A/Ag 0.9955 .. 0.9983. P2 holds (1.7861 at r = 21, falling). P3
      holds. P4 fails: p* is 19, 17, 17, 17 at r = 22..28, above 15,
      then 15 (r = 30..44) and 13 (r = 46..50), falling; B1/B1lead is
      1.1030, 1.0838, 1.0766, 1.0705 at r = 22..28, above 1.06, falling
      to 1.0300 at r = 50.
READS. At fixed e both parities grow as sqrt(l): X/mean ~ K_q cos(5u0)
/(sigma sqrt(2/pi)), sigma^2 ~ (4e - 1)/(2el), K_0 = e/c0 and K_1 the
largest over odd p of B1's lead, which tends to e sqrt(pi/2) exp(-1/2)
as e grows (at p ~ 4e), the takeover theorem's even-rank K. The even p*
falls toward the lead's own peak as the 1/l move of Ubar that a larger p
buys decays. For k = 0 2^(r-g) g, g >= 5, at any p the 1/l move is down
in u (argued, not printed: the mean's shift has the sign of
(4e - 1)(3 - g + q) + 1 - p, positive at g = 3 and small p, and Jensen
is down at U < 1/2):
where g u0 sits just below pi/2 it lifts |cos g u|, as at e = 3; just
above (5u0 = 1.807 at e = 2, 7u0 = 1.769 at e = 4) it shrinks |cos g u|,
which need not make the shape weak (e = 4's part-7 shape stands 9.86 ..
14.03 times the mean at r = 39..51, exacttest's print).

RUN RECORD. 0.2 s, peak commit 20.3 MB.
"""
import math
import os
import sys
import time

os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
from mpmath import mp  # noqa: E402
import circlemean  # noqa: E402
import compbound  # noqa: E402
import exacttest  # noqa: E402

mp.dps = 40
E = 3
CHEB = {5: (5, -20, 16)}   # odd T_g on c^1, c^3, ...


def factors(r, p, g=5, plant=False):
    """B1, C = E cos g u, A and the hand forms at (r, p); plant: be + 1/2."""
    l = r - 2
    J = (4 * E - 1) * l - 1
    q = (r - 5) % 2
    a, b = mp.mpf(l) / 2, mp.mpf(J + 1) / 2
    al, be = mp.mpf(r - 4 + q) / 2, mp.mpf(J + p + plant) / 2
    lb = circlemean.logB
    B1 = mp.exp((p + q - 2) * mp.log(2) - compbound.log_w1(p, q)
                + lb(al, be) - lb(a, b))
    cs = CHEB[g]
    C = sum(v * mp.exp(lb(al, be + mp.mpf(2 * i + 1) / 2) - lb(al, be))
            for i, v in enumerate(cs))
    A = circlemean.mad(a, b, 4 * E)
    U = al / (al + be)
    VU = al * be / ((al + be) ** 2 * (al + be + 1))
    upp = -(1 - 2 * U) / (4 * (U * (1 - U)) ** mp.mpf(1.5))
    ubar = mp.asin(mp.sqrt(U)) + upp * VU / 2
    Vu = VU / (4 * U * (1 - U))
    Ch = mp.cos(g * ubar) * mp.exp(-g * g * Vu / 2)
    if q == 0:
        lead = mp.mpf(2) ** (p - 2) / 2 ** p / (a - 1) \
            * b ** (mp.mpf(p - 1) / 2) * (a + b) ** (-mp.mpf(p - 3) / 2)
    else:
        lead = mp.mpf(2) ** (p - 2) / math.comb(p, p // 2) \
            * mp.sqrt((a + b) / a) * (b / (a + b)) ** (mp.mpf(p - 1) / 2)
    VF = a * b / ((a + b) ** 2 * (a + b + 1))
    Ag = 4 * E * mp.sqrt(VF) * mp.sqrt(2 / mp.pi)
    return {"J": J, "q": q, "B1": B1, "C": C, "A": A, "Ch": Ch,
            "lead": lead, "Ag": Ag, "dU": U - mp.mpf(1) / (4 * E),
            "prod": B1 * abs(C) / A}


def main():
    t0 = time.time()
    c0 = mp.sqrt(1 - mp.mpf(1) / (4 * E))
    lim_B, lim_C = E / c0, c0 / 9
    print("transplant limits: B1 -> e/c0 = %.4f, C5 -> c0/9 = %.4f"
          % (lim_B, lim_C))
    ok, rows = True, []
    for r in range(21, 52):
        J = (4 * E - 1) * (r - 2) - 1
        parts = (5,) + (2,) * (r - 5)
        v, p = exacttest.best_x(J, 0, parts)
        lm = circlemean.log_mean(r, J, E)
        f = factors(r, p)
        ref = mp.exp(mp.mpf(v) - lm)
        d1 = abs(f["prod"] / ref - 1)
        pl = factors(r, p, plant=True)["prod"]
        caught = abs(pl / ref - 1) > 1e-9
        d2, ex = mp.mpf(0), ""
        if r in (21, 22):
            q = f["q"]
            fr = exacttest.exact_pairing(J, 0, parts, p)
            X = mp.mpf(fr.numerator) / fr.denominator \
                / mp.exp(compbound.log_w1(p, q))
            d2 = abs(f["prod"] / (X / mp.exp(lm)) - 1)
            ex = ", exact pairing %.1e" % d2
        good = d1 < 1e-9 and d2 < 1e-9 and caught
        ok &= good
        rows.append((r, p, f))
        print("C1 r = %d p* = %d: product/exacttest - 1 %.1e%s, be + 1/2"
              " plant %.1e %s  %s"
              % (r, p, d1, ex, abs(pl / ref - 1),
                 "caught" if caught else "MISSED",
                 "holds" if good else "FAILS"), flush=True)
    print("C1 %s  [%.1f s]" % ("holds" if ok else "FAILS", time.time() - t0))
    print()
    print(" r  p*  X/mean  B1     B1/lead  B1/lim  C5      C5/lim  C5/C5h"
          "  A       A/Ag    dU")
    kill = []
    for r, p, f in rows:
        k = f["C"] / f["Ch"]
        if not 0.95 <= k <= 1.05:
            kill.append(r)
        print("%2d %3d  %.4f  %.4f %.4f   %.4f  %.4f  %.4f  %.4f  %.4f"
              "  %.4f  %+.5f"
              % (r, p, f["prod"], f["B1"], f["B1"] / f["lead"],
                 f["B1"] / lim_B, f["C"], f["C"] / lim_C, k, f["A"],
                 f["A"] / f["Ag"], f["dU"]))
    print("KILL %s" % ("fires at r = %s" % kill if kill else "misses"))
    print("%.1f s" % (time.time() - t0))
    return ok


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