# cython: cdivision=True
# cython: boundscheck=False
# cython: wraparound=False

# Authors: Gilles Louppe <g.louppe@gmail.com>
#          Peter Prettenhofer <peter.prettenhofer@gmail.com>
#          Brian Holt <bdholt1@gmail.com>
#          Noel Dawe <noel@dawe.me>
#          Satrajit Gosh <satrajit.ghosh@gmail.com>
#          Lars Buitinck
#          Arnaud Joly <arnaud.v.joly@gmail.com>
#          Joel Nothman <joel.nothman@gmail.com>
#          Fares Hedayati <fares.hedayati@gmail.com>
#          Jacob Schreiber <jmschreiber91@gmail.com>
#          Nelson Liu <nelson@nelsonliu.me>
#
# License: BSD 3 clause

cimport numpy as np
import numpy as np
from cpython cimport Py_INCREF, PyObject
from libc.string cimport memcpy

np.import_array()

from libc.stdlib cimport free
from libc.stdlib cimport realloc

# Repeat struct definition for numpy
NODE_DTYPE = np.dtype({
    'names': ['left_child', 'right_child', 'feature'],
    #'formats': [np.intp, np.intp, np.intp],
    'formats': [np.int16, np.int16, np.uint16],
    'offsets': [
        <Py_ssize_t> &(<Node*> NULL).left_child,
        <Py_ssize_t> &(<Node*> NULL).right_child,
        <Py_ssize_t> &(<Node*> NULL).feature
    ]
})

cdef extern from "numpy/arrayobject.h":
    object PyArray_NewFromDescr(object subtype, np.dtype descr,
                                int nd, np.npy_intp* dims,
                                np.npy_intp* strides,
                                void* data, int flags, object obj)


from numpy import float32 as DTYPE

cdef realloc_ptr safe_realloc_node(realloc_ptr* p, size_t nelems) nogil except *:
    # sizeof(realloc_ptr[0]) would be more like idiomatic C, but causes Cython
    # 0.20.1 to crash.
    cdef size_t nbytes = nelems * sizeof(p[0][0])
    if nbytes / sizeof(p[0][0]) != nelems:
        # Overflow in the multiplication
        with gil:
            raise MemoryError("could not allocate (%d * %d) bytes"
                              % (nelems, sizeof(p[0][0])))
    cdef realloc_ptr tmp = <realloc_ptr>realloc(p[0], nbytes)
    if tmp == NULL:
        with gil:
            raise MemoryError("could not allocate %d bytes" % nbytes)

    p[0] = tmp
    return tmp  # for convenience

# =============================================================================
# Tree
# =============================================================================
cdef class Tree:
    def __cinit__(self):
        self.node_count=0
        self.nodes=NULL


    def __dealloc__(self):
        """Destructor."""
        free(self.nodes)

    property nodes:
        def __get__(self):
            return self._get_node_ndarray()

    def __getstate__(self):
        """Getstate re-implementation, for pickling."""
        return self._get_node_ndarray()

    def __setstate__(self, d):
        """Setstate re-implementation, for unpickling."""
        size=d.size
        self.node_count=size
        safe_realloc_node(&self.nodes, size)
        nodes = memcpy(self.nodes, (<np.ndarray> d).data,size * sizeof(Node))


    def import_tree(self,np.ndarray left_child, np.ndarray right_child,np.ndarray feature):
        capacity=left_child.size
        self.node_count=capacity
        safe_realloc_node(&self.nodes, capacity)

        for i in range(capacity):
            self.nodes[i].left_child=left_child[i]
            self.nodes[i].right_child=right_child[i]
            self.nodes[i].feature=feature[i]
            #print(i,self.nodes_new[i])
        pass

    cpdef np.ndarray apply(self, object X,
                           object tree_nodes,object nodes_len,
                           object values_len):
        # Check input
        if not isinstance(X, np.ndarray):
            raise ValueError("X should be in np.ndarray format, got %s"
                             % type(X))

        if X.dtype != DTYPE:
            raise ValueError("X.dtype should be np.float32, got %s" % X.dtype)

        # Extract input
        cdef DTYPE_t[:, :] X_ndarray = X
        cdef Node[:] nodes_array=tree_nodes
        cdef SIZE_t[:] nodes_len_array=nodes_len
        cdef SIZE_t[:] values_len_array=values_len
        # cdef SIZE_t[:] result_arr=result

        cdef SIZE_t n_samples = X.shape[0]
        cdef SIZE_t n_nodes = nodes_len.shape[0]

        cdef np.ndarray[UINT32_t] result = np.zeros((n_nodes*n_samples,), dtype=np.uint32)
        cdef UINT32_t* result_arr = <UINT32_t*> result.data

        # Initialize auxiliary data-structure
        cdef SIZE_t i = 0
        cdef INT16_t node_next
        cdef SIZE_t node_i
        cdef Node* node

        cdef INT64_t node_start=0
        cdef INT64_t value_start=0

        with nogil:
            # if 1:
            for node_i in range(n_nodes):
                for i in range(n_samples):
                    node = &nodes_array[node_start]
                    if X_ndarray[i, node.feature] <= 0.5:
                        node_next=node.left_child
                    else:
                        node_next=node.right_child
                    # print(i,node.left_child,node.right_child,node.feature)

                    # While node not a leaf
                    while node_next>0:
                        node = &nodes_array[node_next+node_start]
                        # print(i,node.left_child,node.right_child,node.feature)
                        if X_ndarray[i, node.feature] <= 0.5:
                            node_next=node.left_child
                        else:
                            node_next=node.right_child

                    result_arr[node_i*n_samples+i]=<UINT32_t>(value_start-node_next)

                node_start+=nodes_len_array[node_i]
                value_start+=values_len_array[node_i]
        return result

    cpdef np.ndarray predict(self, object X):
        # Check input
        if not isinstance(X, np.ndarray):
            raise ValueError("X should be in np.ndarray format, got %s"
                             % type(X))

        if X.dtype != DTYPE:
            raise ValueError("X.dtype should be np.float32, got %s" % X.dtype)

        # Extract input
        cdef DTYPE_t[:, :] X_ndarray = X
        cdef SIZE_t n_samples = X.shape[0]

        # Initialize output
        cdef np.ndarray[SIZE_t] out = np.zeros((n_samples,), dtype=np.int16)
        cdef SIZE_t* out_ptr = <SIZE_t*> out.data

        # Initialize auxiliary data-structure
        cdef Node* node = NULL
        cdef SIZE_t i = 0
        cdef INT16_t node_next

        with nogil:
        #if 1:
            for i in range(n_samples):
                node = self.nodes
                if X_ndarray[i, node.feature] <= 0.5:
                    node_next=node.left_child
                else:
                    node_next=node.right_child

                # While node not a leaf
                while node_next>0:
                    node = &self.nodes[node_next]
                    if X_ndarray[i, node.feature] <= 0.5:
                        node_next=node.left_child
                    else:
                        node_next=node.right_child

                out_ptr[i] = <SIZE_t>(-node_next)

        return out

    cdef np.ndarray _get_node_ndarray(self):
        """Wraps nodes as a NumPy struct array.

        The array keeps a reference to this Tree, which manages the underlying
        memory. Individual fields are publicly accessible as properties of the
        Tree.
        """
        cdef np.npy_intp shape[1]
        shape[0] = <np.npy_intp> self.node_count
        cdef np.npy_intp strides[1]
        strides[0] = sizeof(Node)
        cdef np.ndarray arr
        Py_INCREF(NODE_DTYPE)
        arr = PyArray_NewFromDescr(np.ndarray, <np.dtype> NODE_DTYPE, 1, shape,
                                   strides, <void*> self.nodes,
                                   np.NPY_DEFAULT, None)
        Py_INCREF(self)
        arr.base = <PyObject*> self
        return arr
