# cython: boundscheck=False
# cython: embedsignature=True
#cython: auto_pickle=False
from libc.math cimport floor
from util.urng cimport URNG
import util.max_seed
import numpy

ctypedef double ARRAY_TYPE


cdef bint defines_square_swap_target_equal_source(int edge_1_start, int edge_1_end,
                                                  int edge_2_start, int edge_2_end, ARRAY_TYPE[:,:] adjacency_matrix):
    """
    Parameters
    ----------

    adjacency_matrix : adjacency matrix
      Link from i to j is A[i,j].
    """
    if ((edge_1_start != edge_2_start)
        and (edge_1_end != edge_2_end)
        and (edge_1_start != edge_2_end)
        and (edge_1_end != edge_2_start)
        and (not adjacency_matrix[edge_1_start,edge_2_end])
        and (not adjacency_matrix[edge_2_start,edge_1_end])):
        return True
    else:
        return False


cdef bint defines_square_swap_target_unequal_source(int edge_1_start, int edge_1_end,
                                                    int edge_2_start, int edge_2_end, ARRAY_TYPE[:,:] adjacency_matrix):
    """
    Parameters
    ----------

    adjacency_matrix: adjacency matrix
      Link from i to j is A[i,j].
    """
    if ((edge_1_start != edge_2_start)
        and  (edge_1_end != edge_2_end)
        and (not adjacency_matrix[edge_1_start,edge_2_end])
        and (not adjacency_matrix[edge_2_start,edge_1_end])):
        return True
    else:
        return False


cdef bint defines_directed_triangle_swap(int edge_1_start, int edge_1_end,
                                         int edge_2_start, int edge_2_end, ARRAY_TYPE[:,:] adjacency_matrix):
    """
    Triangle where the end edge_1_end of edge edge_1 is the start edge_2_start of edge edge_2.

    Parameters
    ----------

    adjacency_matrix : adjacency matrix.
      link from i to j is A[i,j]
    """
    if ((edge_1_end == edge_2_start)
        and adjacency_matrix[edge_2_end,edge_1_start]
        and (not adjacency_matrix[edge_1_end,edge_1_start])
        and (not adjacency_matrix[edge_2_end,edge_2_start])
        and (not adjacency_matrix[edge_1_start,edge_2_end])):
        return True
    else:
        return False

cdef void do_square_swap(int edge_1_index, int edge_1_start, int edge_1_end,
                         int  edge_2_index, int edge_2_start, int edge_2_end,
                         ARRAY_TYPE[:,:] adjacency_matrix, long long[:,::1] edge_list, long long[:,:] edge_indices):
    adjacency_matrix[edge_1_start,edge_2_end] = adjacency_matrix[edge_1_start,edge_1_end]
    adjacency_matrix[edge_2_start,edge_1_end] = adjacency_matrix[edge_2_start,edge_2_end]
    adjacency_matrix[edge_1_start,edge_1_end] = 0
    adjacency_matrix[edge_2_start,edge_2_end] = 0
    edge_list[edge_1_index,1] = edge_2_end
    edge_list[edge_2_index,1] = edge_1_end
    edge_indices[edge_1_start,edge_1_end] = -1
    edge_indices[edge_1_start,edge_2_end] = edge_1_index
    edge_indices[edge_2_start,edge_2_end] = -1
    edge_indices[edge_2_start,edge_1_end] = edge_2_index

def my_hash(long long[:,::] a):
    """
    Quick and dirty hash for small arrays.
    """
    a_ = numpy.asarray(a)
    a_.flags.writeable = False
    hash_nr = hash(a_.data)
    return hash_nr

cdef void do_directed_triangle_swap(int edge_1_index, int edge_1_start, int edge_1_end,
                                    int edge_2_index, int edge_2_start, int edge_2_end,
                                    ARRAY_TYPE[:,:] adjacency_matrix, long long[:,::1] edge_list, long long[:,:] edge_indices):
    cdef int gap_edge_start = edge_2_end
    cdef int gap_edge_end = edge_1_start
    cdef int gap_edge_index = edge_indices[gap_edge_start, gap_edge_end]
    edge_list[edge_1_index,0] = edge_1_end
    edge_list[edge_1_index,1] = edge_1_start
    edge_list[edge_2_index,0] = edge_2_end
    edge_list[edge_2_index,1] = edge_2_start
    edge_list[gap_edge_index,0] = gap_edge_end
    edge_list[gap_edge_index,1] = gap_edge_start
    adjacency_matrix[edge_1_end,edge_1_start] = adjacency_matrix[edge_1_start,edge_1_end]
    adjacency_matrix[edge_1_start,edge_1_end] = 0
    adjacency_matrix[edge_2_end,edge_2_start] = adjacency_matrix[edge_2_start,edge_2_end]
    adjacency_matrix[edge_2_start,edge_2_end] = 0
    adjacency_matrix[gap_edge_end,gap_edge_start] = adjacency_matrix[gap_edge_start,gap_edge_end]
    adjacency_matrix[gap_edge_start,gap_edge_end] = 0
    edge_indices[edge_1_start,edge_1_end] = -1
    edge_indices[edge_2_start,edge_2_end] = -1
    edge_indices[gap_edge_start,gap_edge_end] = -1
    edge_indices[edge_1_end,edge_1_start] = edge_1_index
    edge_indices[edge_2_end,edge_2_start] = edge_2_index
    edge_indices[gap_edge_end,gap_edge_start] = gap_edge_index


def adjacency_matrix_to_edge_list(adjacency_matrix_float_type):
    start, end = numpy.nonzero(adjacency_matrix_float_type)
    return numpy.hstack((start[None].T, end[None].T)).astype(numpy.int64)



def shuffle_graph(adjacency_matrix, bint target_population_equals_source_population, unsigned long long nr_draws=0,
                  bint triangle_moves=True, bint count=False):
    """
    Shuffles a graph, **not** modifying the original graph but returning a
    new shuffled graph.


    Parameters
    ----------

    adjacency_matrix: ndarray
        Adjacency matrix. Link from i to j is A[i,j].

    target_population_equals_source_population: bool
        Whether the target population equals the source population.

        * For connections from exc. to exc. and from inh. to inh. this should be true.
        * For connections from exc. to inh. and from inh. to exc. this should be false.
        * If this argument is set to ``True`` no entries on the diagonal of
          the adjacency matrix (self-loops) are allowed.
        * If this argument is set to ``False`` entries on the diagonal of the
          adjacency matrix are allowed.

    nr_draws: unsigned long long
        Number of draws of canonical moves.

    triangle_moves: bool

        * True: include triangle swaps.
        * False: only square swaps.

    count: bool

        * False: do not count occurrences of subgraphs.
        * True: count individual occurrences of subgraphs via hashes
                and store in Python dictionary.



    .. warning::

       Setting count=True is very slow.


    Returns
    -------

    (A, nr_square_moves, nr_triangle_moves): tuple
        Contains the reshuffled matrix, nr square moves, nr triangle moves.


    .. note::

        This graph shuffling draws uniformly (unbiased) from all graphs with the given degree distribution.

        .. image:: uniform_degree_preserving_shuffling.svg


        Let :math:`p(k|l)` be the transition matrix of the Markov chain
        with transitions defined by canonical
        moves [#roberts_unbiased_2012]_ on the neuronal graph.
        For simplicity focus first on the case of an excitatory population only.
        On the associated Markov graph states are represented by
        nodes and edges exist between states with strictly positive transition probability.
        Let :math:`n(k)` denote the number of neighbors of state :math:`k`.
        By finiteness of the graph there exists :math:`\\rho > 0` with :math:`\\max_k \\rho n(k) < 1`.
        The transition matrix can be chosen such that the
        transition from a given state to each neighboring state
        occurs with constant probability :math:`\\rho` for all neighboring states.
        The probability of staying in a state is
        therefore :math:`1 - \\rho n(k)` and the transition matrix

        .. math::

          p(l|k) =
          \\begin{cases}
           \\rho & k \\sim l \\\\
           0 & \\neg k \\sim l \\\\
           1 - \\rho n(k) & l = k
          \\end{cases}.


        The so defined Markov chain is ergodic.
        Aperiodicity holds as long as the Markov graph has
        at least one self-loop which is given in all practical cases.
        Irreducibility holds by connectedness of the Markov graph.
        The Markov graph is connected since the canonical moves
        allow to transition from any graph with a given degree distribution
        to any other graph with the same degree distribution.
        A stationary state exists by symmetry of the transition
        matrix :math:`p(k|l) = p(l|k)`.
        Thus, the uniform distribution :math:`\\pi_k = c` for :math:`c = \\frac{1}{ \\sum^{}_{k} 1}`
        is a stationary distribution.
        By aperiodicity and irreducibility it is also the unique stationary distribution.

        The same argument holds within the inhibitory
        population only and can be similarly applied to the excitatory
        to inhibitory and the inhibitory to excitatory submatrix.


    .. [#roberts_unbiased_2012] Roberts, E. S., and A. C. C. Coolen.
        “Unbiased Degree-Preserving Randomization of Directed Binary Networks.”
        Physical Review E 85, no. 4 (April 5, 2012): 46103. doi:10.1103/PhysRevE.85.046103.
    """
    adjacency_matrix_float_type = numpy.ascontiguousarray(adjacency_matrix).astype(float)
    if count:
        counts = {}

    cdef bint (*defines_square_swap) (int, int, int, int, ARRAY_TYPE[:,:])
    if target_population_equals_source_population:
        defines_square_swap = defines_square_swap_target_equal_source
    else:
        triangle_moves = False
        defines_square_swap = defines_square_swap_target_unequal_source

    cdef:
        URNG rng  = URNG(numpy.random.randint(0, util.max_seed.max_seed), 0., 1.)
        double TRIANGLE_ACCEPTANCE_PROBABILITY = 1./3.  # to correct for oversampling of triangles
        unsigned long long nr_triangle_swaps_performed = 0
        unsigned long long nr_square_swaps_performed = 0
        long long[:,::1] edge_list = adjacency_matrix_to_edge_list(adjacency_matrix_float_type)
        long long[:,:] edge_indices =  (numpy.ones(adjacency_matrix_float_type.shape, dtype=numpy.int64)*-1)
        unsigned long long k = 0

    for k in range(<unsigned long long>edge_list.shape[0]):
        edge_indices[edge_list[k,0], edge_list[k,1]] = k

    cdef:
        int edge_1_index, edge_2_index, edge_1_start, edge_1_end, edge_2_start, edge_2_end
        ARRAY_TYPE[:,:] adjacency_matrix_float_type_memview = adjacency_matrix_float_type

    # do the actual shuffling
    for k in range(nr_draws):
        # draw two edges
        edge_1_index = <int> floor(rng.rvs() * edge_list.shape[0])  # shape is max_index + 1, so ``floor`` covers all edges
        edge_2_index = <int> floor(rng.rvs() * edge_list.shape[0])
        edge_1_start = edge_list[edge_1_index,0]
        edge_1_end = edge_list[edge_1_index,1]
        edge_2_start = edge_list[edge_2_index,0]
        edge_2_end = edge_list[edge_2_index,1]

        # check if they define a valid move and do the move with a certain probability
        # to correct for oversampling of certain moves.
        if defines_square_swap(edge_1_start, edge_1_end, edge_2_start, edge_2_end, adjacency_matrix_float_type_memview):
            do_square_swap(edge_1_index, edge_1_start, edge_1_end,
                           edge_2_index, edge_2_start, edge_2_end,
                           adjacency_matrix_float_type_memview, edge_list, edge_indices)
            nr_square_swaps_performed += 1
        elif triangle_moves:  # triangle moves can be switched off if the source and target population are not the same
            # triangle swap one direction
            if defines_directed_triangle_swap(edge_1_start, edge_1_end,
                                              edge_2_start, edge_2_end,
                                              adjacency_matrix_float_type_memview):
                if rng.rvs() < TRIANGLE_ACCEPTANCE_PROBABILITY:
                    do_directed_triangle_swap(edge_1_index, edge_1_start, edge_1_end,
                                              edge_2_index, edge_2_start, edge_2_end,
                                              adjacency_matrix_float_type_memview, edge_list, edge_indices)
                    nr_triangle_swaps_performed += 1
            # triangle swap the other direction
            elif defines_directed_triangle_swap(edge_2_start, edge_2_end,
                                                edge_1_start, edge_1_end,
                                                adjacency_matrix_float_type_memview):
                if rng.rvs() < TRIANGLE_ACCEPTANCE_PROBABILITY:
                    do_directed_triangle_swap(edge_2_index, edge_2_start, edge_2_end,
                                              edge_1_index, edge_1_start, edge_1_end,
                                              adjacency_matrix_float_type_memview, edge_list, edge_indices)
                    nr_triangle_swaps_performed += 1
        # check for uniform builder
        if count:
            hash_nr = my_hash(adjacency_matrix_float_type_memview)
            try:
                counts[hash_nr] += 1
            except KeyError:
                counts[hash_nr] = 1
    adjacency_matrix_result = adjacency_matrix_float_type.astype(adjacency_matrix.dtype)
    if count:
        return adjacency_matrix_result, nr_square_swaps_performed, nr_triangle_swaps_performed, counts
    else:
        return adjacency_matrix_result, nr_square_swaps_performed, nr_triangle_swaps_performed
