#!/usr/bin/env python3
"""
Finite exact Heath-Brown K=3 and unit-sector crosscheck.

Checks:
1. K=3 coefficient identity;
2. Q0+Q1+Q2+Q3 = Lambda;
3. Q0 = Lambda on U-rough integers;
4. explicit divisor-polynomial formula for Q0.

Not evidence for RH.
"""
import math
import csv
import numpy as np

def mobius_sieve(n):
    primes=[]
    lp=np.zeros(n+1,dtype=int)
    mu=np.zeros(n+1,dtype=int)
    mu[1]=1
    for i in range(2,n+1):
        if lp[i]==0:
            lp[i]=i; primes.append(i); mu[i]=-1
        for p in primes:
            if p>lp[i] or i*p>n: break
            lp[i*p]=p
            mu[i*p]=0 if p==lp[i] else -mu[i]
    return mu

def von_mangoldt(n):
    lam=np.zeros(n+1,float)
    mark=np.ones(n+1,dtype=bool); mark[:2]=False
    for p in range(2,int(n**0.5)+1):
        if mark[p]:
            mark[p*p:n+1:p]=False
    for p in np.flatnonzero(mark):
        p=int(p); z=p; lp=math.log(p)
        while z<=n:
            lam[z]=lp
            if z>n//p: break
            z*=p
    return lam, mark

def conv(a,b,nmax):
    c=np.zeros(nmax+1,float)
    for d in range(1,nmax+1):
        if abs(a[d])<1e-18: continue
        for m in range(1,nmax//d+1):
            if abs(b[m])>1e-18:
                c[d*m]+=a[d]*b[m]
    return c

def divisor_j(nmax,j):
    one=np.ones(nmax+1,float); one[0]=0
    out=np.zeros(nmax+1,float); out[1]=1
    for _ in range(j):
        out=conv(out,one,nmax)
    return out

def run(N):
    X=2*N
    nmax=X-1
    U=math.ceil(X**(1/3))
    mu=mobius_sieve(nmax)
    lam, primes=von_mangoldt(nmax)

    one=np.ones(nmax+1,float); one[0]=0
    delta=np.zeros(nmax+1,float); delta[1]=1
    L=np.zeros(nmax+1,float)
    for n in range(1,nmax+1): L[n]=math.log(n)

    muU=np.zeros(nmax+1,float)
    for n in range(1,nmax+1):
        if n<=U: muU[n]=mu[n]
    nu=muU-delta

    muU2=conv(muU,muU,nmax)
    muU3=conv(muU2,muU,nmax)
    one2=conv(one,one,nmax)

    HB=3*conv(muU,L,nmax)-3*conv(conv(muU2,one,nmax),L,nmax)+conv(conv(muU3,one2,nmax),L,nmax)

    Q0=3*L-3*conv(one,L,nmax)+conv(one2,L,nmax)
    Q1=3*conv(nu,L,nmax)-6*conv(conv(nu,one,nmax),L,nmax)+3*conv(conv(nu,one2,nmax),L,nmax)
    nu2=conv(nu,nu,nmax)
    nu3=conv(nu2,nu,nmax)
    Q2=-3*conv(conv(nu2,one,nmax),L,nmax)+3*conv(conv(nu2,one2,nmax),L,nmax)
    Q3=conv(conv(nu3,one2,nmax),L,nmax)

    d2=divisor_j(nmax,2)
    d3=divisor_j(nmax,3)
    Q0poly=np.zeros(nmax+1,float)
    for n in range(2,nmax+1):
        Q0poly[n]=math.log(n)*(3-1.5*d2[n]+(1/3)*d3[n])

    rough=[]
    rough_err=0.0
    prime_rows=[]
    for n in range(2,nmax+1):
        has_small=False
        for d in range(2,U+1):
            if n%d==0:
                has_small=True
                break
        if not has_small:
            rough.append(n)
            rough_err=max(rough_err,abs(Q0[n]-lam[n]))
            if primes[n]:
                prime_rows.append((n,Q0[n],lam[n]))

    return {
        "N":N,"X":X,"U":U,
        "HB_error":float(np.max(np.abs(HB[1:]-lam[1:]))),
        "sector_error":float(np.max(np.abs((Q0+Q1+Q2+Q3)[1:]-lam[1:]))),
        "Q0poly_error":float(np.max(np.abs(Q0[2:]-Q0poly[2:]))),
        "rough_error":rough_err,
        "rough_count":len(rough),
        "rough_primes":prime_rows[:10]
    }

rows=[run(N) for N in [50,100,180]]
with open("campaign08_heath_brown_k3_crosscheck.csv","w",encoding="utf-8",newline="") as f:
    w=csv.writer(f)
    w.writerow(["N","X","U","HB_identity_error","sector_sum_error","Q0_polynomial_error","rough_Q0_equals_Lambda_error","rough_count"])
    for r in rows:
        w.writerow([r["N"],r["X"],r["U"],r["HB_error"],r["sector_error"],r["Q0poly_error"],r["rough_error"],r["rough_count"]])

print("PASS")
print("max_HB_error",max(r["HB_error"] for r in rows))
print("max_sector_error",max(r["sector_error"] for r in rows))
print("max_Q0_poly_error",max(r["Q0poly_error"] for r in rows))
print("max_rough_error",max(r["rough_error"] for r in rows))