import numpy as np
import cosmotool as ct
from tqdm import tqdm

Nmc=1
L=1.0
Ng=256
Np=128**3
Nbins=50
alpha = (L/Ng)**3
mass = (1.0*Ng**3) / Np

ik = np.fft.fftfreq(Ng, d=L/Ng)*2*np.pi 
kx,ky,kz=np.meshgrid(ik,ik,ik[:(Ng//2+1)],indexing='ij')
k = np.sqrt(kx**2+ky**2+kz**2)
var_in_k =Np * mass**2 * alpha**2 / L**3
Hw,_ = np.histogram(k, range=(0,k.max()), bins=Nbins)

W= lambda q: (np.sinc(q/(2.*np.pi)*L/Ng)**2)

Href,_ = np.histogram(k, weights=(W(kx)*W(ky)*W(kz))**2, range=(0,k.max()), bins=Nbins)
Href = (Href/Hw)[1:]

rho = np.zeros(Nbins-1)
rho_stag = np.zeros(Nbins-1)

for i in tqdm( range(Nmc)):
  x = np.random.uniform(size=Np*3).reshape((Np,3))
  d = ct.project_cic(x, None, Ng, L, True, False)
  d_stag = 0.5*(d+ ct.project_cic(x+L/(2*Ng), None, Ng, L, True, False))
  d *= mass
  d_stag *= mass
  d -= 1
  d_stag -= 1
  
  dhat = np.fft.rfftn(d) * alpha
  
  H,b = np.histogram(k, weights=(dhat.real**2+dhat.imag**2)/L**3, range=(0,k.max()), bins=Nbins)
  dhat = np.fft.rfftn(d_stag) * alpha
  Hstag,b = np.histogram(k, weights=(dhat.real**2+dhat.imag**2)/L**3, range=(0,k.max()), bins=Nbins)

  H = H / Hw
  Hstag = Hstag / Hw
  bc = 0.5*(b[1:]+b[:-1])
  bc = bc[1:]

  rho += H[1:]
  rho_stag += Hstag[1:]

rho /= Nmc
rho_stag /= Nmc

