
# A cythoned version of the original function of the stats_process function,
# based on _kinetics.pyx

from __future__ import print_function,division
import cython
import numpy
import numpy as np
import h5py
from scipy.sparse import csgraph
import warnings
from collections import Counter
cimport numpy
cimport numpy as np
cimport scipy.linalg
cimport scipy.linalg.cython_lapack as cl
import scipy.linalg
from libc.math cimport isnan

ctypedef numpy.uint16_t index_t
ctypedef numpy.float64_t weight_t
ctypedef numpy.uint8_t bool_t
ctypedef numpy.int64_t trans_t
ctypedef numpy.uintp_t uint_t # 32 bits on 32-bit systems, 64 bits on 64-bit systems
ctypedef unsigned short Ushort
ctypedef double complex Cdouble

weight_dtype = numpy.float64
index_dtype = numpy.uint16
bool_dtype = numpy.bool_
intc_dtype = numpy.intc

@cython.boundscheck(False)
@cython.wraparound(False)
cpdef stats_process(numpy.ndarray[index_t, ndim=2] bin_assignments,
                    numpy.ndarray[weight_t, ndim=1] weights,
                    numpy.ndarray[weight_t, ndim=2] fluxes,
                    numpy.ndarray[weight_t, ndim=1] populations,
                    numpy.ndarray[trans_t, ndim=2] trans,
                    numpy.ndarray[index_t, ndim=2] mask,
                    str interval='timepoint'                        ):
    cdef:
        Py_ssize_t i,k
        index_t ibin,fbin,nsegs,npts
    nsegs = bin_assignments.shape[0]
    npts = bin_assignments.shape[1]

    if interval == 'timepoint':
        for i in xrange(0,npts - 1):
            for k in xrange(nsegs):
                ibin = bin_assignments[k,i]
                fbin = bin_assignments[k, i + 1]

                if mask[k, 0] == 1:
                    continue

                w = weights[k]

                fluxes[ibin, fbin] += w
                trans[ibin, fbin] += 1
                populations[ibin] += w
        return

    if interval == 'iteration':
        for k in xrange(nsegs):
            # Should this be 0?
            # .... this should super be 0.  What?
            ibin = bin_assignments[k,0]
            fbin = bin_assignments[k, npts - 1]

            if mask[k, 0] == 1:
                continue

            w = weights[k]

            fluxes[ibin, fbin] += w
            trans[ibin, fbin] += 1
            populations[ibin] += w
        return

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cpdef int normalize(weight_t[:,:] m, Py_ssize_t nfbins) nogil:

    cdef:
        weight_t row_sum
        Py_ssize_t x, y

    for y in range(nfbins):
        row_sum = 0
        for x in range(nfbins):
            row_sum += m[y,x]
        if row_sum != 0:
            for x in range(nfbins):
                m[y,x] /= row_sum
    return 0

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cpdef reweight_for_c(rows, cols, obs, flux, insert, indices, nstates, nbins, state_labels, state_map, nfbins, istate, jstate, stride, bin_last_state_map, bin_state_map, return_obs, obs_threshold=1):



    # Instead of pulling in start and stop, we'll pull in a list of indices.
    # This way, it should support the bootstrap.
    cdef:
        int[:] _rows, _cols, _obs, _ins, _nrows, _ncols, _nobs, _visited
        long[:]  _bin_last_state_map
        weight_t[:] _flux, _total_pop, _rw_bin_probs, _nflux, _rw_color_probs, _rw_state_probs
        double[:] _eigvals, _eigvalsi
        int n_trans, _nstates, lind, _nfbins, _stride, _obs_threshold, nnz, nlind, i, j, _istate, _jstate
        Ushort[:] _indices, _bin_state_map
        weight_t _return_value

        Ushort[:] _new_indices

        weight_t[:,:] _total_fluxes, _transition_matrix, _rw_state_flux, _strong_transition_matrix
        double[:,:] _WORK, _eigvecs
        int[:,:] _total_obs, _graph
        #bint _return_flux, _return_states, _return_color
        str _return_obs
        #double[:] eigvals, eigvalsi
        #double[:,:] eigvecs, WORK



    # CREATE NUMPY ARRAYS
    # This is a temporary measure that fixes some segfaults, which implies I'm probably off by
    # a little bit.  Memory heavy, but whatever.
    # It breaks depending on things, so I need to root that out.  Clearly, nnz is larger than that.
    _flux = flux
    nnz = len(flux)
    lind = indices.shape[0]
    nlind = indices.shape[0]*stride
    total_fluxes = np.zeros((nfbins, nfbins), weight_dtype)
    total_obs = np.zeros((nfbins, nfbins), intc_dtype)
    transition_matrix = np.zeros((nfbins, nfbins), weight_dtype)
    strong_transition_matrix = np.zeros((nfbins, nfbins), weight_dtype)
    visited = np.zeros((nfbins), intc_dtype)
    graph = np.zeros((nfbins, nfbins+1), dtype=intc_dtype)
    rw_bin_probs = np.zeros(nfbins, weight_dtype)
    new_indices = np.zeros(((nlind)), dtype=indices.dtype)
    rw_state_flux = np.zeros((nstates, nstates), np.float64)
    state_flux = np.zeros((nstates, nstates), weight_dtype)
    eigvals = np.zeros((nfbins), np.float64)
    eigvalsi = np.zeros((nfbins), np.float64)
    eigvecs = np.zeros((nfbins, nfbins), np.float64)
    WORK = np.zeros((nfbins*4, nfbins*4), np.float64)
    rw_color_probs = np.zeros((nstates), weight_dtype)
    rw_state_probs = np.zeros((nbins), weight_dtype)


    # CREATE MEMORYVIEWS
    # These are for what we sent in...
    _rows = rows
    _cols = cols
    _obs = obs
    _flux = flux
    _ins = insert
    _nstates = nstates
    # ... these are for functions we'll be using.
    _total_fluxes = total_fluxes
    _total_obs = total_obs
    _transition_matrix = transition_matrix
    _strong_transition_matrix = strong_transition_matrix
    _nfbins = nfbins
    _stride = stride
    _obs_threshold = obs_threshold
    _indices = indices
    _new_indices = new_indices
    _rw_state_flux = rw_state_flux
    _rw_bin_probs = rw_bin_probs
    _eigvals = eigvals
    _eigvalsi = eigvalsi
    _eigvals = eigvals
    _eigvecs = eigvecs
    _WORK = WORK
    _graph = graph
    _visited = visited
    _rw_color_probs = rw_color_probs
    _rw_state_probs = rw_state_probs
    _bin_last_state_map = bin_last_state_map
    _bin_state_map = bin_state_map
    _istate = istate
    _jstate = jstate
    _return_obs = return_obs


    #NOGIL
    # Reconstruct dataset.  We're just passing the same thing back and forth between functions.
    with nogil:
        for i in range(_nfbins):
            for j in range(1, _nfbins+1):
                _graph[i, j] = _nfbins
        regenerate_subsampled_indices(_indices, _new_indices, lind, _stride)
        accumulate_fluxes(_rows, _cols, _obs, _flux, _ins, _new_indices, nnz, _transition_matrix, nlind)
        accumulate_obs(_rows, _cols, _obs, _flux, _ins, _new_indices, nnz, _total_obs, nlind)

        remove_under_obs(_transition_matrix, _total_obs, _obs_threshold, _nfbins)
        normalize(_transition_matrix, _nfbins)
        steadystate_solve(_transition_matrix, _strong_transition_matrix, _rw_bin_probs, _nfbins, _eigvals, _eigvalsi, _eigvecs, _WORK, _graph, _visited)

        for i in range(_nfbins):
            _rw_color_probs[_bin_last_state_map[i]] += _rw_bin_probs[i]
            _rw_state_probs[_bin_state_map[i]] += _rw_bin_probs[i]


        calc_state_flux(_transition_matrix, _rw_bin_probs, _bin_last_state_map, _bin_state_map, _nstates, _rw_state_flux, _nfbins)

    # This allows us to use the same function for all three types.
    # Return conditional fluxes.
    if _return_obs == b'F':
        _return_value = _rw_state_flux[_istate,_jstate]
        if isnan(_return_value) is True:
            return 0.0
        else:
            return _return_value
    # Return state probabilities.
    elif _return_obs == b'S':
        _return_value = _rw_state_probs[_istate]
        if isnan(_return_value) is True:
            return 0.0
        else:
            return _return_value
    # Return color (ensemble) probabilities
    elif _return_obs == b'C':
        _return_value = _rw_color_probs[_istate]
        if isnan(_return_value) is True:
            return 0.0
        else:
            return _return_value
    # Return the rates.
    elif _return_obs == b'R':
        if _rw_color_probs[_istate] != 0.0:
            _return_value = (_rw_state_flux[_istate,_jstate] / (_rw_color_probs[_istate] / (_rw_color_probs[_istate] + _rw_color_probs[_jstate])))
            if isnan(_return_value) is True:
                return 0.0
            else:
                return _return_value
        else:
            # We have no ensemble probability, and as such, cannot have a flux.
            return 0.0
    # Return the populations.
    elif _return_obs == b'P':
        return rw_bin_probs

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cpdef int regenerate_subsampled_indices(Ushort[:] iin, Ushort[:] iout, int ilen, int stride) nogil:

    cdef:
        int i, si

    # go over the range of all indices within iin
    for i in range(ilen):
        # Run over the length of the stride.
        for si in range(stride):
            iout[(i*stride)+si] = iin[i] + si

    return 0

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cpdef int accumulate_fluxes(int[:] hrows, int[:] hcols, int[:] hobs, weight_t[:] hflux, int[:] hins, Ushort[:] iterations, Py_ssize_t nnz, weight_t[:,:] total_fluxes, int itermax) nogil:

    cdef:
        index_t curriter, elem, iiter, ipop
        long ilem

    curriter = 0

    for iter in range(itermax):
        iiter = iterations[iter]
        for ilem in range(hins[iiter], hins[iiter+1]):
            # Not sure if this is necessary, here...
            if ilem < nnz and iiter+1 < itermax:
                total_fluxes[hrows[ilem], hcols[ilem]] += hflux[ilem]

    return 0

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cpdef int accumulate_obs(int[:] hrows, int[:] hcols, int[:] hobs, weight_t[:] hflux, int[:] hins, Ushort[:] iterations, Py_ssize_t nnz, int[:,:] total_obs, int itermax) nogil:

    cdef:
        index_t curriter, elem, iiter, ipop
        long ilem

    curriter = 0

    for iter in range(itermax):
        iiter = iterations[iter]
        for ilem in range(hins[iiter], hins[iiter+1]):
            if ilem < nnz and iiter+1 < itermax:
                total_obs[hrows[ilem], hcols[ilem]] += hobs[ilem]

    return 0

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cpdef int remove_under_obs(weight_t[:,:] flux, int[:,:] obs, int threshold, int nbins) nogil:

    cdef:
        int x, y

    for x in range(nbins):
        for y in range(nbins):
            if obs[x,y] < threshold:
                flux[x,y] = 0


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cpdef int calc_state_flux(weight_t[:, :] trans_matrix, weight_t[:] bin_probs, long[:] bin_last_state_map, Ushort[:] bin_state_map, int nstates, weight_t[:,:] state_flux, int K_shape) nogil:

    cdef:
        int i, j, ii, jj


    for i in range(K_shape):
        for j in range(K_shape):

            ii = bin_last_state_map[i]
            jj = bin_state_map[j]

            if jj != nstates:
                state_flux[ii, jj] += (trans_matrix[i, j] * bin_probs[i])

    return 0

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cpdef int steadystate_solve(weight_t[:,:] K, weight_t[:,:] K_mod, weight_t[:] bin_prob, int K_shape, double[:] eigvals, double[:] eigvalsi, double[:,:] eigvecs, double[:,:] WORK, int[:,:] graph, int[:] visited) nogil:

    cdef:
        int[:] components
        int[:,:] _graph
        double max, eigsum
        int n_components, components_assignments, largest_component, maxi, x, y, n, INFO, LWORK, i, j

        # POINTERS

        int  *_INFO, *_K_shape, *_LWORK
        double *_K_mod, *_eigvals, *_eigvecs, *_WORK, *_eigvalsi
    _K_shape = &K_shape
    _INFO = &INFO
    _LWORK = &LWORK
    INFO = 0
    LWORK = K_shape * 4
    _K_mod = &K_mod[0,0]
    _eigvals = &eigvals[0]
    _eigvecs = &eigvecs[0,0]
    _eigvalsi = &eigvalsi[0]
    _WORK = &WORK[0,0]
    _graph = graph

    for i in range(K_shape):
        if visited[i] == 0:
            visited[i] = 1
            return_strong_component(K, K_shape, _graph, i, i, visited)
    n = 0
    for i in range(K_shape):
        if graph[i, 0] >= graph[n, 0]:
            n = i
    # I suspect this may be giving us issues?
    #components = _graph[n, :K_shape+1]

    maxi = 0
    eigsum = 0.0
    # This all works!
    for x in range(K_shape):
        #i = components[x+1]
        i = graph[n, x+1]
        for y in range(K_shape):
            #j = components[y+1]
            j = graph[n, y+1]
            if i != K_shape and j != K_shape:
                K_mod[i, j] = K[i, j]
    normalize(K_mod, K_shape)
    cl.dgeev('N', 'V', _K_shape, _K_mod, _K_shape, _eigvals, _eigvalsi, _eigvecs, _K_shape, _eigvecs, _K_shape, _WORK, _LWORK, _INFO)
    for x in range(K_shape):
        if x == 0:
            max = eigvals[0]
            maxi = x
        else:
            if max < eigvals[x]:
                max = eigvals[x]
                maxi = x
    # We need to go over the whole range and pick out non K_shape elements.
    # This probably no longer needs to be done, now...
    for i in range(K_shape):
        #x = components[i+1]
        x = graph[n, i+1]
        if x != K_shape:
            #eigsum += eigvecs[maxi, components[i+1]]
            eigsum += eigvecs[maxi, x]
    for i in range(K_shape):
        #x = components[i+1]
        x = graph[n, i+1]
        if x != K_shape:
            #bin_prob[components[i+1]] = eigvecs[maxi, components[i+1]]
            #bin_prob[components[i+1]] /= eigsum
            bin_prob[x] = eigvecs[maxi, x]
            bin_prob[x] /= eigsum

    return 0

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cpdef int return_strong_component(weight_t[:,:] K, int K_shape, int[:, :] graph, int i, int z, int[:] visited) nogil:

    cdef:
        int j, y

    if graph[z, 0] == 0:
        graph[z, 0] += 1
        graph[z, 1] = i
    for j in xrange(K_shape):
        if i != j:
            if K[i, j] > 0.0:
                # Strongly connected!
                if visited[j] == 0:
                    graph[z, 0] += 1
                    y = graph[z, 0]
                    graph[z, y] = j
                    # We only want to call it when we haven't visited it before.
                    # We don't want to call, THEN modify and check.  Otherwise, we could be doing many calls.
                    visited[j] = 1
                    return_strong_component(K, K_shape, graph, j, z, visited)


    return 0
