# distutils: include_dirs = LIB_DIR
# distutils: libraries = STD_LIBS
"""
Particle Deposition onto Cells




"""


cimport numpy as np

import numpy as np

cimport cython
from libc.math cimport sqrt

from yt.utilities.lib.fp_utils cimport *

from .oct_container cimport Oct, OctInfo, OctreeContainer

from yt.utilities.lib.misc_utilities import OnceIndirect


cdef append_axes(np.ndarray arr, int naxes):
    if arr.ndim == naxes:
        return arr
    # Avoid copies
    arr2 = arr.view()
    arr2.shape = arr2.shape + (1,) * (naxes - arr2.ndim)
    return arr2

cdef class ParticleDepositOperation:
    def __init__(self, nvals, kernel_name):
        # nvals is a tuple containing the active dimensions of the
        # grid to deposit onto and the number of grids,
        # (nx, ny, nz, ngrids)
        self.nvals = nvals
        self.update_values = 0 # This is the default
        self.sph_kernel = get_kernel_func(kernel_name)

    def initialize(self, *args):
        raise NotImplementedError

    def finalize(self, *args):
        raise NotImplementedError

    @cython.boundscheck(False)
    @cython.wraparound(False)
    def process_octree(self, OctreeContainer octree,
                     np.ndarray[np.int64_t, ndim=1] dom_ind,
                     np.ndarray[np.float64_t, ndim=2] positions,
                     fields = None, int domain_id = -1,
                     int domain_offset = 0, lvlmax = None):
        cdef int nf, i, j
        if fields is None:
            fields = []
        nf = len(fields)
        cdef np.float64_t[::cython.view.indirect, ::1] field_pointers
        if nf > 0: field_pointers = OnceIndirect(fields)
        cdef np.float64_t pos[3]
        cdef np.float64_t[:] field_vals = np.empty(nf, dtype="float64")
        cdef int dims[3]
        dims[0] = dims[1] = dims[2] = octree.nz
        cdef OctInfo oi
        cdef np.int64_t offset, moff
        cdef Oct *oct
        cdef np.int8_t use_lvlmax
        moff = octree.get_domain_offset(domain_id + domain_offset)
        if lvlmax is None:
            use_lvlmax = False
            lvlmax = []
        else:
            use_lvlmax = True
        cdef np.ndarray[np.int32_t, ndim=1] lvlmaxval = np.asarray(lvlmax, dtype=np.int32)

        for i in range(positions.shape[0]):
            # We should check if particle remains inside the Oct here
            for j in range(nf):
                field_vals[j] = field_pointers[j,i]
            for j in range(3):
                pos[j] = positions[i, j]
            # This line should be modified to have it return the index into an
            # array based on whatever cutting of the domain we have done.  This
            # may or may not include the domain indices that we have
            # previously generated.  This way we can support not knowing the
            # full octree structure.  All we *really* care about is some
            # arbitrary offset into a field value for deposition.
            if not use_lvlmax:
                oct = octree.get(pos, &oi)
            else:
                oct = octree.get(pos, &oi, max_level=lvlmaxval[i])
            # This next line is unfortunate.  Basically it says, sometimes we
            # might have particles that belong to octs outside our domain.
            # For the distributed-memory octrees, this will manifest as a NULL
            # oct.  For the non-distributed memory octrees, we'll simply see
            # this as a domain_id that is not the current domain id.  Note that
            # this relies on the idea that all the particles in a region are
            # all fed to sequential domain subsets, which will not be true with
            # RAMSES, where we *will* miss particles that live in ghost
            # regions on other processors.  Addressing this is on the TODO
            # list.
            if oct == NULL or (domain_id > 0 and oct.domain != domain_id):
                continue
            # Note that this has to be our local index, not our in-file index.
            offset = dom_ind[oct.domain_ind - moff]
            if offset < 0: continue
            # Check that we found the oct ...
            self.process(dims, i, oi.left_edge, oi.dds,
                         offset, pos, field_vals, oct.domain_ind)
            if self.update_values == 1:
                for j in range(nf):
                    field_pointers[j][i] = field_vals[j]

    @cython.boundscheck(False)
    @cython.wraparound(False)
    def process_grid(self, gobj,
                     np.ndarray[np.float64_t, ndim=2, cast=True] positions,
                     fields = None):
        cdef int nf, i, j
        if fields is None:
            fields = []
        if positions.shape[0] == 0: return
        nf = len(fields)
        cdef np.float64_t[:] field_vals = np.empty(nf, dtype="float64")
        cdef np.float64_t[::cython.view.indirect, ::1] field_pointers
        if nf > 0: field_pointers = OnceIndirect(fields)
        cdef np.float64_t pos[3]
        cdef np.int64_t gid = getattr(gobj, "id", -1)
        cdef np.float64_t dds[3]
        cdef np.float64_t left_edge[3]
        cdef np.float64_t right_edge[3]
        cdef int dims[3]
        for i in range(3):
            dds[i] = gobj.dds[i]
            left_edge[i] = gobj.LeftEdge[i]
            right_edge[i] = gobj.RightEdge[i]
            dims[i] = gobj.ActiveDimensions[i]
        for i in range(positions.shape[0]):
            # Now we process
            for j in range(nf):
                field_vals[j] = field_pointers[j,i]
            for j in range(3):
                pos[j] = positions[i, j]
            continue_loop = False
            for j in range(3):
                if pos[j] < left_edge[j] or pos[j] > right_edge[j]:
                    continue_loop = True
            if continue_loop:
                continue
            self.process(dims, i, left_edge, dds, 0, pos, field_vals, gid)
            if self.update_values == 1:
                for j in range(nf):
                    field_pointers[j][i] = field_vals[j]

    cdef int process(self, int dim[3], int ipart, np.float64_t left_edge[3],
                     np.float64_t dds[3], np.int64_t offset,
                     np.float64_t ppos[3], np.float64_t[:] fields,
                     np.int64_t domain_ind) except -1 nogil:
        with gil:
            raise NotImplementedError

cdef class CountParticles(ParticleDepositOperation):
    cdef np.int64_t[:,:,:,:] count
    def initialize(self):
        # Create a numpy array accessible to python
        self.count = append_axes(
            np.zeros(self.nvals, dtype="int64", order='F'), 4)

    @cython.cdivision(True)
    @cython.boundscheck(False)
    cdef int process(self, int dim[3], int ipart,
                     np.float64_t left_edge[3],
                     np.float64_t dds[3],
                     np.int64_t offset, # offset into IO field
                     np.float64_t ppos[3], # this particle's position
                     np.float64_t[:] fields,
                     np.int64_t domain_ind
                     ) except -1 nogil:
        # here we do our thing; this is the kernel
        cdef int ii[3]
        cdef int i
        for i in range(3):
            ii[i] = <int>((ppos[i] - left_edge[i])/dds[i])
        self.count[ii[2], ii[1], ii[0], offset] += 1
        return 0

    def finalize(self):
        arr = np.asarray(self.count)
        arr.shape = self.nvals
        return arr.astype("float64")

deposit_count = CountParticles

cdef class SimpleSmooth(ParticleDepositOperation):
    # Note that this does nothing at the edges.  So it will give a poor
    # estimate there, and since Octrees are mostly edges, this will be a very
    # poor SPH kernel.
    cdef np.float64_t[:,:,:,:] data
    cdef np.float64_t[:,:,:,:] temp

    def initialize(self):
        self.data = append_axes(
            np.zeros(self.nvals, dtype="float64", order='F'), 4)
        self.temp = append_axes(
            np.zeros(self.nvals, dtype="float64", order='F'), 4)

    @cython.cdivision(True)
    @cython.boundscheck(False)
    cdef int process(self, int dim[3], int ipart,
                     np.float64_t left_edge[3],
                     np.float64_t dds[3],
                     np.int64_t offset,
                     np.float64_t ppos[3],
                     np.float64_t[:] fields,
                     np.int64_t domain_ind
                     ) except -1 nogil:
        cdef int ii[3]
        cdef int ib0[3]
        cdef int ib1[3]
        cdef int i, j, k, half_len
        cdef np.float64_t idist[3]
        cdef np.float64_t kernel_sum, dist
        # Smoothing length is fields[0]
        kernel_sum = 0.0
        for i in range(3):
            ii[i] = <int>((ppos[i] - left_edge[i])/dds[i])
            half_len = <int>(fields[0]/dds[i]) + 1
            ib0[i] = ii[i] - half_len
            ib1[i] = ii[i] + half_len
            if ib0[i] >= dim[i] or ib1[i] <0:
                return 0
            ib0[i] = iclip(ib0[i], 0, dim[i] - 1)
            ib1[i] = iclip(ib1[i], 0, dim[i] - 1)
        for i from ib0[0] <= i <= ib1[0]:
            idist[0] = (ii[0] - i) * dds[0]
            idist[0] *= idist[0]
            for j from ib0[1] <= j <= ib1[1]:
                idist[1] = (ii[1] - j) * dds[1]
                idist[1] *= idist[1]
                for k from ib0[2] <= k <= ib1[2]:
                    idist[2] = (ii[2] - k) * dds[2]
                    idist[2] *= idist[2]
                    dist = idist[0] + idist[1] + idist[2]
                    # Calculate distance in multiples of the smoothing length
                    dist = sqrt(dist) / fields[0]
                    with gil:
                        self.temp[k,j,i,offset] = self.sph_kernel(dist)
                    kernel_sum += self.temp[k,j,i,offset]
        # Having found the kernel, deposit accordingly into gdata
        for i from ib0[0] <= i <= ib1[0]:
            for j from ib0[1] <= j <= ib1[1]:
                for k from ib0[2] <= k <= ib1[2]:
                    dist = self.temp[k,j,i,offset] / kernel_sum
                    self.data[k,j,i,offset] += fields[1] * dist
        return 0

    def finalize(self):
        return self.odata

deposit_simple_smooth = SimpleSmooth

cdef class SumParticleField(ParticleDepositOperation):
    cdef np.float64_t[:,:,:,:] sum
    def initialize(self):
        self.sum = append_axes(
            np.zeros(self.nvals, dtype="float64", order='F'), 4)

    @cython.cdivision(True)
    @cython.boundscheck(False)
    cdef int process(self, int dim[3], int ipart,
                     np.float64_t left_edge[3],
                     np.float64_t dds[3],
                     np.int64_t offset,
                     np.float64_t ppos[3],
                     np.float64_t[:] fields,
                     np.int64_t domain_ind
                     ) except -1 nogil:
        cdef int ii[3]
        cdef int i
        for i in range(3):
            ii[i] = <int>((ppos[i] - left_edge[i]) / dds[i])
        self.sum[ii[2], ii[1], ii[0], offset] += fields[0]
        return 0

    def finalize(self):
        sum = np.asarray(self.sum)
        sum.shape = self.nvals
        return sum

deposit_sum = SumParticleField

cdef class StdParticleField(ParticleDepositOperation):
    # Thanks to Britton and MJ Turk for the link
    # to a single-pass STD
    # http://www.cs.berkeley.edu/~mhoemmen/cs194/Tutorials/variance.pdf
    cdef np.float64_t[:,:,:,:] mk
    cdef np.float64_t[:,:,:,:] qk
    cdef np.float64_t[:,:,:,:] i
    def initialize(self):
        # we do this in a single pass, but need two scalar
        # per cell, M_k, and Q_k and also the number of particles
        # deposited into each one
        # the M_k term
        self.mk = append_axes(
            np.zeros(self.nvals, dtype="float64", order='F'), 4)
        self.qk = append_axes(
            np.zeros(self.nvals, dtype="float64", order='F'), 4)
        self.i = append_axes(
            np.zeros(self.nvals, dtype="float64", order='F'), 4)

    @cython.cdivision(True)
    @cython.boundscheck(False)
    cdef int process(self, int dim[3], int ipart,
                     np.float64_t left_edge[3],
                     np.float64_t dds[3],
                     np.int64_t offset,
                     np.float64_t ppos[3],
                     np.float64_t[:] fields,
                     np.int64_t domain_ind
                     ) except -1 nogil:
        cdef int ii[3]
        cdef int i
        cdef float k, mk, qk
        for i in range(3):
            ii[i] = <int>((ppos[i] - left_edge[i])/dds[i])
        k = self.i[ii[2], ii[1], ii[0], offset]
        mk = self.mk[ii[2], ii[1], ii[0], offset]
        qk = self.qk[ii[2], ii[1], ii[0], offset]
        if k == 0.0:
            # Initialize cell values
            self.mk[ii[2], ii[1], ii[0], offset] = fields[0]
        else:
            self.mk[ii[2], ii[1], ii[0], offset] = mk + (fields[0] - mk) / k
            self.qk[ii[2], ii[1], ii[0], offset] = \
                qk + (k - 1.0) * (fields[0] - mk) * (fields[0] - mk) / k
        self.i[ii[2], ii[1], ii[0], offset] += 1
        return 0

    def finalize(self):
        # This is the standard variance
        # if we want sample variance divide by (self.oi - 1.0)
        i = np.asarray(self.i)
        std2 = np.asarray(self.qk) / i
        std2[i == 0.0] = 0.0
        std2.shape = self.nvals
        return np.sqrt(std2)

deposit_std = StdParticleField

cdef class CICDeposit(ParticleDepositOperation):
    cdef np.float64_t[:,:,:,:] field
    cdef public object ofield
    def initialize(self):
        if not all(_ > 1 for _ in self.nvals[:-1]):
            from yt.utilities.exceptions import YTBoundsDefinitionError
            raise YTBoundsDefinitionError(
                "CIC requires minimum of 2 zones in all spatial dimensions.",
                self.nvals[:-1])
        self.field = append_axes(
            np.zeros(self.nvals, dtype="float64", order='F'), 4)

    @cython.cdivision(True)
    @cython.boundscheck(False)
    cdef int process(self, int dim[3], int ipart,
                     np.float64_t left_edge[3],
                     np.float64_t dds[3],
                     np.int64_t offset, # offset into IO field
                     np.float64_t ppos[3], # this particle's position
                     np.float64_t[:] fields,
                     np.int64_t domain_ind
                     ) except -1 nogil:

        cdef int i, j, k
        cdef int ind[3]
        cdef np.float64_t rpos[3]
        cdef np.float64_t rdds[3][2]

        # Compute the position of the central cell
        for i in range(3):
            rpos[i] = (ppos[i]-left_edge[i])/dds[i]
            rpos[i] = fclip(rpos[i], 0.5001, dim[i]-0.5001)
            ind[i] = <int> (rpos[i] + 0.5)
            # Note these are 1, then 0
            rdds[i][1] = (<np.float64_t> ind[i]) + 0.5 - rpos[i]
            rdds[i][0] = 1.0 - rdds[i][1]

        for i in range(2):
            for j in range(2):
                for k in range(2):
                    self.field[ind[2] - k, ind[1] - j, ind[0] - i, offset] += \
                        fields[0]*rdds[0][i]*rdds[1][j]*rdds[2][k]

        return 0

    def finalize(self):
        rv = np.asarray(self.field)
        rv.shape = self.nvals
        return rv

deposit_cic = CICDeposit

cdef class WeightedMeanParticleField(ParticleDepositOperation):
    # Deposit both mass * field and mass into two scalars
    # then in finalize divide mass * field / mass
    cdef np.float64_t[:,:,:,:] wf
    cdef np.float64_t[:,:,:,:] w
    def initialize(self):
        self.wf = append_axes(
            np.zeros(self.nvals, dtype='float64', order='F'), 4)
        self.w = append_axes(
            np.zeros(self.nvals, dtype='float64', order='F'), 4)

    @cython.cdivision(True)
    @cython.boundscheck(False)
    cdef int process(self, int dim[3], int ipart,
                     np.float64_t left_edge[3],
                     np.float64_t dds[3],
                     np.int64_t offset,
                     np.float64_t ppos[3],
                     np.float64_t[:] fields,
                     np.int64_t domain_ind
                     ) except -1 nogil:
        cdef int ii[3]
        cdef int i
        for i in range(3):
            ii[i] = <int>((ppos[i] - left_edge[i]) / dds[i])
        self.w[ii[2], ii[1], ii[0], offset] += fields[1]
        self.wf[ii[2], ii[1], ii[0], offset] += fields[0] * fields[1]
        return 0

    def finalize(self):
        wf = np.asarray(self.wf)
        w = np.asarray(self.w)
        with np.errstate(divide='ignore', invalid='ignore'):
            rv = wf / w
        rv.shape = self.nvals
        return rv

deposit_weighted_mean = WeightedMeanParticleField

cdef class MeshIdentifier(ParticleDepositOperation):
    # This is a tricky one!  What it does is put into the particle array the
    # value of the oct or block (grids will always be zero) identifier that a
    # given particle resides in
    def initialize(self):
        self.update_values = 1

    @cython.cdivision(True)
    @cython.boundscheck(False)
    cdef int process(self, int dim[3], int ipart,
                      np.float64_t left_edge[3],
                      np.float64_t dds[3],
                      np.int64_t offset,
                      np.float64_t ppos[3],
                      np.float64_t[:] fields,
                      np.int64_t domain_ind
                      ) except -1 nogil:
        fields[0] = domain_ind
        return 0

    def finalize(self):
        return

deposit_mesh_id = MeshIdentifier

cdef class CellIdentifier(ParticleDepositOperation):
    cdef np.int64_t[:] indexes, cell_index
    # This method stores the offset of the grid containing each particle
    # and compute the index of the cell in the oct.
    def initialize(self, int npart):
        self.indexes = np.zeros(npart, dtype=np.int64, order='F') - 1
        self.cell_index = np.zeros(npart, dtype=np.int64, order='F') - 1

    @cython.cdivision(True)
    @cython.boundscheck(False)
    cdef int process(self, int dim[3], int ipart,
                      np.float64_t left_edge[3],
                      np.float64_t dds[3],
                      np.int64_t offset,
                      np.float64_t ppos[3],
                      np.float64_t[:] fields,
                      np.int64_t domain_ind
                      ) except -1 nogil:
        cdef int i, icell
        self.indexes[ipart] = offset

        icell = 0
        for i in range(3):
            if ppos[i] > left_edge[i] + dds[i]:
                icell |= 4 >> i

        # Compute cell index
        self.cell_index[ipart] = icell

        return 0

    def finalize(self):
        return np.asarray(self.indexes), np.asarray(self.cell_index)

deposit_cell_id = CellIdentifier

cdef class NNParticleField(ParticleDepositOperation):
    cdef np.float64_t[:,:,:,:] nnfield
    cdef np.float64_t[:,:,:,:] distfield
    def initialize(self):
        self.nnfield = append_axes(
            np.zeros(self.nvals, dtype="float64", order='F'), 4)
        self.distfield = append_axes(
            np.zeros(self.nvals, dtype="float64", order='F'), 4)
        self.distfield[:] = np.inf

    @cython.cdivision(True)
    @cython.boundscheck(False)
    cdef int process(self, int dim[3], int ipart,
                     np.float64_t left_edge[3],
                     np.float64_t dds[3],
                     np.int64_t offset,
                     np.float64_t ppos[3],
                     np.float64_t[:] fields,
                     np.int64_t domain_ind
                     ) except -1 nogil:
        # This one is a bit slow.  Every grid cell is going to be iterated
        # over, and we're going to deposit particles in it.
        cdef int i, j, k
        cdef np.float64_t r2
        cdef np.float64_t gpos[3]
        gpos[0] = left_edge[0] + 0.5 * dds[0]
        for i in range(dim[0]):
            gpos[1] = left_edge[1] + 0.5 * dds[1]
            for j in range(dim[1]):
                gpos[2] = left_edge[2] + 0.5 * dds[2]
                for k in range(dim[2]):
                    r2 = ((ppos[0] - gpos[0])*(ppos[0] - gpos[0]) +
                          (ppos[1] - gpos[1])*(ppos[1] - gpos[1]) +
                          (ppos[2] - gpos[2])*(ppos[2] - gpos[2]))
                    if r2 < self.distfield[k,j,i,offset]:
                        self.distfield[k,j,i,offset] = r2
                        self.nnfield[k,j,i,offset] = fields[0]
                    gpos[2] += dds[2]
                gpos[1] += dds[1]
            gpos[0] += dds[0]
        return 0

    def finalize(self):
        nn = np.asarray(self.nnfield)
        nn.shape = self.nvals
        return nn

deposit_nearest = NNParticleField
