import numpy as np

# CIC interpolation
def cic(x, Lbox, Lboxcut = 0, Nmesh = 128, weights = None):

 if weights == None: weights = 1
 wm = np.mean(weights)
 ws = np.mean(weights**2)

 d = np.mod(x/(Lbox+2*Lboxcut)*Nmesh,1)

 box = ([Lboxcut,Lbox+Lboxcut],[Lboxcut,Lbox+Lboxcut],[Lboxcut,Lbox+Lboxcut])

 rho = np.histogramdd(x, range = box, bins = Nmesh, weights = weights*(1-d[:,0])*(1-d[:,1])*(1-d[:,2]))[0] \
     + np.roll(np.histogramdd(x, range = box, bins = Nmesh, weights = weights*d[:,0]*(1-d[:,1])*(1-d[:,2]))[0],1,0) \
     + np.roll(np.histogramdd(x, range = box, bins = Nmesh, weights = weights*(1-d[:,0])*d[:,1]*(1-d[:,2]))[0],1,1) \
     + np.roll(np.histogramdd(x, range = box, bins = Nmesh, weights = weights*(1-d[:,0])*(1-d[:,1])*d[:,2])[0],1,2) \
     + np.roll(np.roll(np.histogramdd(x, range = box, bins = Nmesh, weights = weights*d[:,0]*d[:,1]*(1-d[:,2]))[0],1,0),1,1) \
     + np.roll(np.roll(np.histogramdd(x, range = box, bins = Nmesh, weights = weights*d[:,0]*(1-d[:,1])*d[:,2])[0],1,0),1,2) \
     + np.roll(np.roll(np.histogramdd(x, range = box, bins = Nmesh, weights = weights*(1-d[:,0])*d[:,1]*d[:,2])[0],1,1),1,2) \
     + np.roll(np.roll(np.roll(np.histogramdd(x, range = box, bins = Nmesh, weights = weights*d[:,0]*d[:,1]*d[:,2])[0],1,0),1,1),1,2)

 rho /= wm
 
 rho = rho/rho.mean() - 1.
 
 return (rho, wm, ws)


# Power spectra & correlation functions
def powcor(d1, d2, Lbox, Nbin = 10, scale = 'lin', cor = False, cic = True, dim = 1):

 Nmesh = len(d1)

 # CIC correction
 if cic:
  wid = np.indices(np.shape(d1))
  wid[np.where(wid >= Nmesh/2)] -= Nmesh
  wid = wid*np.pi/Nmesh + 1e-100
  wcic = np.prod(np.sin(wid)/wid,0)**2

 # Shell average power spectrum
 dk = 2*np.pi/Lbox
 Pk = np.conj(d1)*d2*(Lbox/Nmesh**2)**3
 if cic: Pk /= wcic**2
 
 (Nm, km, Pkm, SPkm) = shellavg(np.real(Pk), dk, Nmesh, Nbin = Nbin, xmin = 0., xmax = Nmesh*dk/2, scale = scale, dim = dim)
 
 # Inverse Fourier transform and shell average correlation function
 if cor:
  if cic: Pk *= wcic**2 # Undo cic-correction in correlation function
  dx = Lbox/Nmesh
  Xr = np.fft.irfftn(Pk)*(Nmesh/Lbox)**3

  (Nmx, rm, Xrm, SXrm) = shellavg(np.real(Xr), dx, Nmesh, Nbin = Nbin/2, xmin = dx, xmax = 140., scale = scale, dim = dim)
  
  return ((Nm, km, Pkm, SPkm),(Nmx, rm, Xrm, SXrm))
 
 else: return (Nm, km, Pkm, SPkm)


# Shell averaging
def shellavg(f, dx, Nmesh, Nbin = 10, xmin = 0., xmax = 1., scale = 'lin', dim = 1):
 
 x = np.indices(np.shape(f))
 x[np.where(x >= Nmesh/2)] -= Nmesh
 f = f.flatten()
 
 if scale == 'lin': bins = xmin+(xmax-xmin)* np.linspace(0,1,Nbin+1)
 if scale == 'log': bins = xmin*(xmax/xmin)**np.linspace(0,1,Nbin+1)

 if dim == 1: # 1D
  x = dx*np.sqrt(np.sum(x**2,0)).flatten()
  Nm = np.histogram(x, bins = bins)[0]
  xm = np.histogram(x, bins = bins, weights = x)[0]/Nm
  fm = np.histogram(x, bins = bins, weights = f)[0]/Nm
  fs = np.sqrt((np.histogram(x, bins = bins, weights = f**2)[0]/Nm - fm**2)/(Nm-1))
  return (Nm, xm, fm, fs)

 elif dim == 2: # 2D
  xper = dx*np.sqrt(x[0,:,:,:]**2 + x[1,:,:,:]**2 + 1e-100).flatten()
  xpar = dx*np.abs(x[2,:,:,:]).flatten()
  x = dx*np.sqrt(np.sum(x**2,0)).flatten()
  Nm = np.histogram2d(xper, xpar, bins = [bins,bins])[0]
  xmper = np.histogram2d(xper, xpar, bins = [bins,bins], weights = xper)[0]/Nm
  xmpar = np.histogram2d(xper, xpar, bins = [bins,bins], weights = xpar)[0]/Nm
  fm = np.histogram2d(xper, xpar, bins = [bins,bins], weights = f)[0]/Nm
  return (Nm, xmper, xmpar, fm)
