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 extern from "cic.hpp" namespace "CosmoTool":

  ctypedef float CICType
  ctypedef float Coordinates[3]
  
  cdef cppclass CICParticles:
    float mass
    Coordinates coords

  cdef cppclass CICFilter:
  
    CICFilter(np.uint32_t resolution, double L) nogil
    
    void resetMesh() nogil
    void putParticles(CICParticles* particles, np.uint32_t N) nogil
    void getDensityField(CICType *& field, np.uint32_t& res) nogil

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

  cdef cppclass SimuData:
    np.float_t BoxSize
    np.float_t time
    np.float_t Hubble

    np.float_t Omega_M
    np.float_t Omega_Lambda
    np.int64_t TotalNumPart
    np.int64_t NumPart
    np.int64_t *Id
    float *Pos[3]
    float *Vel[3]
    int *type
    float *Mass

    bool noAuto

  cdef const int NEED_GADGET_ID 
  cdef const int NEED_POSITION
  cdef const int NEED_VELOCITY
  cdef const int NEED_TYPE
  cdef const int NEED_MASS

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

  SimuData *loadGadgetMulti(const char *fname, int id, int flags) except +
  void cxx_writeGadget "CosmoTool::writeGadget" (const char * s, SimuData *data) except +


cdef extern from "loadRamses.hpp" namespace "CosmoTool":
  SimuData *loadRamsesSimu(const char *basename, int id, int cpuid, bool dp, int flags) except +


class PySimulationBase(object):
  def getPositions(self):
    raise NotImplementedError("getPositions is not implemented")
    
  def getVelocities(self):
    raise NotImplementedError("getVelocities is not implemented")
    
  def getIdentifiers(self):
    raise NotImplementedError("getIdentifiers is not implemented")

  def getTypes(self):
    raise NotImplementedError("getTypes is not implemented")

  def getOmega_M(self):
    raise NotImplementedError("getOmega_M is not implemented")
  
  def getOmega_Lambda(self):
    raise NotImplementedError("getOmega_Lambda is not implemented")

  def getTime(self):
    raise NotImplementedError("getTime is not implemented")

  def getHubble(self):
    raise NotImplementedError("getHubble is not implemented")

  def getBoxsize(self):
    raise NotImplementedError("getBoxsize is not implemented")

  def getMasses(self):
    raise NotImplementedError("getMasses is not implemented")

cdef class Simulation:

  cdef list positions
  cdef list velocities
  cdef object identifiers
  cdef object types
  cdef object masses

  cdef SimuData *data

  property BoxSize:
    def __get__(Simulation self):
        return self.data.BoxSize
        
  property time:
    def __get__(Simulation self):
        return self.data.time

  property Hubble:
    def __get__(Simulation self):
        return self.data.Hubble

  property Omega_M:
    def __get__(Simulation self):
        return self.data.Omega_M
        
  property Omega_Lambda:
    def __get__(Simulation self):
        return self.data.Omega_Lambda
        
  property positions:
    def __get__(Simulation self):
        return self.positions
        
  property velocities:
    def __get__(Simulation self):
        return self.velocities

  property identifiers:
    def __get__(Simulation self):
        return self.identifiers

  property types:
    def __get__(Simulation self):
        return self.types

  property masses:
    def __get__(Simulation self):
        return self.masses
        
  property numParticles:
    def __get__(Simulation self):
        return self.data.NumPart

  def __cinit__(Simulation self):
    self.data = <SimuData *>0
  
  def __dealloc__(Simulation self):
    if self.data != <SimuData *>0:
      print("Clearing simulation data")
      del self.data


class _PySimulationAdaptor(PySimulationBase):

  def __init__(self,sim):
    self.simu = sim

  def getBoxsize(self):
    return self.simu.BoxSize
    
  def getPositions(self):
    return self.simu.positions

  def getTypes(self):
    return self.simu.types
    
  def getVelocities(self):
    return self.simu.velocities
    
  def getIdentifiers(self):
    return self.simu.identifiers

  def getTime(self):
    return self.simu.time
    
  def getHubble(self):
    return self.simu.Hubble
    
  def getOmega_M(self):
    return self.simu.Omega_M
    
  def getOmega_Lambda(self):
    return self.simu.Omega_Lambda

  def getMasses(self):
    return self.simu.masses

cdef class ArrayWrapper:
    cdef void* data_ptr
    cdef int size
    cdef int type_array
    
    cdef set_data(self, int size, int type_array, void* data_ptr):
        """ Set the data of the array

This cannot be done in the constructor as it must recieve C-level
arguments.

Parameters:
-----------
size: int
Length of the array.
data_ptr: void*
Pointer to the data

    """
        self.data_ptr = data_ptr
        self.size = size
        self.type_array = type_array
     
    def __array__(self):
        """ Here we use the __array__ method, that is called when numpy
tries to get an array from the object."""
        cdef np.npy_intp shape[1]
        
        shape[0] = <np.npy_intp> self.size
        # Create a 1D array, of length 'size'
        ndarray = np.PyArray_SimpleNewFromData(1, shape, self.type_array, self.data_ptr)
        return ndarray
     
    def __dealloc__(self):
        """ Frees the array. This is called by Python when all the
references to the object are gone. """
        pass
 
cdef object wrap_array(void *p, np.uint64_t s, int typ):
    cdef np.ndarray ndarray
    cdef ArrayWrapper wrapper

    wrapper = ArrayWrapper()
    wrapper.set_data(s, typ, p)
    ndarray = np.array(wrapper, copy=False)
    ndarray.base = <PyObject*> wrapper
    Py_INCREF(wrapper)
    
    return ndarray


cdef object wrap_float_array(float *p, np.uint64_t s):
    return wrap_array(<void *>p, s, np.NPY_FLOAT32)

cdef object wrap_int64_array(np.int64_t* p, np.uint64_t s):
    return wrap_array(<void *>p, s, np.NPY_INT64)

cdef object wrap_int_array(int* p, np.uint64_t s):
    return wrap_array(<void *>p, s, np.NPY_INT)

cdef object wrap_simudata(SimuData *data, int flags):
  cdef Simulation simu

  simu = Simulation()
  simu.data = data
  if flags & NEED_POSITION:
     simu.positions = [wrap_float_array(data.Pos[i], data.NumPart) for i in xrange(3)]
  else:
     simu.positions = None

  if flags & NEED_VELOCITY:
     simu.velocities = [wrap_float_array(data.Vel[i], data.NumPart) for i in xrange(3)]
  else:
     simu.velocities = None

  if flags & NEED_GADGET_ID:
     simu.identifiers = wrap_int64_array(data.Id, data.NumPart)
  else:
     simu.identifiers = None

  if flags & NEED_TYPE:
     simu.types = wrap_int_array(data.type, data.NumPart)
  else:
     simu.types = None

  if flags & NEED_MASS:
     simu.masses = wrap_float_array(data.Mass, data.NumPart)
  else:
     simu.masses = None

  return simu

def loadGadget(str filename, int snapshot_id, bool loadPosition = True, bool loadVelocity = False, bool loadId = False, bool loadType = False, bool loadMass=False):
  """loadGadget(filename, cpu_id, loadPosition=True, loadVelocity=False, loadId=False, loadType=False)
It loads a gadget-1 snapshot and return a PySimulationBase object. If cpu_id is negative then the snapshot is considered not to be part of a set of snapshots written by different cpu. Otherwise the filename is modified to reflectt this cpu_id.
"""

  cdef int flags
  cdef SimuData *data
  cdef Simulation simu

  flags = 0
  if loadPosition:
     flags |= NEED_POSITION
  if loadVelocity:
     flags |= NEED_VELOCITY
  if loadId:
     flags |= NEED_GADGET_ID
  if loadType:
     flags |= NEED_TYPE
  if loadMass:
     flags |= NEED_MASS

  data = loadGadgetMulti(filename, snapshot_id, flags)
  if data == <SimuData*>0:
    return None

  return _PySimulationAdaptor(wrap_simudata(data, flags))

def writeGadget(str filename, object simulation):

  cdef SimuData simdata
  cdef np.ndarray[np.float32_t, ndim=1] pos, vel
  cdef np.ndarray[np.int64_t, ndim=1] ids
  cdef np.int64_t NumPart
  cdef int j
  
  if not isinstance(simulation,PySimulationBase):
    raise TypeError("Second argument must be of type SimulationBase")
  
  NumPart = simulation.positions[0].size
  simdata.noAuto = True
  
  for j in xrange(3):
    pos = simulation.getPositions()[j]
    vel = simulation.getVelocities()[j]
    
    if pos.size != NumPart or vel.size != NumPart:
      raise ValueError("Invalid number of particles")
    
    simdata.Pos[j] = <float *>pos.data
    simdata.Vel[j] = <float *>vel.data
  
  ids = simulation.getIdentifiers()
  simdata.Id = <np.int64_t *>ids.data
  simdata.BoxSize = simulation.getBoxsize()
  simdata.time = simulation.getTime()
  simdata.Hubble = simulation.getHubble()
  simdata.Omega_M = simulation.getOmega_M()
  simdata.Omega_Lambda = simulation.getOmega_Lambda()
  simdata.TotalNumPart = NumPart
  simdata.NumPart = NumPart

  cxx_writeGadget(filename, &simdata)

def loadRamses(str basepath, int snapshot_id, int cpu_id, bool doublePrecision = False, bool loadPosition = True, bool loadVelocity = False):
  """ loadRamses(basepath, snapshot_id, cpu_id, doublePrecision=False, loadPosition=True, loadVelocity=False)
  Loads the indicated snapshot based on the cpu id, snapshot id and basepath. It is important to specify the correct precision in doublePrecision.
"""
  cdef int flags
  cdef SimuData *data
  cdef Simulation simu

  flags = 0
  if loadPosition:
     flags |= NEED_POSITION
  if loadVelocity:
     flags |= NEED_VELOCITY

  data = loadRamsesSimu(basepath, snapshot_id, cpu_id, doublePrecision, flags)
  if data == <SimuData*>0:
    return None

  return _PySimulationAdaptor(wrap_simudata(data, flags))
  
 
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)
    
  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)
      
@cython.boundscheck(False)
def leanCic(float[:,:] particles, float L, int Resolution):
  cdef CICParticles p
  cdef CICFilter *cic
  cdef np.uint64_t i
  cdef CICType *field
  cdef np.uint32_t dummyRes
  cdef np.ndarray[np.float64_t, ndim=3] out_field
  cdef np.uint64_t j

  cic = new CICFilter(Resolution, L)
  cic.resetMesh()
  
  if particles.shape[1] != 3:
    raise ValueError("Particles must be Nx3 array")
  
  p.mass = 1
  for i in xrange(particles.shape[0]):
    p.coords[0] = particles[i,0]
    p.coords[1] = particles[i,1]
    p.coords[2] = particles[i,2]
    cic.putParticles(&p, 1)
    
  field = <CICType*>0
  dummyRes = 0
  cic.getDensityField(field, dummyRes)
  
  out_field = np.empty((dummyRes, dummyRes, dummyRes), dtype=np.float64)
  for j in xrange(out_field.size):
    out_field[j] = field[j]
  
  del cic
  return out_field