from libcpp cimport bool
from libcpp cimport string as cppstring
import numpy as np
cimport numpy as np
from cpython cimport PyObject, Py_INCREF
cimport cython

np.import_array()

cdef extern from "cosmopower.hpp" namespace "CosmoTool":

  cdef enum CosmoFunction "CosmoTool::CosmoPower::CosmoFunction":
      POWER_EFSTATHIOU "CosmoTool::CosmoPower::POWER_EFSTATHIOU",
      HU_WIGGLES "CosmoTool::CosmoPower::HU_WIGGLES",
      HU_BARYON "CosmoTool::CosmoPower::HU_BARYON",
      OLD_POWERSPECTRUM,
      POWER_BARDEEN "CosmoTool::CosmoPower::POWER_BARDEEN",
      POWER_SUGIYAMA "CosmoTool::CosmoPower::POWER_SUGIYAMA",
      POWER_BDM,
      POWER_TEST

  cdef cppclass CosmoPower:
    double n
    double K0
    double V_LG_CMB
    
    double CMB_VECTOR[3]
    double h
    double SIGMA8
    double OMEGA_B
    double OMEGA_C
    double omega_B
    double omega_C
    double Theta_27
    double OMEGA_0
    double Omega
    double beta
    double OmegaEff
    double Gamma0
    double normPower
    
        
    CosmoPower()
    void setFunction(CosmoFunction)
    void updateCosmology()
    void updatePhysicalCosmology()
    void normalize()
    void setNormalization(double)
    double power(double)
 
cdef class CosmologyPower:

  cdef CosmoPower power

  def __init__(self,**cosmo):
  
    self.power = CosmoPower()
    self.power.OMEGA_B = cosmo['omega_B_0']
    self.power.OMEGA_C = cosmo['omega_M_0']-cosmo['omega_B_0']
    self.power.h = cosmo['h']
    if 'ns' in cosmo:
      self.power.n = cosmo['ns'] 
    
    assert self.power.OMEGA_C > 0
    
    self.power.updateCosmology()
    
  def normalize(self,s8):
    self.power.SIGMA8 = s8
    self.power.normalize()
    
    
  def setFunction(self,funcname):
    cdef CosmoFunction f
    
    f = POWER_EFSTATHIOU
  
    if funcname=='EFSTATHIOU':
      f = POWER_EFSTATHIOU
    elif funcname=='HU_WIGGLES':
      f = HU_WIGGLES
    elif funcname=='HU_BARYON':
      f = HU_BARYON
    elif funcname=='BARDEEN':
      f = POWER_BARDEEN
    elif funcname=='SUGIYAMA':
      f = POWER_SUGIYAMA

    self.power.setFunction(f)
    
  cdef double _compute(self, double k):
    k *= self.power.h
    return self.power.power(k) * self.power.h**3
    
  def compute(self, k):
    cdef np.ndarray out
    cdef double kval
    cdef tuple i
    
    if isinstance(k, np.ndarray):
      out = np.empty(k.shape, dtype=np.float64)
      for i,kval in np.ndenumerate(k):
        out[i] = self._compute(kval)
      return out
    else:
      return self._compute(k)