# Author: Christopher Moody <chrisemoody@gmail.com>
# Author: Nick Travers <nickt@squareup.com>
# Implementation by Chris Moody & Nick Travers
# See http://homepage.tudelft.nl/19j49/t-SNE.html for reference
# implementations and papers describing the technique


import numpy as np
cimport numpy as cnp
from libc.stdio cimport printf
from libc.math cimport log
from libc.stdlib cimport malloc, free
from libc.time cimport clock, clock_t
from cython.parallel cimport prange, parallel

from ..neighbors._quad_tree cimport _QuadTree

cnp.import_array()


cdef char* EMPTY_STRING = ""

# Smallest strictly positive value that can be represented by floating
# point numbers for different precision levels. This is useful to avoid
# taking the log of zero when computing the KL divergence.
cdef float FLOAT32_TINY = np.finfo(np.float32).tiny

# Useful to void division by zero or divergence to +inf.
cdef float FLOAT64_EPS = np.finfo(np.float64).eps

# This is effectively an ifdef statement in Cython
# It allows us to write printf debugging lines
# and remove them at compile time
cdef enum:
    DEBUGFLAG = 0

cdef float compute_gradient(float[:] val_P,
                            float[:, :] pos_reference,
                            cnp.int64_t[:] neighbors,
                            cnp.int64_t[:] indptr,
                            float[:, :] tot_force,
                            _QuadTree qt,
                            float theta,
                            int dof,
                            long start,
                            bint compute_error,
                            int num_threads) noexcept nogil:
    # Having created the tree, calculate the gradient
    # in two components, the positive and negative forces
    cdef:
        long i, coord
        int ax
        long n_samples = pos_reference.shape[0]
        int n_dimensions = qt.n_dimensions
        clock_t t1 = 0, t2 = 0
        double sQ
        float error
        int take_timing = 1 if qt.verbose > 15 else 0

    if qt.verbose > 11:
        printf("[t-SNE] Allocating %li elements in force arrays\n",
               n_samples * n_dimensions * 2)
    cdef float* neg_f = <float*> malloc(sizeof(float) * n_samples * n_dimensions)
    cdef float* pos_f = <float*> malloc(sizeof(float) * n_samples * n_dimensions)

    if take_timing:
        t1 = clock()
    sQ = compute_gradient_negative(pos_reference, neg_f, qt, dof, theta, start,
                                   num_threads)
    if take_timing:
        t2 = clock()
        printf("[t-SNE] Computing negative gradient: %e ticks\n", ((float) (t2 - t1)))

    if take_timing:
        t1 = clock()
    error = compute_gradient_positive(val_P, pos_reference, neighbors, indptr,
                                      pos_f, n_dimensions, dof, sQ, start,
                                      qt.verbose, compute_error, num_threads)
    if take_timing:
        t2 = clock()
        printf("[t-SNE] Computing positive gradient: %e ticks\n",
               ((float) (t2 - t1)))
    for i in prange(start, n_samples, nogil=True, num_threads=num_threads,
                    schedule='static'):
        for ax in range(n_dimensions):
            coord = i * n_dimensions + ax
            tot_force[i, ax] = pos_f[coord] - (neg_f[coord] / sQ)

    free(neg_f)
    free(pos_f)
    return error


cdef float compute_gradient_positive(float[:] val_P,
                                     float[:, :] pos_reference,
                                     cnp.int64_t[:] neighbors,
                                     cnp.int64_t[:] indptr,
                                     float* pos_f,
                                     int n_dimensions,
                                     int dof,
                                     double sum_Q,
                                     cnp.int64_t start,
                                     int verbose,
                                     bint compute_error,
                                     int num_threads) noexcept nogil:
    # Sum over the following expression for i not equal to j
    # grad_i = p_ij (1 + ||y_i - y_j||^2)^-1 (y_i - y_j)
    # This is equivalent to compute_edge_forces in the authors' code
    # It just goes over the nearest neighbors instead of all the data points
    # (unlike the non-nearest neighbors version of `compute_gradient_positive')
    cdef:
        int ax
        long i, j, k
        long n_samples = indptr.shape[0] - 1
        float C = 0.0
        float dij, qij, pij
        float exponent = (dof + 1.0) / 2.0
        float float_dof = (float) (dof)
        float* buff
        clock_t t1 = 0, t2 = 0
        float dt

    if verbose > 10:
        t1 = clock()

    with nogil, parallel(num_threads=num_threads):
        # Define private buffer variables
        buff = <float *> malloc(sizeof(float) * n_dimensions)

        for i in prange(start, n_samples, schedule='static'):
            # Init the gradient vector
            for ax in range(n_dimensions):
                pos_f[i * n_dimensions + ax] = 0.0
            # Compute the positive interaction for the nearest neighbors
            for k in range(indptr[i], indptr[i+1]):
                j = neighbors[k]
                dij = 0.0
                pij = val_P[k]
                for ax in range(n_dimensions):
                    buff[ax] = pos_reference[i, ax] - pos_reference[j, ax]
                    dij += buff[ax] * buff[ax]
                qij = float_dof / (float_dof + dij)
                if dof != 1:  # i.e. exponent != 1
                    qij = qij ** exponent
                dij = pij * qij

                # only compute the error when needed
                if compute_error:
                    qij = qij / sum_Q
                    C += pij * log(max(pij, FLOAT32_TINY) / max(qij, FLOAT32_TINY))
                for ax in range(n_dimensions):
                    pos_f[i * n_dimensions + ax] += dij * buff[ax]

        free(buff)
    if verbose > 10:
        t2 = clock()
        dt = ((float) (t2 - t1))
        printf("[t-SNE] Computed error=%1.4f in %1.1e ticks\n", C, dt)
    return C


cdef double compute_gradient_negative(float[:, :] pos_reference,
                                      float* neg_f,
                                      _QuadTree qt,
                                      int dof,
                                      float theta,
                                      long start,
                                      int num_threads) noexcept nogil:
    cdef:
        int ax
        int n_dimensions = qt.n_dimensions
        int offset = n_dimensions + 2
        long i, j, idx
        long n_samples = pos_reference.shape[0]
        long n = n_samples - start
        long dta = 0
        long dtb = 0
        float size, dist2s, mult
        float exponent = (dof + 1.0) / 2.0
        float float_dof = (float) (dof)
        double qijZ, sum_Q = 0.0
        float* force
        float* neg_force
        float* pos
        clock_t t1 = 0, t2 = 0, t3 = 0
        int take_timing = 1 if qt.verbose > 20 else 0

    with nogil, parallel(num_threads=num_threads):
        # Define thread-local buffers
        summary = <float*> malloc(sizeof(float) * n * offset)
        pos = <float *> malloc(sizeof(float) * n_dimensions)
        force = <float *> malloc(sizeof(float) * n_dimensions)
        neg_force = <float *> malloc(sizeof(float) * n_dimensions)

        for i in prange(start, n_samples, schedule='static'):
            # Clear the arrays
            for ax in range(n_dimensions):
                force[ax] = 0.0
                neg_force[ax] = 0.0
                pos[ax] = pos_reference[i, ax]

            # Find which nodes are summarizing and collect their centers of mass
            # deltas, and sizes, into vectorized arrays
            if take_timing:
                t1 = clock()
            idx = qt.summarize(pos, summary, theta*theta)
            if take_timing:
                t2 = clock()
            # Compute the t-SNE negative force
            # for the digits dataset, walking the tree
            # is about 10-15x more expensive than the
            # following for loop
            for j in range(idx // offset):

                dist2s = summary[j * offset + n_dimensions]
                size = summary[j * offset + n_dimensions + 1]
                qijZ = float_dof / (float_dof + dist2s)  # 1/(1+dist)
                if dof != 1:  # i.e. exponent != 1
                    qijZ = qijZ ** exponent

                sum_Q += size * qijZ   # size of the node * q
                mult = size * qijZ * qijZ
                for ax in range(n_dimensions):
                    neg_force[ax] += mult * summary[j * offset + ax]
            if take_timing:
                t3 = clock()
            for ax in range(n_dimensions):
                neg_f[i * n_dimensions + ax] = neg_force[ax]
            if take_timing:
                dta += t2 - t1
                dtb += t3 - t2
        free(pos)
        free(force)
        free(neg_force)
        free(summary)
    if take_timing:
        printf("[t-SNE] Tree: %li clock ticks | ", dta)
        printf("Force computation: %li clock ticks\n", dtb)

    # Put sum_Q to machine EPSILON to avoid divisions by 0
    sum_Q = max(sum_Q, FLOAT64_EPS)
    return sum_Q


def gradient(float[:] val_P,
             float[:, :] pos_output,
             cnp.int64_t[:] neighbors,
             cnp.int64_t[:] indptr,
             float[:, :] forces,
             float theta,
             int n_dimensions,
             int verbose,
             int dof=1,
             long skip_num_points=0,
             bint compute_error=1,
             int num_threads=1):
    # This function is designed to be called from external Python
    # it passes the 'forces' array by reference and fills that's array
    # up in-place
    cdef float C
    cdef int n
    n = pos_output.shape[0]
    assert val_P.itemsize == 4
    assert pos_output.itemsize == 4
    assert forces.itemsize == 4
    m = "Forces array and pos_output shapes are incompatible"
    assert n == forces.shape[0], m
    m = "Pij and pos_output shapes are incompatible"
    assert n == indptr.shape[0] - 1, m
    if verbose > 10:
        printf("[t-SNE] Initializing tree of n_dimensions %i\n", n_dimensions)
    cdef _QuadTree qt = _QuadTree(pos_output.shape[1], verbose)
    if verbose > 10:
        printf("[t-SNE] Inserting %li points\n", pos_output.shape[0])
    qt.build_tree(pos_output)
    if verbose > 10:
        # XXX: format hack to workaround lack of `const char *` type
        # in the generated C code that triggers error with gcc 4.9
        # and -Werror=format-security
        printf("[t-SNE] Computing gradient\n%s", EMPTY_STRING)

    C = compute_gradient(val_P, pos_output, neighbors, indptr, forces,
                         qt, theta, dof, skip_num_points, compute_error,
                         num_threads)

    if verbose > 10:
        # XXX: format hack to workaround lack of `const char *` type
        # in the generated C code
        # and -Werror=format-security
        printf("[t-SNE] Checking tree consistency\n%s", EMPTY_STRING)
    m = "Tree consistency failed: unexpected number of points on the tree"
    assert qt.cells[0].cumulative_size == qt.n_points, m
    if not compute_error:
        C = np.nan
    return C