"""Davenport-Heilbronn analogue of von Mangoldt: -f'/f(s) = sum Lambda_f(n) n^{-s}; compute Lambda_f(n), n<=X, then
P^f_x(s) = sum_{n<=x} Lambda_f(n) n^{-s} on a log grid."""
import numpy as np, sys, time, math
X=int(float(sys.argv[1])); T0=time.time()
KAP=(np.sqrt(10-2*np.sqrt(5))-2)/(np.sqrt(5)-1)
c=np.array([0.0,1.0,KAP,-KAP,-1.0]); n=np.arange(X+1); a=c[n%5]; a[0]=0
L=np.zeros(X+1); acc=np.zeros(X+1); logn=np.log(np.maximum(n,1))
for d in range(2,X+1):
    v=a[d]*logn[d]-acc[d]; L[d]=v
    if v!=0.0 and 2*d<=X:
        k=np.arange(2,X//d+1); acc[d*k]+=v*a[k]
print('Lambda_f computed %.0fs'%(time.time()-T0),flush=True)
np.save('LambdaDH_%d.npy'%X,L)
S=np.array([0.5+10j,0.5+50j,0.5+200j,0.75+50j])
XG=np.unique(np.floor(np.logspace(2,math.log10(X),600)))+0.5
P=np.zeros((len(XG),len(S)),complex); acc=np.zeros(len(S),complex); CH=2000000; gi=0
for lo in range(1,X+1,CH):
    hi=min(lo+CH,X+1); terms=L[lo:hi,None]*np.exp(-np.outer(logn[lo:hi],S)); cs=np.cumsum(terms,axis=0)
    while gi<len(XG) and XG[gi]<hi:
        k=int(XG[gi])-lo; P[gi]=acc+cs[k]; gi+=1
    acc=acc+cs[-1]
np.savez('PdataDH_%d.npz'%X,S=S,XG=XG,P=P)
print('max |Lambda_f(n)|/log n for n<=X:',np.max(np.abs(L[2:])/logn[2:]),' done %.0fs'%(time.time()-T0))