from libcpp cimport bool
from libcpp cimport string as cppstring
from libcpp.vector cimport vector as cppvector
from cython.parallel cimport prange
import numpy as np
cimport numpy as np
from cpython cimport PyObject, Py_INCREF
cimport cython

np.import_array()

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) nogil except +
  void cxx_writeGadget "CosmoTool::writeGadget" (const char * s, SimuData *data) except +

cdef extern from "safe_gadget.hpp":
  SimuData *loadGadgetMulti_safe(cppstring.string s, int flags) nogil
  SimuData **alloc_simudata(int num) nogil
  void del_simudata(SimuData **d) nogil

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


  property totalNumParticles:
    def __get__(Simulation self):
        return self.data.TotalNumPart

  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 np.uint64_t size
    cdef int type_array
    
    cdef set_data(self, np.uint64_t 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 loadParallelGadget(object filename_list, bool loadPosition = True, bool loadVelocity = False, bool loadId = False, bool loadType = False, bool loadMass=False):
  """loadGadget(filename list, loadPosition=True, loadVelocity=False, loadId=False, loadType=False)

  Arguments:
    - filename list: a list or tuple of filenames to load in parallel
  
  Keywords:
    - loadPosition: indicate to load positions
    - loadVelocity: indicate to load velocities
    - loadId: indicate to load id
    - loadType: indicate to load particle types
    
It loads a gadget-1 snapshot and return a PySimulationBase object.
"""
  cdef int flags, i, num_files
  cdef list out_arrays
  cdef SimuData ** data
  cdef SimuData * local_data
  cdef Simulation simu
  cdef cppvector[cppstring.string] filenames

  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

  num_files = len(filename_list)
  filenames.resize(num_files)
  data = alloc_simudata(num_files)
  for i,l in enumerate(filename_list):
    filenames[i] = l
  
  with nogil:
    for i in prange(num_files):
      local_data = loadGadgetMulti_safe(filenames[i], flags)
      data[i] = local_data
#      data[i] = loadGadgetMulti(filenames[i].c_str(), -1, flags)
  
  out_arrays = []
  for i in xrange(num_files):
    if data[i] == <SimuData*>0:        
      out_arrays.append(None)
    else:
      out_arrays.append(_PySimulationAdaptor(wrap_simudata(data[i], flags)))

  del_simudata(data)

  return out_arrays

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))