"""
Plot an approximation to  pi(x) - li(x)  for x near 1.398 x 10^316,
using the explicit formula with the first 2,001,052 zeta zeros (Odlyzko).

Quantity plotted, with u = log x and assuming the Riemann hypothesis:

  D(u) = (pi(x) - li(x)) * u / sqrt(x)
       ~ -(1 + 2/u + 8/u^2)
         - 2 * sum_{0 < gamma <= T} w(gamma) * Re[ e^{i gamma u} / rho * (1 + 1/(rho u)) ],
  rho = 1/2 + i gamma.

The first term is the contribution of prime squares, (u/sqrt x)(1/2) li(sqrt x),
expanded to order 1/u^2. Each zero contributes li(x^rho), expanded to second order.
Terms like li(x^(1/3)) are smaller by a factor of about x^(-1/6) and are dropped.
So D > 0 means pi(x) > li(x), and D is measured in units of sqrt(x)/log x.

Two versions are drawn:
  sharp    : w = 1, all 2,001,052 zeros (T ~ 1.13e6).
             Pointwise approximation; the RMS error from the omitted zeros is
             about sqrt(log(T/2pi)/(pi T)) ~ 0.002 in these units.
  smoothed : the Bays-Hudson choice w = exp(-gamma^2/(2 alpha)), alpha = 1e10,
             1,000,000 zeros. This is a Gaussian average of D over a width of
             about 1/sqrt(alpha) ~ 1e-5 in u.

The sum over zeros is evaluated on a uniform grid with a nonuniform FFT (finufft):
  sum_j c_j exp(i gamma_j (u_c + k h)) = sum_j [c_j e^{i gamma_j u_c}] e^{i k (gamma_j h)},
which is a "type 1" NUFFT with nonuniform points gamma_j h and integer modes k.

Requirements: numpy, mpmath, matplotlib, finufft, and a zeros file.
Get zeros6 from https://www-users.cse.umn.edu/~odlyzko/zeta_tables/ (or the
PyPI package passagemath-database-odlyzko-zeta) and save the imaginary parts as
zeros.npy, or as zeros6 text with one zero per line.
"""
import sys, time
import numpy as np
import mpmath as mp
import finufft
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

mp.mp.dps = 40


def load_zeros(path="zeros.npy"):
    if path.endswith(".npy"):
        return np.load(path)
    return np.loadtxt(path)


def D_on_grid(gammas, x_lo, x_hi, npts, alpha=None, eps=1e-13):
    """Return x/10^316 values, u values, and D(u) on a uniform grid in u."""
    u_lo = mp.log(mp.mpf(x_lo))
    u_hi = mp.log(mp.mpf(x_hi))
    u_c = (u_lo + u_hi) / 2                     # high-precision window centre
    h = float((u_hi - u_lo) / npts)
    k = np.arange(-npts // 2, npts - npts // 2)  # finufft mode ordering

    g = gammas
    rho = 0.5 + 1j * g
    uc_f = float(u_c)
    w = np.ones_like(g) if alpha is None else np.exp(-g**2 / (2 * alpha))
    coef = w / rho * (1 + 1 / (rho * uc_f))     # 1/(rho u) varies by ~1e-6 over the window

    # Phase gamma*u_c mod 2pi.  Split u_c = a + b, with a having few significant
    # bits so that gamma*a is (nearly) exact in floating point; remaining error ~1e-9 rad.
    a = float(mp.floor(u_c * 2**20) / 2**20)
    b = float(u_c - a)
    phase = np.mod(g * a, 2 * np.pi) + g * b
    c = coef * np.exp(1j * phase)

    S = finufft.nufft1d1(g * h, c.astype(np.complex128), npts, eps=eps, isign=+1)
    u = uc_f + k * h
    D = -(1 + 2 / u + 8 / u**2) - 2 * S.real
    xs = np.array([float(mp.e**(u_c + int(kk) * mp.mpf(h)) / mp.mpf(10)**316) for kk in k])
    return xs, u, D


def main(zeros_path="zeros.npy", prefix="skewes"):
    g_all = load_zeros(zeros_path)
    print(f"{len(g_all)} zeros, largest {g_all[-1]:.3f}")

    panels = [  # (x_lo, x_hi, npts, title, filename)
        ("1.3977e316", "1.3989e316", 8000, r"$(\pi(x) - \mathrm{li}(x)) \cdot \log x \,/\, \sqrt{x}$ for $1.3977 \times 10^{316} \leq x \leq 1.3989 \times 10^{316}$", f"{prefix}_narrow.png"),
        ("1.3980e316", "1.3984e316", 8000, r"$(\pi(x) - \mathrm{li}(x)) \cdot \log x \,/\, \sqrt{x}$ for $1.3980 \times 10^{316} \leq x \leq 1.3984 \times 10^{316}$", f"{prefix}_narrower.png"),
        ("1.3955e316", "1.3995e316", 16000, r"$(\pi(x) - \mathrm{li}(x)) \cdot \log x \,/\, \sqrt{x}$ for $1.3955 \times 10^{316} \leq x \leq 1.3995 \times 10^{316}$", f"{prefix}_wide.png"),
    ]
    for lo, hi, n, title, out in panels:
        fig, ax = plt.subplots(figsize=(12, 4.5))
        t = time.time()
        xs, u, D = D_on_grid(g_all, lo, hi, n)
        xs2, _, D2 = D_on_grid(g_all[:1_000_000], lo, hi, n, alpha=1e10)
        print(f"{title}: {time.time()-t:.1f}s; max D = {D.max():.4f} at x = {xs[D.argmax()]:.7f}e316")
        ax.plot(xs, D, lw=0.6, color="C0", label="2,001,052 zeros, no smoothing")
        ax.plot(xs2, D2, lw=1.2, color="C3",
                label="10⁶ zeros, Gaussian smoothing α = 10¹⁰ (as in Bays–Hudson)")
        ax.axhline(0, color="k", lw=0.8)
        ax.fill_between(xs, 0, D, where=D > 0, color="C2", alpha=0.35, lw=0,
                        label="π(x) > li(x)")
        ax.set_xlim(xs[0], xs[-1])
        ax.ticklabel_format(axis="x", useOffset=False)
        ax.set_title(title)
        ax.set_xlabel("x / 10³¹⁶")
        ax.set_ylabel("(π(x) − li(x)) · log x / √x")
        ax.legend(loc="lower right", fontsize=8)
        fig.tight_layout()
        fig.savefig(out, dpi=200)
        plt.close(fig)
        print("wrote", out)


if __name__ == "__main__":
    main(*sys.argv[1:])