# distutils: language = c
# cython: wraparound=False, cdivision=True, boundscheck=False, initializedcheck=False

cimport cython
import numpy as np
cimport numpy as cnp

from dipy.segment.clustering import ClusterCentroid, ClusterMapCentroid
from dipy.segment.clustering import TreeCluster, TreeClusterMap


from libc.math cimport fabs
from dipy.segment.cythonutils cimport Data2D, Shape, shape2tuple,\
    tuple2shape, same_shape, create_memview_2d, free_memview_2d

cdef extern from "math.h" nogil:
    double fabs(double x)

cdef extern from "stdlib.h" nogil:
    ctypedef unsigned long size_t
    void free(void *ptr)
    void *malloc(size_t elsize)
    void *calloc(size_t nelem, size_t elsize)
    void *realloc(void *ptr, size_t elsize)
    void *memset(void *ptr, int value, size_t num)

DTYPE = np.float32
DEF BIGGEST_DOUBLE = 1.7976931348623157e+308  # np.finfo('f8').max
DEF BIGGEST_INT = 2147483647  # np.iinfo('i4').max
DEF BIGGEST_FLOAT = 3.4028235e+38  # np.finfo('f4').max
DEF SMALLEST_FLOAT = -3.4028235e+38  # np.finfo('f4').max


cdef print_node(CentroidNode* node, prepend=""):
    if node == NULL:
        return ""
    cdef Data2D centroid
    centroid = <float[:node.centroid_shape.dims[0],:node.centroid_shape.dims[1]]> &node.centroid[0][0,0]
    txt = "{}".format(np.asarray(centroid).tolist())
    txt += " {" + ",".join(map(str, np.asarray(<int[:node.size]> node.indices))) + "}"
    txt += " children({})".format(node.nb_children)
    txt += " count({})".format(node.size)
    txt += " thres({})".format(node.threshold)
    txt += "\n"

    cdef int i
    for i in range(node.nb_children):
        txt += prepend
        if i == node.nb_children-1:
            # Last child
            txt += "`-- " + print_node(node.children[i], prepend + "    ")
        else:
            txt += "|-- " + print_node(node.children[i], prepend + "|   ")

    return txt


cdef void aabb_creation(Data2D streamline, float* aabb) nogil:
    """ Creates AABB enveloping the given streamline.

        Notes
        -----
        This currently assumes streamline is made of 3D points.
    """
    cdef:
        int N = streamline.shape[0], D = streamline.shape[1]
        int n, d
        float min_[3]
        float max_[3]

    for d in range(D):
        min_[d] = BIGGEST_FLOAT
        max_[d] = SMALLEST_FLOAT
        for n in range(N):

            if max_[d] < streamline[n, d]:
                max_[d] = streamline[n, d]

            if min_[d] > streamline[n, d]:
                min_[d] = streamline[n, d]

        aabb[d + 3] = (max_[d] - min_[d]) / 2.0 # radius
        aabb[d] = min_[d] + aabb[d + 3]  # center


cdef inline int aabb_overlap(float* aabb1, float* aabb2, float padding=0.) nogil:
    """ SIMD optimized AABB-AABB test

    Optimized by removing conditional branches
    """
    cdef:
        int x = fabs(aabb1[0] - aabb2[0]) <= (aabb1[3] + aabb2[3] + padding)
        int y = fabs(aabb1[1] - aabb2[1]) <= (aabb1[4] + aabb2[4] + padding)
        int z = fabs(aabb1[2] - aabb2[2]) <= (aabb1[5] + aabb2[5] + padding)

    return x & y & z


cdef CentroidNode* create_empty_node(Shape centroid_shape, float threshold) nogil:
    # Important: because the CentroidNode structure contains an uninitialized memview,
    # we need to zero-initialize the allocated memory (calloc or via memset),
    # otherwise during assignment CPython will try to call _PYX_XDEC_MEMVIEW on it and segfault.
    cdef CentroidNode* node = <CentroidNode*> calloc(1, sizeof(CentroidNode))
    node.centroid = create_memview_2d(centroid_shape.size, centroid_shape.dims)
        #node.updated_centroid = <float[:centroid_shape.dims[0], :centroid_shape.dims[1]]> calloc(centroid_shape.size, sizeof(float))

    node.father = NULL
    node.children = NULL
    node.nb_children = 0
    node.aabb[0] = 0
    node.aabb[1] = 0
    node.aabb[2] = 0
    node.aabb[3] = BIGGEST_FLOAT
    node.aabb[4] = BIGGEST_FLOAT
    node.aabb[5] = BIGGEST_FLOAT
    node.threshold = threshold
    node.indices = NULL
    node.size = 0
    node.centroid_shape = centroid_shape
    return node




cdef class QuickBundlesX(object):

    def __init__(self, features_shape, levels_thresholds, Metric metric):
        self.metric = metric
        self.features_shape = tuple2shape(features_shape)

        self.nb_levels = len(levels_thresholds)
        self.thresholds = <double*> malloc(self.nb_levels*sizeof(double))

        cdef int i
        for i in range(self.nb_levels):
            self.thresholds[i] = levels_thresholds[i]

        self.root = create_empty_node(self.features_shape, self.thresholds[0])

        self.level = None
        self.clusters = None
        self.stats.stats_per_layer = <QuickBundlesXStatsLayer*> calloc(self.nb_levels, sizeof(QuickBundlesXStatsLayer))
        # Important: because the CentroidNode structure contains an uninitialized memview,
        # we need to zero-initialize the allocated memory (calloc or via memset),
        # otherwise during assignment CPython will try to call _PYX_XDEC_MEMVIEW on it and segfault.
        self.current_streamline = <StreamlineInfos*> calloc(1, sizeof(StreamlineInfos))
        self.current_streamline.features = create_memview_2d(self.features_shape.size, self.features_shape.dims)
        self.current_streamline.features_flip = create_memview_2d(self.features_shape.size, self.features_shape.dims)

    def __dealloc__(self):
        self.traverse_postorder(self.root, self._dealloc_node)
        self.root = NULL

        if self.thresholds != NULL:
            free(self.thresholds)
            self.thresholds = NULL

        if self.stats.stats_per_layer != NULL:
            free(self.stats.stats_per_layer)
            self.stats.stats_per_layer = NULL

        if self.current_streamline != NULL:
            free(self.current_streamline)
            self.current_streamline = NULL

    cdef int _add_child(self, CentroidNode* node) nogil:
        cdef double threshold = 0.0  # Leaf node doesn't need threshold.
        if node.level+1 < self.nb_levels:
            threshold = self.thresholds[node.level+1]

        cdef CentroidNode* child = create_empty_node(self.features_shape, threshold)
        child.level = node.level+1

        # Add new child.
        child.father = node
        node.children = <CentroidNode**> realloc(node.children, (node.nb_children+1)*sizeof(CentroidNode*))
        node.children[node.nb_children] = child
        node.nb_children += 1

        return node.nb_children-1

    cdef void _update_node(self, CentroidNode* node, StreamlineInfos* streamline_infos) nogil:
        cdef Data2D element = streamline_infos.features[0]
        cdef int C = node.size
        cdef cnp.npy_intp n, d

        if streamline_infos.use_flip:
            element = streamline_infos.features_flip[0]

        # Update centroid
        cdef Data2D centroid = node.centroid[0]
        cdef cnp.npy_intp N = centroid.shape[0], D = centroid.shape[1]
        for n in range(N):
            for d in range(D):
                centroid[n, d] = ((centroid[n, d] * C) + element[n, d]) / (C+1)

        # Update list of indices
        node.indices = <int*> realloc(node.indices, (C+1)*sizeof(int))
        node.indices[C] = streamline_infos.idx
        node.size += 1

        # Update AABB
        aabb_creation(centroid, node.aabb)

    cdef void _insert_in(self, CentroidNode* node, StreamlineInfos* streamline_infos, int[:] path) nogil:
        cdef:
            float dist, dist_flip
            cnp.npy_intp k
            NearestCluster nearest_cluster

        self._update_node(node, streamline_infos)

        if node.level == self.nb_levels:
            return

        nearest_cluster.id = -1
        nearest_cluster.dist = BIGGEST_DOUBLE
        nearest_cluster.flip = 0

        for k in range(node.nb_children):
            # Check streamline's aabb colides with the current child.
            self.stats.stats_per_layer[node.level].nb_aabb_calls += 1
            if aabb_overlap(node.children[k].aabb, streamline_infos.aabb, node.threshold):
                self.stats.stats_per_layer[node.level].nb_mdf_calls += 1
                dist = self.metric.c_dist(node.children[k].centroid[0], streamline_infos.features[0])

                # Keep track of the nearest cluster
                if dist < nearest_cluster.dist:
                    nearest_cluster.dist = dist
                    nearest_cluster.id = k
                    nearest_cluster.flip = 0

                self.stats.stats_per_layer[node.level].nb_mdf_calls += 1
                dist_flip = self.metric.c_dist(node.children[k].centroid[0], streamline_infos.features_flip[0])
                if dist_flip < nearest_cluster.dist:
                    nearest_cluster.dist = dist_flip
                    nearest_cluster.id = k
                    nearest_cluster.flip = 1

        if nearest_cluster.dist > node.threshold:
            # No near cluster, create a new one.
            nearest_cluster.id = self._add_child(node)

        streamline_infos.use_flip = nearest_cluster.flip
        path[node.level] = nearest_cluster.id
        self._insert_in(node.children[nearest_cluster.id], streamline_infos, path)

    cpdef object insert(self, Data2D datum, int datum_idx):
        self.metric.feature.c_extract(datum, self.current_streamline.features[0])
        self.metric.feature.c_extract(datum[::-1], self.current_streamline.features_flip[0])
        self.current_streamline.idx = datum_idx

        aabb_creation(self.current_streamline.features[0], self.current_streamline.aabb)
        path = -1 * np.ones(self.nb_levels, dtype=np.int32)
        self._insert_in(self.root, self.current_streamline, path)
        return path

    def __str__(self):
        return print_node(self.root)

    cdef void traverse_postorder(self, CentroidNode* node, void (*visit)(QuickBundlesX, CentroidNode*)):
        cdef int i
        for i in range(node.nb_children):
            self.traverse_postorder(node.children[i], visit)
        visit(self, node)

    cdef void _dealloc_node(self, CentroidNode* node):
        free_memview_2d(node.centroid)

        if node.children != NULL:
            free(node.children)
            node.children = NULL

        free(node.indices)
        node.indices = NULL

        # No need to free node.father, only the current node.
        free(node)

    cdef object _build_tree_clustermap(self, CentroidNode* node):
        cdef Data2D centroid
        centroid = <float[:self.features_shape.dims[0],:self.features_shape.dims[1]]> &node.centroid[0][0,0]
        tree_cluster = TreeCluster(threshold=node.threshold,
                                   centroid=np.asarray(centroid),
                                   indices=np.asarray(<int[:node.size]> node.indices).copy())
        cdef int i
        for i in range(node.nb_children):
            tree_cluster.add(self._build_tree_clustermap(node.children[i]))

        return tree_cluster

    def get_tree_cluster_map(self):
        return TreeClusterMap(self._build_tree_clustermap(self.root))

    def get_stats(self):
        stats_per_level = []
        for i in range(self.nb_levels):
            stats_per_level.append({'nb_mdf_calls': self.stats.stats_per_layer[i].nb_mdf_calls,
                                    'nb_aabb_calls': self.stats.stats_per_layer[i].nb_aabb_calls,
                                    'threshold': self.thresholds[i]})

        stats = {'stats_per_level': stats_per_level}

        return stats


cdef class Clusters:
    """ Provides Cython functionalities to interact with clustering outputs.

    This class allows to create clusters and assign elements to them.
    Assignements of a cluster are represented as a list of element indices.
    """
    def __init__(Clusters self):
        self._nb_clusters = 0
        self.clusters_indices = NULL
        self.clusters_size = NULL

    def __dealloc__(Clusters self):
        """ Deallocates memory created with `c_create_cluster` and `c_assign`. """
        for i in range(self._nb_clusters):
            free(self.clusters_indices[i])
            self.clusters_indices[i] = NULL

        free(self.clusters_indices)
        self.clusters_indices = NULL
        free(self.clusters_size)
        self.clusters_size = NULL

    cdef int c_size(Clusters self) nogil:
        """ Returns the number of clusters. """
        return self._nb_clusters

    cdef void c_assign(Clusters self, int id_cluster, int id_element, Data2D element) nogil except *:
        """ Assigns an element to a cluster.

        Parameters
        ----------
        id_cluster : int
            Index of the cluster to which the element will be assigned.
        id_element : int
            Index of the element to assign.
        element : 2d array (float)
            Data of the element to assign.
        """
        cdef cnp.npy_intp C = self.clusters_size[id_cluster]
        self.clusters_indices[id_cluster] = <int*> realloc(self.clusters_indices[id_cluster], (C+1)*sizeof(int))
        self.clusters_indices[id_cluster][C] = id_element
        self.clusters_size[id_cluster] += 1

    cdef int c_create_cluster(Clusters self) nogil except -1:
        """ Creates a cluster and adds it at the end of the list.

        Returns
        -------
        id_cluster : int
            Index of the new cluster.
        """
        self.clusters_indices = <int**> realloc(self.clusters_indices, (self._nb_clusters+1)*sizeof(int*))
        self.clusters_indices[self._nb_clusters] = <int*> calloc(0, sizeof(int))
        self.clusters_size = <int*> realloc(self.clusters_size, (self._nb_clusters+1)*sizeof(int))
        self.clusters_size[self._nb_clusters] = 0

        self._nb_clusters += 1
        return self._nb_clusters - 1


cdef class ClustersCentroid(Clusters):
    """ Provides Cython functionalities to interact with clustering outputs
    having the notion of cluster's centroid.

    This class allows to create clusters, assign elements to them and
    update their centroid.

    Parameters
    ----------
    centroid_shape : int, tuple of int
        Information about the shape of the centroid.
    eps : float, optional
        Consider the centroid has not changed if the changes per dimension
        are less than this epsilon. (Default: 1e-6)
    """
    def __init__(ClustersCentroid self, centroid_shape, float eps=1e-6, *args, **kwargs):
        Clusters.__init__(self, *args, **kwargs)
        if isinstance(centroid_shape, int):
            centroid_shape = (1, centroid_shape)

        if not isinstance(centroid_shape, tuple):
            raise ValueError("'centroid_shape' must be a tuple or a int.")

        self._centroid_shape = tuple2shape(centroid_shape)
        self.centroids = NULL
        self._updated_centroids = NULL
        self.eps = eps

    def __dealloc__(ClustersCentroid self):
        """ Deallocates memory created with `c_create_cluster` and `c_assign`.

        Notes
        -----
        The `__dealloc__` method of the superclass is automatically called:
        http://docs.cython.org/src/userguide/special_methods.html#finalization-method-dealloc
        """
        cdef cnp.npy_intp i
        for i in range(self._nb_clusters):
            free_memview_2d(self.centroids[i].features)
            free_memview_2d(self._updated_centroids[i].features)


        free(self.centroids)
        self.centroids = NULL
        free(self._updated_centroids)
        self._updated_centroids = NULL

    cdef void c_assign(ClustersCentroid self, int id_cluster, int id_element, Data2D element) nogil except *:
        """ Assigns an element to a cluster.

        In addition of keeping element's index, an updated version of the
        cluster's centroid is computed. The centroid is the average of all
        elements in a cluster.

        Parameters
        ----------
        id_cluster : int
            Index of the cluster to which the element will be assigned.
        id_element : int
            Index of the element to assign.
        element : 2d array (float)
            Data of the element to assign.
        """
        cdef Data2D updated_centroid = self._updated_centroids[id_cluster].features[0]
        cdef cnp.npy_intp C = self.clusters_size[id_cluster]
        cdef cnp.npy_intp n, d

        cdef cnp.npy_intp N = updated_centroid.shape[0], D = updated_centroid.shape[1]
        for n in range(N):
            for d in range(D):
                updated_centroid[n, d] = ((updated_centroid[n, d] * C) + element[n, d]) / (C+1)

        Clusters.c_assign(self, id_cluster, id_element, element)

    cdef int c_update(ClustersCentroid self, cnp.npy_intp id_cluster) nogil except -1:
        """ Update the centroid of a cluster.

        Parameters
        ----------
        id_cluster : int
            Index of the cluster of which its centroid will be updated.

        Returns
        -------
        int
            Tells whether the centroid has changed or not, i.e. converged.
        """
        cdef Data2D centroid = self.centroids[id_cluster].features[0]
        cdef Data2D updated_centroid = self._updated_centroids[id_cluster].features[0]
        cdef cnp.npy_intp N = updated_centroid.shape[0], D = centroid.shape[1]
        cdef cnp.npy_intp n, d
        cdef int converged = 1

        for n in range(N):
            for d in range(D):
                converged &= fabs(centroid[n, d] - updated_centroid[n, d]) < self.eps
                centroid[n, d] = updated_centroid[n, d]

        #cdef float * aabb = &self.centroids[id_cluster].aabb[0]

        aabb_creation(centroid, self.centroids[id_cluster].aabb)

        return converged

    cdef int c_create_cluster(ClustersCentroid self) nogil except -1:
        """ Creates a cluster and adds it at the end of the list.

        Returns
        -------
        id_cluster : int
            Index of the new cluster.
        """
        self.centroids = <Centroid*> realloc(self.centroids, (self._nb_clusters+1) * sizeof(Centroid))
        # Zero-initialize the Centroid structure
        memset(&self.centroids[self._nb_clusters], 0, sizeof(Centroid))

        self._updated_centroids = <Centroid*> realloc(self._updated_centroids, (self._nb_clusters+1) * sizeof(Centroid))
        # Zero-initialize the new Centroid structure
        memset(&self._updated_centroids[self._nb_clusters], 0, sizeof(Centroid))

        self.centroids[self._nb_clusters].features = create_memview_2d(self._centroid_shape.size, self._centroid_shape.dims)
        self._updated_centroids[self._nb_clusters].features = create_memview_2d(self._centroid_shape.size, self._centroid_shape.dims)

        aabb_creation(self.centroids[self._nb_clusters].features[0], self.centroids[self._nb_clusters].aabb)

        return Clusters.c_create_cluster(self)


cdef class QuickBundles(object):
    def __init__(QuickBundles self, features_shape, Metric metric, double threshold,
                 int max_nb_clusters=BIGGEST_INT):
        self.metric = metric
        self.features_shape = tuple2shape(features_shape)
        self.threshold = threshold
        self.max_nb_clusters = max_nb_clusters
        self.clusters = ClustersCentroid(features_shape)
        self.features = np.empty(features_shape, dtype=DTYPE)
        self.features_flip = np.empty(features_shape, dtype=DTYPE)

        self.stats.nb_mdf_calls = 0
        self.stats.nb_aabb_calls = 0

    cdef NearestCluster find_nearest_cluster(QuickBundles self, Data2D features) nogil except *:
        """ Finds the nearest cluster of a datum given its `features` vector.

        Parameters
        ----------
        features : 2D array
            Features of a datum.

        Returns
        -------
        `NearestCluster` object
            Nearest cluster to `features` according to the given metric.
        """
        cdef:
            cnp.npy_intp k
            double dist
            NearestCluster nearest_cluster
            float aabb[6]

        nearest_cluster.id = -1
        nearest_cluster.dist = BIGGEST_DOUBLE
        nearest_cluster.flip = 0

        for k in range(self.clusters.c_size()):

            self.stats.nb_mdf_calls += 1
            dist = self.metric.c_dist(self.clusters.centroids[k].features[0], features)

            # Keep track of the nearest cluster
            if dist < nearest_cluster.dist:
                nearest_cluster.dist = dist
                nearest_cluster.id = k


        return nearest_cluster

    cdef int assignment_step(QuickBundles self, Data2D datum, int datum_id) nogil except -1:
        """ Compute the assignment step of the QuickBundles algorithm.

        It will assign a datum to its closest cluster according to a given
        metric. If the distance between the datum and its closest cluster is
        greater than the specified threshold, a new cluster is created and the
        datum is assigned to it.

        Parameters
        ----------
        datum : 2D array
            The datum to assign.
        datum_id : int
            ID of the datum, usually its index.

        Returns
        -------
        int
            Index of the cluster the datum has been assigned to.
        """
        cdef:
            Data2D features_to_add = self.features
            NearestCluster nearest_cluster, nearest_cluster_flip
            Shape features_shape = self.metric.feature.c_infer_shape(datum)

        # Check if datum is compatible with the metric
        if not same_shape(features_shape, self.features_shape):
            with gil:
                raise ValueError("All features do not have the same shape! QuickBundles requires this to compute centroids!")

        # Check if datum is compatible with the metric
        if not self.metric.c_are_compatible(features_shape, self.features_shape):
            with gil:
                raise ValueError("Data features' shapes must be compatible according to the metric used!")

        # Find nearest cluster to datum
        self.metric.feature.c_extract(datum, self.features)
        nearest_cluster = self.find_nearest_cluster(self.features)

        # Find nearest cluster to s_i_flip if metric is not order invariant
        if not self.metric.feature.is_order_invariant:
            self.metric.feature.c_extract(datum[::-1], self.features_flip)
            nearest_cluster_flip = self.find_nearest_cluster(self.features_flip)

            # If we found a lower distance using a flipped datum,
            #  add the flipped version instead
            if nearest_cluster_flip.dist < nearest_cluster.dist:
                nearest_cluster.id = nearest_cluster_flip.id
                nearest_cluster.dist = nearest_cluster_flip.dist
                features_to_add = self.features_flip

        # Check if distance with the nearest cluster is below some threshold
        # or if we already have the maximum number of clusters.
        # If the former or the latter is true, assign datum to its nearest cluster
        # otherwise create a new cluster and assign the datum to it.
        if not (nearest_cluster.dist < self.threshold or self.clusters.c_size() >= self.max_nb_clusters):
            nearest_cluster.id = self.clusters.c_create_cluster()

        self.clusters.c_assign(nearest_cluster.id, datum_id, features_to_add)
        return nearest_cluster.id

    cdef void update_step(QuickBundles self, int cluster_id) nogil except *:
        """ Compute the update step of the QuickBundles algorithm.

        It will update the centroid of a cluster given its index.

        Parameters
        ----------
        cluster_id : int
            ID of the cluster to update.

        """
        self.clusters.c_update(cluster_id)

    def get_stats(self):
        stats = {'nb_mdf_calls': self.stats.nb_mdf_calls,
                 'nb_aabb_calls': self.stats.nb_aabb_calls}

        return stats

    cdef object _build_clustermap(self):
        clusters = ClusterMapCentroid()
        cdef int k
        for k in range(self.clusters.c_size()):
            cluster = ClusterCentroid(np.asarray(self.clusters.centroids[k].features[0]).copy())
            cluster.indices = np.asarray(<int[:self.clusters.clusters_size[k]]> self.clusters.clusters_indices[k]).copy()
            clusters.add_cluster(cluster)

        return clusters

    def get_cluster_map(self):
        return self._build_clustermap()


def evaluate_aabb_checks():
    cdef:
        Data2D feature1 = np.array([[1, 0, 0], [1, 1, 0], [1 + np.sqrt(2)/2., 1 + np.sqrt(2)/2., 0]], dtype='f4')
        Data2D feature2 = np.array([[1, 0, 0], [1, 1, 0], [1 + np.sqrt(2)/2., 1 + np.sqrt(2)/2., 0]], dtype='f4') + np.array([0.5, 0, 0], dtype='f4')
        float[6] aabb1
        float[6] aabb2
        int res

    aabb_creation(feature1, &aabb1[0])
    aabb_creation(feature2, &aabb2[0])

    res = aabb_overlap(&aabb1[0], &aabb2[0])

    return np.asarray(aabb1), np.asarray(aabb2), res

