"""
Rigorous enclosure of t^2 * E(xi) for the explicit second-order Ramanujan theorem.

Every quantity is an arb ball; every inequality used is listed in Appendix A of second_order_ramanujan.tex
with the same labels as Appendix A of second_order_ramanujan.tex (U1..U9, M1..M7, A1, A2).  Assumes RH (explicit formula for psi,
Schoenfeld's bound) and the unconditional bounds of Buthe (Math. Comp. 2018, Thm 2).

Parametrisation: real xi >= 1e9, t = log xi, s = 1/t, eps = F(xi,1) = log(1+1/xi)/t,
xi_k = root of F(z,k) = eps, N = prod_k prod_{p<=xi_k} p  (this is every CA number).
"""
from flint import arb, ctx
import sympy as sp, pickle, mpmath as mm

ctx.prec = 160
A = arb
ZERO = A(0)
def iv(a, b):
    a = A(a); b = A(b)
    return a.union(b)
def hull(*xs):
    r = xs[0]
    for y in xs[1:]: r = r.union(y)
    return r

PI = A.pi(); LOG2 = A(2).log(); SQ2 = A(2).sqrt(); GAMMA = A.const_euler()
BETA = 2 + GAMMA - (4*PI).log()          # sum 1/|rho|^2
GAMMA1 = A("14.134725")                   # lower bound for |rho| = |1-rho|, first zero 14.1347251...
C = A("2.0529")                           # sup_{0<y<=1e19}(y-theta(y))/sqrt(y)  [cconst.py: 2.05281.. on (0,1423]; Buthe 1.95 on [1423,1e19]]
Y = A(10)**19                             # Buthe range end
LOGY = Y.log()
THETA_UP = A(1) + A("1e-7")               # theta(y) <= THETA_UP*y for all y>0 (Buthe (1.7) + Schoenfeld)
C3 = A(3).root(3)
KAPPA = SQ2*(2+LOG2) - 2
V_LO, V_HI = A("0.47"), A("0.55")         # Lemma (shell 2): v in [V_LO,V_HI]
L = LOG2

def E1(z):            # exponential integral E_1(z), z>0, rigorous (arb_hypgeom)
    return z.expint(A(1))

def Qenv(z):          # z e^z E1(z) in [1-1/z, 1]  (enveloping asymptotic series)
    return iv(1 - 1/z, A(1))

def Qser(z, n=6):
    """z e^z E1(z) = sum_{k<n} (-1)^k k!/z^k + R, R between 0 and (-1)^n n!/z^n  (z>0)"""
    r = A(0); f = A(1); zi = 1/z; p = A(1)
    for k in range(n):
        r += (-1)**k*f*p
        f *= (k+1); p *= zi
    rem = (-1)**n*f*p
    return r + iv(min(rem.lower(),0), max(rem.upper(),0))

def Sschoen(y):       # Schoenfeld: |theta(y)-y| <= Sschoen(y) sqrt(y), y>=599, under RH
    return y.log()**2/(8*PI)

def eta(y):
    """theta(y)-y in [-etaL sqrt y, etaU sqrt y] for the ball y (y>0)."""
    if y < Y:
        return C, ZERO
    S = Sschoen(A(y.upper()))
    if y > Y:
        return S, S
    return (C if C > S else S), S

def theta_iv(y):
    """enclosure of theta(y) for the ball y>0"""
    eL, eU = eta(y)
    lo = y - eL*y.sqrt()
    hi = y + eU*y.sqrt()
    return iv(max(lo.lower(), 0), hi.upper())

# ---------------------------------------------------------------- sympy polynomials
_sym = pickle.load(open(__file__.replace('engine.py', 'sym.pkl'), 'rb'))
_s, _v, _mu, _kap = sp.symbols('s v mu kap')
def _compile(expr):
    expr = sp.sympify(sp.parse_expr(expr) if isinstance(expr, str) else expr)
    poly = sp.Poly(sp.expand(expr), _s, _v, _mu, _kap)
    terms = []
    for (i, j, k, m), c in poly.terms():
        cs = sp.nsimplify(c)
        terms.append(((i, j, k, m), cs))
    return terms
def _coef(c):
    # exact sympy coefficient in Q(log2, sqrt2) -> arb
    c = sp.expand(c)
    val = A(0)
    for term in sp.Add.make_args(c):
        num = A(1)
        for f in sp.Mul.make_args(term):
            if f.is_Rational: num *= A(int(f.p))/A(int(f.q))
            elif f == sp.log(2): num *= LOG2
            elif f.is_Pow and f.base == sp.log(2): num *= LOG2**int(f.exp)
            elif f == sp.sqrt(2): num *= SQ2
            elif f.is_Pow and f.base == 2 and f.exp == sp.Rational(1, 2): num *= SQ2
            elif f.is_Pow and f.base == 2 and f.exp == sp.Rational(-1, 2): num /= SQ2
            elif f.is_Pow and f.base == 2 and f.exp.is_Rational: num *= A(2)**A(int(f.exp.p))/A(1) if f.exp.q == 1 else (A(2)**(A(int(f.exp.p))/int(f.exp.q)))
            else: raise ValueError(("coef", f))
        val += num
    return val
def _load(name):
    expr = sp.parse_expr(_sym[name], local_dict={'s': _s, 'v': _v, 'mu': _mu, 'kap': _kap}) if False else eval(_sym[name], sp.__dict__ | {'Symbol': sp.Symbol, 'Integer': sp.Integer, 'Rational': sp.Rational})
    poly = sp.Poly(sp.expand(expr), _s, _v, _mu, _kap)
    return [((i, j, k, m), _coef(c)) for (i, j, k, m), c in poly.terms()]
P1T = _load('P1'); NUMT = _load('NumS2'); DENT = _load('Den')
def peval(terms, s, v, mu, kap):
    # evaluate grouped in s (Horner by collecting coefficient balls per power of s)
    from collections import defaultdict
    byk = defaultdict(lambda: A(0))
    for (i, j, k, m), c in terms:
        byk[i] += c * v**j * mu**k * kap**m
    deg = max(byk)
    r = A(0)
    for i in range(deg, -1, -1):
        r = r*s + byk.get(i, A(0))
    return r

def mu_iv(s, v):
    # mu = (log(1+w)-w)/s^2 in [-(l/2-vs)^2/(2(1+w)^2), -(l/2-vs)^2/2],  w = -l s/2 + v s^2 < 0
    a2 = (L/2 - v*s)**2
    w = -L/2*s + v*s**2
    assert (-L/2 + v*s) < 0     # w = s(-l/2 + v s) <= 0
    return iv((-a2/(2*(1+w)**2)).lower(), (-a2/2).upper())

KAP = iv(0, 1)
VV = iv(V_LO, V_HI)

def P1_eval(s):
    return peval(P1T, s, ZERO, ZERO, KAP)

def Mcomb_eval(s):
    v = VV; mu = mu_iv(s, v)
    return -peval(NUMT, s, v, mu, KAP)/peval(DENT, s, v, mu, KAP) - SQ2*v

# ---------------------------------------------------------------- shells
def Fk(z, k):     # F(z,k) = log(1+1/(z+...+z^k))/log z
    S = z*(z**k - 1)/(z - 1)
    return (1/S).log1p()/z.log()

def eps_of(xi):
    return (1/xi).log1p()/xi.log()

mm.mp.dps = 60
_rootcache = {}
def xi_k_point(ta, k, side):
    """certified bound for xi_k at xi = e^{ta} (exact real point): side='lo' -> z with z <= xi_k, 'hi' -> z >= xi_k"""
    key = (ta, k, side)
    if key in _rootcache: return _rootcache[key]
    t = mm.mpf(ta); eps = mm.log1p(mm.e**(-t))/t
    # solve log1p(1/S(e^u)) = eps*u for u = log z
    f = lambda u: mm.log1p(1/(mm.e**u*(mm.e**(k*u)-1)/(mm.e**u-1))) - eps*u
    u0 = (t + mm.log(k))/k
    u = mm.findroot(f, u0)
    xi = A(ta).exp(); epsA = eps_of(xi)
    for rel in ["1e-40", "1e-30", "1e-20", "1e-12", "1e-8"]:
        z = A(mm.nstr(mm.e**u, 55)) * (1 - A(rel) if side == 'lo' else 1 + A(rel))
        F = Fk(z, k)
        ok = (F > epsA) if side == 'lo' else (F < epsA)
        if ok:
            _rootcache[key] = z; return z
    raise RuntimeError(("root", ta, k))

def xk_point_lo(ta, k): return xi_k_point(ta, k, 'lo')
def xk_point_hi(tb, k): return xi_k_point(tb, k, 'hi')
def xi_k_cell(ta, tb, k):
    return iv(xi_k_point(ta, k, 'lo').lower(), xi_k_point(tb, k, 'hi').upper())

# ---------------------------------------------------------------- integrals
def UJ(b, T):
    """U * int_x^inf s^b w(s) ds = e^{-T(1/2-b)} [1 + b/(1-b) Q((1-b)T)]"""
    z = (1 - b)*T
    return (-(A(1)/2 - b)*T).exp() * (1 + b/(1 - b)*Qenv(z))

def USch(b, k, T, logX0):
    """U * int_{X0}^inf s^b (log s/k)^2/(8 pi) w(s) ds, with log X0 = logX0 >= T"""
    U = (T/2).exp()*T
    oneb = 1 - b
    X0pow = ((b - 1)*logX0).exp()
    return U/(8*PI*k*k) * X0pow * (logX0/oneb + 1/oneb**2 + 1/oneb)

def U_int_Ek(k, T):
    """U * int_x^inf E(s^{1/k}) w(s) ds  (E = theta - id)"""
    b = A(1)/(2*k)
    logYk = k*LOGY
    if T > logYk:
        neg = ZERO; logX0 = T
    else:
        neg = C*UJ(b, T); logX0 = logYk
    sch = USch(b, k, T, logX0)
    return iv((-neg - sch).lower(), sch.upper())

def int_sqrt_mg(m, Ly, Lz):
    """int_y^z sqrt(s) (-g_m'(s)) ds, g_m = s^{-m}/log s ; Ly=log y, Lz=log z (balls)"""
    a1 = A(m) - A(1)/2
    I1 = m*(E1(a1*Ly) - E1(a1*Lz))
    I2 = (-a1*Ly).exp()/Ly - (-a1*Lz).exp()/Lz - a1*(E1(a1*Ly) - E1(a1*Lz))
    return I1 + I2

def int_sqrt_log2_mg(m, Ly, Lz):
    """int_y^z sqrt(s) log^2 s (-g_m') ds = int (m log s + 1) s^{-m-1/2} ds"""
    a1 = A(m) - A(1)/2
    def F(Lq):
        p = (-a1*Lq).exp()
        return p*(m*(Lq/a1 + 1/a1**2) + 1/a1)
    return F(Ly) - F(Lz)

def int_E_mg(m, Ly, Lz):
    """enclosure of int_y^z E(s) (-g_m'(s)) ds  (integrands >= 0 are over-integrated when limits straddle Y)"""
    lo, hi = ZERO, ZERO
    if not (Ly > LOGY):
        Lz1 = Lz if (Lz < LOGY) else LOGY
        cpart = C*int_sqrt_mg(m, Ly, Lz1)
        lo -= A(max(cpart.upper(), 0))
    if not (Lz < LOGY):
        Ly1 = Ly if (Ly > LOGY) else LOGY
        spart = int_sqrt_log2_mg(m, Ly1, Lz)/(8*PI)
        spart = A(max(spart.upper(), 0))
        lo -= spart; hi += spart
    return iv(lo.lower(), hi.upper())

def Ebox(y):
    eL, eU = eta(y)
    return iv((-eL*y.sqrt()).lower(), (eU*y.sqrt()).upper())

def prime_sum(m, y, z):
    """enclosure of sum_{y<p<=z} p^{-m}  (y<z balls)"""
    if not (y.lower() < z.upper()): return ZERO
    Ly, Lz = y.log(), z.log()
    main = E1((m - 1)*Ly) - E1((m - 1)*Lz)
    g = lambda q, Lq: (-m*Lq).exp()/Lq
    bd = g(z, Lz)*Ebox(z) - g(y, Ly)*Ebox(y)
    tot = main + bd + int_E_mg(m, Ly, Lz)
    return iv(max(tot.lower(), 0), tot.upper())

def prime_tail_up(m, y):
    """sum_{p>y} p^{-m} <= THETA_UP y^{1-m} (m/((m-1)log y) + 1/((m-1) log^2 y))"""
    Ly = y.log()
    return THETA_UP*((1 - m)*Ly).exp()*(A(m)/((m - 1)*Ly) + 1/((m - 1)*Ly**2))

# ---------------------------------------------------------------- the cell
M_UP, M_LO = 60, 14
KDEF = 7              # deficit prime sums for shells k=3..KDEF, stationarity bound below xi_{KDEF+1}
KROOT = 6             # certified roots for xi_3..xi_KROOT

def cell(ta, tb, detail=False):
    T = iv(A(ta), A(tb)); s = 1/T; T2 = T*T
    X = T.exp(); sqx = (T/2).exp(); U = sqx*T
    out = {}
    # ---------- xi, e = E(x)/x
    eX = Ebox(X)/X
    # ---------- U-group  (delta_u)
    out['U1_P1'] = P1_eval(s)
    z3 = 2*T/3
    out['U2_cube'] = T2*(-T/6).exp()/2*(Qser(z3) - 1)
    out['U3_sq_err'] = T2*U_int_Ek(2, T)
    out['U4_cu_err'] = T2*U_int_Ek(3, T)
    up = A(0); lo = A(0)
    import math as _m
    MU = max(M_UP, int(_m.ceil(2*float(tb)/_m.log(2))))
    for k in range(4, MU + 1):
        bk = A(1)/k
        termU = UJ(bk, T)
        logYk = k*LOGY
        logX0 = T if (T > logYk) else logYk
        sch = USch(A(1)/(2*k), k, T, logX0)
        up += termU + sch
        if k <= M_LO:
            lo += termU - C*UJ(A(1)/(2*k), T) - sch
    # k > MU: theta(s^{1/k}) = 0 unless s >= 2^k >= 2^{MU+1} >= x^2;  sum <= THETA_UP (log s/log 2) s^{1/(MU+1)}
    bM = A(1)/(MU + 1); LX0 = (MU + 1)*LOG2
    up += U*THETA_UP/LOG2*(1 + 1/LX0)*((bM - 1)*LX0).exp()/(1 - bM)
    lo = A(max(lo.lower(), 0))
    out['U5_k4plus'] = T2*iv(lo.lower(), up.upper())
    rz = 2*BETA*(1 + 6*s)/GAMMA1
    out['U6_zeros'] = iv((-rz).lower(), rz.upper())
    out['U7_const'] = T2*iv((PI*2).log()*(-T/2).exp() - (-5*T/2).exp(), (PI*2).log()*(-T/2).exp())
    a = A(eX.abs_upper())
    assert a < A(1)/2 and 2*a/T < A(1)/2
    out['U8_bracket'] = T2*iv(0, (U*(2*a*a/T + 8*a*a/T2)).upper())
    Tp = (1 + 1/(X - 1))/2*THETA_UP*(-T/2).exp()*(2 + s)
    out['U9_tail'] = T2*iv((-Tp).lower(), 0)
    dU = sum(out[k] for k in out if k.startswith('U'))
    UG = 2 - 2*s + A(3)/2*(-T/6).exp() + dU/T2
    # ---------- shells
    w = -L/2*s + VV*s*s
    xi2 = (2*X).sqrt()*(1 + w)
    L2 = xi2.log()
    xk = {2: xi2}
    for k in range(3, KROOT + 1):
        xk[k] = xi_k_cell(ta, tb, k)
    # crude bounds for larger k:  xi^{1/k} <= xi_k <= (k xi)^{1/k}
    def xk_get(k):
        if k in xk: return xk[k]
        lo_ = (T/k).exp(); hi_ = ((T + A(k).log())/k).exp()
        return iv(lo_.lower(), hi_.upper())
    xi3 = xk[3]; L3 = xi3.log()
    # R4 = sum_{k>=4} theta(xi_k)
    R4 = A(0)
    k = 4
    tbA = A(tb)
    while True:
        hi_k = ((tbA + A(k).log())/k).exp()
        if hi_k < 2: break
        y = xk_get(k)
        th = theta_iv(y)
        th = iv(max(th.lower(), 0), th.upper())
        if y.upper() < 2: th = ZERO
        R4 += th
        k += 1
        assert k < 100000
    R0 = xi2 + xi3 + R4
    E2 = Ebox(xi2); E3 = Ebox(xi3)
    thX = X*(1 + eX)
    # M1 analytic
    out['M1_main'] = Mcomb_eval(s)
    # M2 cube residuals: -U E1(2 L3) + (1/2)3^{1/3} e^{-T/6}  and  -(B(R0)-xi2/sqrt x) + 3^{1/3} e^{-T/6}
    # nu3 = xi3/(3x)^{1/3}; U E1(2L3) = (3^{1/3}/2) e^{-T/6} nu3^{-2} Q(2L3)/D3,  D3 = 1 + s log3 + 3 s log nu3
    nu3 = iv((xk_point_lo(ta,3)/(3*A(tb).exp()).root(3)).lower(), (xk_point_hi(tb,3)/(3*A(ta).exp()).root(3)).upper())
    D3 = 1 + s*A(3).log() + 3*s*nu3.log()
    eT6 = (-T/6).exp()
    cubeA = C3/2*eT6*(1 - Qser(2*L3)/(nu3**2*D3))          # -U E1(2L3) + (1/2)3^{1/3}e^{-T/6}
    r2 = xi2/X; r4 = R4/X; r0 = r2 + C3*(-2*T/3).exp()*nu3 + r4
    q = r0/(1 + eX)
    phi = iv((1 - q/2).lower(), (1 - q/2 + q*q/3).upper())   # log(1+q)/q
    Psi = (1 + eX + r0).sqrt()*phi/(1 + eX)
    # B(R0) - xi2/sqrt x - 3^{1/3}e^{-T/6} = 3^{1/3}e^{-T/6}(nu3 Psi - 1) + sqrt(x) r4 Psi + sqrt2 (1+w)(Psi - 1)
    cubeB = C3*eT6*(nu3*Psi - 1) + sqx*r4*Psi + SQ2*(1 + w)*(Psi - 1)
    out['M2_cube'] = T2*(cubeA - cubeB)
    # M3 tiny positive: U E1(T) + U E1(2 L2)
    out['M3_tiny'] = T2*(U*E1(T) + U*E1(2*L2))
    # M4: -U g(x) E(x),  g(x) = x^{-2}/t
    out['M4_bdx'] = T2*(-U*Ebox(X)/(X*X*T))
    # M5 joint boundary terms at xi2, xi3
    Rstar = R0 + E2 + E3
    Bp = (1 + (Rstar/thX).log1p()/2)/(thX + Rstar).sqrt()
    g2 = 1/(xi2**2*L2); g32 = 1/(xi2**3*L2); g33 = 1/(xi3**3*L3)
    K2 = U*g2 - U*g32 - Bp
    K3 = U*g33 - Bp
    out['M5_joint'] = T2*(E2*K2 + E3*K3)
    # M6 integrated fluctuations on (xi2,x] and (xi3,xi2]
    int2 = int_E_mg(2, L2, T)
    int3 = int_E_mg(3, L3, L2)
    out['M6_fluct'] = T2*(-U*int2 - U*int3)
    # M7 higher-order losses
    A2x = prime_tail_up(4, xi2)/(1 - xi2**-2)
    A3x = prime_tail_up(6, xi3)/(1 - xi3**-3)
    A4lo = A(0); A4up = A(0)
    for kk in range(3, KDEF + 1):
        yk = xk_get(kk + 1); zk = xk_get(kk)
        Sk = prime_sum(kk + 1, yk, zk)
        A4lo += A(Sk.lower())
        A4up += A(Sk.upper())/(1 - yk**(-(kk + 1)))
    ysm = xk_get(KDEF + 1)
    epsU = 1/(X*T)
    A4up += A(64)/63*epsU*(ysm + ysm.log() + A("0.76"))
    out['M7_loss'] = T2*iv((-U*(A2x + A3x + A4up)).lower(), (-U*A4lo).upper())
    # ---------- assembly
    # UA crude (for Y only)
    S2 = E1(L2) - E1(T) + Ebox(X)/(X*X*T) - g2*E2 + int2
    S3 = E1(2*L3) - E1(2*L2) + g32*E2 - g33*E3 + int3
    UA = U*(S2 + S3 + iv(0, (A2x + A3x + A4up).upper()))
    Wb = iv(-BETA, BETA)
    Yv = Wb + Wb/T + UG - UA
    R = R0 + E2 + E3
    onepd = (1 + eX + R/X).sqrt()*(1 + eX.log1p()/T)
    d = onepd - 1
    out['A1_dY'] = T2*d*Yv
    yv = Yv/U
    ay = A(yv.abs_upper()); qmax = ay*ay*ay.exp()/2
    out['A2_quad'] = T2*U*onepd*iv(0, qmax.upper())
    tot = sum(out.values(), A(0))
    out['TOTAL'] = tot
    return out

def summary(o):
    return {k: (float(v.lower()), float(v.upper())) for k, v in o.items()}

if __name__ == "__main__":
    import math, time, sys
    t0 = math.log(1e9)
    for (a, b) in [(t0, t0 + 0.002), (23.0, 23.002), (30, 30.01), (60, 60.05), (150, 150.5), (299, 300)]:
        st = time.time(); o = cell(a, b)
        print(f"cell [{a:.4f},{b:.4f}]  {time.time()-st:.2f}s")
        for k, (lo, hi) in summary(o).items(): print(f"   {k:12s} [{lo: .5f}, {hi: .5f}]")