"""Rigorous two-sided bounds for the CDF  C_n(y) = P(S_n <= y),  S_n = X_1+...+X_n,  X ~ g^2/m2 on [0,T],
g(t) = 1/(c+bt),  c = theta/log k, T = beta/log k, b = k-1  (theta, beta exact binary64 numbers; theta = gamma).

    C_n(y) = (1/pi) Re int_0^inf phi(xi)^n K_y(xi) dxi,   K_y(xi) = (1 - e^{-i xi y})/(i xi)     (Lemma 5.1).

[0, Xi1]: composite Gauss-Legendre, m = 40 nodes on each unit interval; nodes rounded to binary64 (node-shift
          term, using |d/dxi (phi^n K_y)| <= n mu y + y^2/2), phi^n at the nodes in Arb; binary64 summation with a
          rigorous rounding budget.  Quadrature error per unit interval, for an integrand bounded by M on the
          Bernstein ellipse with parameter rho = 1 + sqrt(2) (inside the strip |Im xi| <= eta0 = 1/2):
              (1/2) (64/15) M rho^{-2m} / (1 - rho^{-2}),   M = (1 + mu (e^{eta0 T}-1)/T)^n * y e^{eta0 y}
          (Trefethen, SIAM Rev. 50 (2008), Thm 4.5).
[Xi1, inf): Lipschitz cells (|phi'| <= mu = E X), then the second-mean-value bound
          |phi(xi)| <= F(a) + 2 sqrt(2) f(a)/xi with a = 4 sqrt(2)/xi on geometric cells, then the power law
          |phi(xi)| <= 6 sqrt(2) f(0)/xi; always |K_y| <= 2/xi, so the tail bound is independent of y.
The binary64 summation assumes that numpy's sin and cos are accurate to within four ulp (see sincos_check.py)."""
import sys, math, time
import numpy as np
from flint import arb, acb, ctx
import os; sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from phi import make_phi
from fcert import setup

U = 2.0 ** -53

class CDF:
    def __init__(self, k, theta, beta, ns, Xi1=None, m=40, eta0=0.5, verbose=False):
        t0 = time.time()
        ctx.prec = 256
        self.k, self.theta, self.beta, self.ns = k, theta, beta, tuple(ns)
        S = setup(k, theta, beta); self.S = S
        c, T, b, m2, mu = S["c"], S["T"], S["b"], S["m2"], S["mu"]
        self.Tf = float(T.mid())
        phi, _ = make_phi(k, theta, beta); self.phi = phi
        nmin = min(ns)
        if Xi1 is None:
            Xi1 = 50
            while True:
                p = phi(acb(arb(Xi1)))
                lnv = float(p.abs_upper().log().mid()) * nmin if float(p.abs_upper()) < 1 else 0.0
                if lnv < -130 or Xi1 > 60000:
                    break
                Xi1 = int(Xi1 * 1.2) + 1
        self.Xi1 = Xi1
        gl = [arb.legendre_p_root(m, i, weight=True) for i in range(m)]
        xs, ws, wr, shift = [], [], [], 0.0
        A = {n: ([], [], []) for n in ns}
        for a0 in range(0, Xi1):
            for (x, wgt) in gl:
                xb = arb(a0) + (x + 1) / 2
                xf = float(xb.mid())
                dxi = float((arb(xf) - xb).abs_upper())
                ww = wgt / 2
                p = phi(acb(arb(xf)))
                xs.append(xf); ws.append(float(ww.mid())); wr.append(float(ww.rad()) + abs(float(ww.mid())) * U)
                shift += float(ww.upper()) * dxi
                for n in ns:
                    pp = p ** n
                    R, I, RR = A[n]
                    R.append(float(pp.real.mid())); I.append(float(pp.imag.mid()))
                    RR.append(float(pp.real.rad()) + float(pp.imag.rad()) + 2 * U * (abs(float(pp.real.mid())) + abs(float(pp.imag.mid()))))
        # float accumulation of the node-shift sum: inflate by the recursive-summation bound (referee note)
        shift *= (1 + 4 * len(xs) * U)
        self.xi = np.array(xs); self.w = np.array(ws); self.wr = np.array(wr); self.shift = shift
        self.A = {n: tuple(np.array(v) for v in A[n]) for n in ns}
        self.t_nodes = time.time() - t0
        # quadrature constants
        s_ = 4 * eta0
        self.rho = (s_ + math.sqrt(s_ * s_ + 4)) / 2
        self.MG = 1 + mu * ((arb(eta0) * T).exp() - 1) / T
        self.m, self.eta0, self.mu = m, eta0, mu
        self.tail = {n: self._tail(n) for n in ns}
        self.t_all = time.time() - t0
        if verbose:
            print("CDF setup: Xi1=%d nodes=%d tails=%s secs=%.1f" % (Xi1, len(xs), {n: self.tail[n] for n in ns}, self.t_all), flush=True)

    def gl_err(self, n, ymax):
        y = arb(ymax)
        M = self.MG ** n * y * (arb(self.eta0) * y).exp()
        rho = arb(self.rho)
        # Gauss-Legendre error on [-1,1] for |f| <= M on the Bernstein ellipse E_rho:
        #   sum_{j >= 2m, j even} |a_j| |I(T_j) - I_m(T_j)| <= sum 2 M rho^-j (32/15) = (64/15) M rho^(-2m) / (1 - rho^-2)
        # (Trefethen's argument; the factor 1/(1 - rho^-2) = rho^2/(rho^2 - 1)); times 1/2 for a unit interval.
        e = self.Xi1 * arb(64) / 15 * M * rho ** (-2 * self.m) / (1 - rho ** (-2)) / 2
        e += arb(self.shift) * (arb(n) * self.mu * y + y * y / 2)
        return float(e.upper())

    def _tail(self, n):
        S, phi, mu, Tf = self.S, self.phi, self.mu, self.Tf
        alpha = arb(4) * arb(2).sqrt()
        tot = arb(0); xi = arb(self.Xi1); steps = 0       # xi is always an exact point
        while True:
            p = phi(acb(xi)).abs_upper()
            gap = 1 - p
            if gap <= 0:
                raise RuntimeError("phi too close to 1")
            dx = max(gap / mu / 4, arb(1))
            xn = arb((xi + dx).mid())                      # exact: consecutive cells share their endpoints
            tot += 2 * (xn / xi).log() * (p + mu * (xn - xi)) ** n
            xi = xn; steps += 1
            a = alpha / xi
            if float(a.upper()) <= Tf:
                Bv = S["F"](a) + 2 * arb(2).sqrt() * S["f"](a) / xi
                if float(Bv.upper()) < 1 and float((Bv.log() * n).upper()) < -150:
                    break
            if steps > 500000:
                raise RuntimeError("middle range too long")
        r = arb(1.01); f0 = S["f"](arb(0)); D = f0 * (alpha + 2 * arb(2).sqrt())
        for _ in range(200000):
            if float((D / xi).upper()) < 0.5:
                tot += 2 / arb(n) * (D / xi) ** n
                return float(tot.upper())
            xn = arb((xi * r).mid())                       # exact
            Bv = S["F"](alpha / xi) + 2 * arb(2).sqrt() * S["f"](alpha / xn) / xi
            Bv = Bv if float(Bv.upper()) < 1 else arb(1)
            tot += 2 * (xn / xi).log() * Bv ** n
            xi = xn
        raise RuntimeError("tail did not finish")

    def bounds(self, n, ys, chunk=128):
        """ys: binary64 array of points in (0, ymax]. Returns (L, U) arrays with L <= C_n(y) <= U (rigorous)."""
        ys = np.asarray(ys, dtype=float)
        ymax = float(ys.max())
        xi, w, wr = self.xi, self.w, self.wr
        Ar, Ai, Arad = self.A[n]
        absA = np.hypot(Ar, Ai) * (1 + 4 * U) + Arad
        N = len(xi)
        gam = 1.05 * N * U
        extra = self.gl_err(n, ymax) + self.tail[n]
        L = np.empty(len(ys)); Uu = np.empty(len(ys))
        inv = 1.0 / xi
        for i0 in range(0, len(ys), chunk):
            y = ys[i0:i0 + chunk][:, None]
            arg = xi[None, :] * y
            s = np.sin(arg); c1 = 1.0 - np.cos(arg)
            terms = w * (Ar * s + Ai * c1) * inv
            Ssum = terms.sum(axis=1)
            K_abs = np.minimum(y, 2.0 * inv[None, :])
            dK = 16 * U * (y + inv[None, :])
            err = (np.abs(w) * (Arad * K_abs + absA * dK) + wr * absA * K_abs).sum(axis=1) * (1 + 1e-12) \
                  + 2 * gam * np.abs(terms).sum(axis=1)
            tot_err = (err + extra) * (1 + 1e-12)
            # division by pi: math.pi < pi, so for x >= 0: x/pi <= x/math.pi <= (x/pi)(1 + 4e-17); one rounding
            # of the division and one of the product are covered by the factors (1 -+ 4e-16).
            lo = Ssum - tot_err
            L[i0:i0 + chunk] = np.where(lo > 0, lo / math.pi * (1 - 4e-16), 0.0)    # C >= 0 always
            Uu[i0:i0 + chunk] = (Ssum + tot_err) / math.pi * (1 + 4e-16)
        if np.any(Uu < 0):
            raise RuntimeError("negative upper bound for a CDF: inconsistent")
        return L, Uu