# Compare engine enclosures with exact values at the sieved checkpoints xi in [1e9,1e10].
import sys, math, mpmath as mm
sys.path.insert(0,'.')
import engine as E
mm.mp.dps = 40
zs = [mm.mpf(l) for l in open('../zeros.txt').read().split()]
TB = float(2+mm.euler-mm.log(4*mm.pi) - sum(2/(mm.mpf(1)/4+g*g) for g in zs))   # rigorous: needs only completeness of the zero list
def WW2(x):
    lx = mm.log(x); W = mm.mpf(0); W2 = mm.mpf(0)
    for g in zs:
        rho = mm.mpc(0.5, g); ph = mm.expj(g*lx)
        W += 2*mm.re(ph/(rho*(1-rho))); W2 += -2*mm.re(ph/(1-rho)**2)
    return W, W2
# small primes up to 2e5
NS = 300000; sv = bytearray([1])*(NS+1); sv[0]=sv[1]=0
for i in range(2,int(NS**.5)+1):
    if sv[i]: sv[i*i::i] = bytearray(len(sv[i*i::i]))
P = [p for p in range(NS+1) if sv[p]]
rows = [l.split() for l in open('../sieve/cp_1e9_1e10.txt')]
step = int(sys.argv[1]) if len(sys.argv)>1 else 20
bad = 0; out=[]
for r in rows[::step]:
    xi = int(r[0]); th = mm.mpf(r[1]); M = mm.mpf(r[2]); Q2 = mm.mpf(r[3])
    t = mm.log(xi); eps = mm.log1p(mm.mpf(1)/xi)/t
    def xik(k):
        f = lambda u: mm.log1p(1/(mm.e**u*(mm.e**(k*u)-1)/(mm.e**u-1))) - eps*u
        return mm.e**mm.findroot(f,(mm.log(1.5),(t+mm.log(k))/k+1),solver='anderson')
    xs = {}; k = 2
    while True:
        if mm.log1p(1/(2*(2**k-1)))/mm.log(2) < eps: break
        z = xik(k)
        xs[k] = z; k += 1
    xi2 = xs[2]
    a = {}
    for p in P:
        if p > xi2: break
        a[p] = max(k for k in xs if p <= xs[k])
    logN = th + sum(sum(mm.log(p) for p in P if p <= xs[k]) for k in xs)
    Lam = (Q2 - sum(mm.log(1-mm.mpf(p)**-2) for p in P if p <= xi2)) + sum(mm.log(1-mm.mpf(p)**(-a[p]-1)) for p in a)
    logsig = -M + Lam
    g = mm.euler
    c = (mm.e**(logsig-g) - mm.log(logN))*mm.sqrt(logN)          # c_sigma / e^gamma
    W, W2 = WW2(xi)
    main = 2-2*mm.sqrt(2) + W + (E.float(E.KAPPA) if False else (mm.sqrt(2)*(2+mm.log(2))-2) + W2)/t + 1.5*(1-mm.cbrt(3))*xi**(-mm.mpf(1)/6)
    Et2 = (c - main)*t**2
    # zero-free pieces
    U = mm.sqrt(xi)*t
    u = -(g + mm.log(mm.log(th)) + M)
    UG = U*u - W - W2/t
    du = (UG - 2 + 2/t - 1.5*xi**(-mm.mpf(1)/6))*t**2
    A_ = -Lam; R = logN - th
    B = mm.sqrt(logN)*mm.log(1+R/th)
    dsh = (-(U*A_ - mm.sqrt(2) + mm.sqrt(2)*(2+mm.log(2)/2)/t - mm.cbrt(3)/2*xi**(-mm.mpf(1)/6))
           -(B - mm.sqrt(2) + mm.sqrt(2)*mm.log(2)/(2*t) - mm.cbrt(3)*xi**(-mm.mpf(1)/6)))*t**2
    ta = float(t) - 1e-11; tb = float(t) + 1e-11
    o = E.cell(ta, tb)
    encU = sum((o[k] for k in o if k.startswith('U')), E.A(0))
    encM = sum((o[k] for k in o if k.startswith('M')), E.A(0))
    tot = o['TOTAL']
    trunc = TB*(1+1/t)*t**2   # zero truncation allowance (W,W2 with 3000 zeros), in t^2 units
    chk = lambda val, enc, tol: float(enc.lower()) - tol <= float(val) <= float(enc.upper()) + tol
    okU = chk(du, encU, trunc); okM = chk(dsh, encM, 1e-9); okT = chk(Et2, tot, trunc)
    bad += (not okU) + (not okM) + (not okT)
    print(f"xi={xi:.4e} t^2E={float(Et2):8.4f} enc=[{float(tot.lower()):8.3f},{float(tot.upper()):7.3f}] | du={float(du):7.3f} [{float(encU.lower()):7.3f},{float(encU.upper()):7.3f}] | shells={float(dsh):8.4f} [{float(encM.lower()):8.3f},{float(encM.upper()):7.3f}] {'OK' if okU and okM and okT else 'VIOLATION'}", flush=True)
print("violations:", bad)