import multiprocessing
import pyfftw
import weakref
import numpy as np
import cosmolopy as cpy

class CubeFT(object):
  def __init__(self, L, N, max_cpu=-1):
  
      self.N = N
      self.align = pyfftw.simd_alignment
      self.L = L
      self.max_cpu = multiprocessing.cpu_count() if max_cpu < 0 else max_cpu
      self._dhat = pyfftw.n_byte_align_empty((self.N,self.N,self.N/2+1), self.align, dtype='complex64')
      self._density = pyfftw.n_byte_align_empty((self.N,self.N,self.N), self.align, dtype='float32')
      self.irfft = pyfftw.FFTW(self._dhat, self._density, axes=(0,1,2), direction='FFTW_BACKWARD', threads=self.max_cpu, normalize_idft=False)
      self.rfft = pyfftw.FFTW(self._density, self._dhat, axes=(0,1,2), threads=self.max_cpu, normalize_idft=False)

  def rfft(self):
      return self.rfft()*(self.L/self.N)**3
      
  def irfft(self):
      return self.irfft()/self.L**3
   
  def get_dhat(self):
      return self._dhat
  def set_dhat(self, in_dhat):
      self._dhat[:] = in_dhat 
  dhat = property(get_dhat, set_dhat, None)
      
  def get_density(self):
      return self._density      
  def set_density(self, d):
      self._density[:] = d 
  density = property(get_density, set_density, None)
      

class CosmoGrowth(object):

  def __init__(self, **cosmo):
    self.cosmo = cosmo

  def D(self, a):
    return cpy.perturbation.fgrowth(1/a-1, self.cosmo['omega_M_0'], unnormed=True)

  def compute_E(self, a):
    om = self.cosmo['omega_M_0']
    ol = self.cosmo['omega_lambda_0']
    ok = self.cosmo['omega_k_0']

    E = np.sqrt(om/a**3 + ol + ok/a**2)
  
    H2 = -3*om/a**4  - 2*ok/a**3

    Eprime = 0.5*H2/E

    return E,Eprime

  def Ddot(self, a):
    E,Eprime = self.compute_E(a)
    D = self.D(a)
    Ddot_D = Eprime/E + 2.5 * self.cosmo['omega_M_0']/(a**3*E**2*D)
    Ddot_D *= a
    return Ddot_D

  def compute_velmul(self, a):
    E,_ = self.compute_E(a)
    velmul = self.Ddot(a)
    velmul *= 100 * a * E
    return velmul





class LagrangianPerturbation(object):

    def __init__(self,density,L, fourier=False, supersample=1, max_cpu=-1):
    
        self.L = L
        self.N = density.shape[0]
        
        self.max_cpu = max_cpu
        self.cube = CubeFT(self.L, self.N, max_cpu=max_cpu)
        
        if not fourier:
          self.cube.density = density
          self.dhat = self.cube.rfft().copy()
        else:
          self.dhat = density.copy()

        if supersample > 1:
          self.upgrade_sampling(supersample)
        self.ik = np.fft.fftfreq(self.N, d=L/self.N)*2*np.pi
        self.cache = {}#weakref.WeakValueDictionary()

    def upgrade_sampling(self, supersample):
        N2 = self.N * supersample
        N = self.N
        dhat_new = np.zeros((N2, N2, N2/2+1), dtype=np.complex128)

        hN = N/2
        dhat_new[:hN, :hN, :hN+1] = self.dhat[:hN, :hN, :]
        dhat_new[:hN, (N2-hN):N2, :hN+1] = self.dhat[:hN, hN:, :]
        dhat_new[(N2-hN):N2, (N2-hN):N2, :hN+1] = self.dhat[hN:, hN:, :]
        dhat_new[(N2-hN):N2, :hN, :hN+1] = self.dhat[hN:, :hN, :]

        self.dhat = dhat_new
        self.N = N2
        self.cube = CubeFT(self.L, self.N, max_cpu=self.max_cpu)
        
    def _gradient(self, phi, direction):
        self.cube.dhat = self._kdir(direction)*1j*phi 
        return self.cube.irfft()
            
    def lpt1(self, direction=0):
        k2 = self._get_k2()
        k2[0,0,0] = 1
    
        return self._gradient(self.dhat/k2, direction)
        
    def new_shape(self,direction, q=3, half=False):
        N0 = (self.N/2+1) if half else self.N
        return ((1,)*direction) + (N0,) + ((1,)*(q-1-direction))

    def _kdir(self, direction, q=3):
        if direction != q-1:
            return self.ik.reshape(self.new_shape(direction, q=q))
        else:
            return self.ik[:self.N/2+1].reshape(self.new_shape(direction, q=q, half=True))

    def _get_k2(self, q=3):
        if 'k2' in self.cache:
            return self.cache['k2']
            
        k2 = self._kdir(0, q=q)**2
        for d in xrange(1,q):
            k2 = k2 + self._kdir(d, q=q)**2
            
        self.cache['k2'] = k2
        return k2
        
    def _do_irfft(self, array):
        self.cube.dhat = array
        return self.cube.irfft()
        
    def _do_rfft(self, array):
        self.cube.density = array
        return self.cube.rfft()

    def lpt2(self, direction=0):
        k2 = self._get_k2()
        k2[0,0,0] = 1

        if 'lpt2_potential' not in self.cache:
            print("Rebuilding potential...")
            div_phi2 = np.zeros((self.N,self.N,self.N), dtype=np.float64)
            for j in xrange(3):
                q = self._do_irfft( self._kdir(j)**2*self.dhat / k2 ).copy()
                for i in xrange(j+1, 3):
                    div_phi2 += q * self._do_irfft( self._kdir(i)**2*self.dhat / k2 )
                    div_phi2 -= (self._do_irfft(self._kdir(j)*self._kdir(i)*self.dhat / k2 ) )**2

            div_phi2 *= 1/self.L**6
            phi2_hat = -self._do_rfft(div_phi2) / k2
            #self.cache['lpt2_potential'] = phi2_hat
            del div_phi2
        else:
            phi2_hat = self.cache['lpt2_potential']
            
        return self._gradient(phi2_hat, direction)