# -*- coding: utf-8 -*-
#!python
#cython: boundscheck=False
#cython: cdivision=True
#cython: embedsignature=True

"""
Created on Mon May  26 8:53:00 2014

@author: emmanuel


Algorithms similar to Till Kretschmar's master thesis (2013).


This module implement the weak edge path builder algorithms to detect
layer structure.
"""
import itertools as it

import numpy as np
import scipy as sp

from libc.stdlib cimport malloc, free
from util.urng cimport URNG
from util.max_seed import max_seed
from libc.time cimport time
ctypedef unsigned long long CounterType
from connectome.pes.strong_path_enumeration_sampling import extension_ratio_from_nr_paths_sampled


def graph_to_multi_graph_list(arr):
    """
    Parameters
    ----------
        arr : array
            Directed adjacency matrix
            
    Returns
    -------
        edges :
            edges[k][l][m], if it exists means that vertex k
            is connected to vertex l. edges[k][l][m] == 1 if
            the connection is in direction k->l and -1
            it the direction is l->k
    """
    # construct list of edges
    N = arr.shape[0]
    graph = [ [] for k in range(N) ]
    edges = []
    tmp = 0
    for k, l in it.product(range(N), repeat=2):
        if k != l and arr[k,l]:
            graph[k].append( (l, -1, tmp) )
            graph[l].append( (k, 1, tmp ) )
            edges.append( (l, k) )
            tmp += 1
    return graph, edges


ctypedef struct graphstruct:
    int *** graph
    int nr_nodes
    int *nr_edges

cdef graphstruct graph_to_c_graph(graph):
    """
    Convert python adjacency list to c graph_struct
    with pointer arithmetics
    """
    cdef graphstruct cgs
    cgs.graph = <int ***>malloc(  len(graph) * sizeof(int**))
    cgs.nr_nodes = len(graph)
    cgs.nr_edges = <int *>malloc(cgs.nr_nodes * sizeof(int))
    for k in range(len(graph)):
        cgs.graph[k] = <int **>malloc(len(graph[k]) * sizeof(int*))
        cgs.nr_edges[k] = len(graph[k])
        for l in range(cgs.nr_edges[k]):
            cgs.graph[k][l] = <int *>malloc(3 * sizeof(int))
            for m in range(3):
                cgs.graph[k][l][m] = graph[k][l][m]
    return cgs

cdef void free_graphstruct(graphstruct cgs):
    """
    Free the memory of the graphstruct
    """
    for k in range(cgs.nr_nodes):
        for l in range(cgs.nr_edges[k]):
            free(cgs.graph[k][l])
        free(cgs.graph[k])
    free(cgs.graph)
    free(cgs.nr_edges)
    


cdef int vertex_in_path(int[::1] path, int vertex, int current_length):
    """
    Returns
    -------
    
    1 : if ``vertex`` is in the ``path`` of length ``current_length``
    
    0 : otherwise
    
    """
    cdef int k
    for k in range(current_length):
        if path[k] == vertex:
            return 1
    return 0


cdef int cycle_type(int start_vertex, int end_vertex, int counter,
                    int[:,::1] adjacency_matrix,int current_length):
    """
    Checks if a path is actually a cycle and if so, if it is
    a balanced or unbalanced cycles


    Parameters
    ----------
    
    start_vertex : int
        start vertex of tha path
        
    end_vertex : int
        the end vertex of tha path. Togehter, start_vertex and
        end_vertex form the two "danling" ends of the path
        
    counter : int
        (# edges directed from start -> end )
        minus
        (# edges directes from end -> start)
    
    adjacency_matrix : ndarray
        directed adjacency matrix of the graph
    
    current_length:
        length of the path.
    
    
    Returns
    -------

    0 : no cycle
    1 : balanced cycle 
    2 : unbalanced cycle
    """
    if (adjacency_matrix[start_vertex,end_vertex] == 0 
        and adjacency_matrix[end_vertex,start_vertex] == 0):
        return 0 #no cycle
    # case: note even a cycle => therefore not an unbalanced cycle
    elif current_length == 2:
        if (adjacency_matrix[start_vertex,end_vertex] != 0 
            and adjacency_matrix[end_vertex,start_vertex] != 0):
            return 1 # balanced cycle
        else:
            return 0 #no cycle
    # at least one edge between start_vertex and end_vertex exists
    # check for edge: end -> start
    elif adjacency_matrix[start_vertex,end_vertex] != 0 and counter != -1:
        return 2 # unbalanced
     # check for edge: start -> end
    elif adjacency_matrix[end_vertex,start_vertex] != 0 and counter != 1:
        return 2 #unbalanced
    else:
        return 1 #balanced
    


cdef void extend_path(int start_node, int end_node, int extend_both_directions,
                int[::1] path_node_set, int counter, int initial_edge_number,
                graphstruct *graph, int max_length, int current_length,
                int[:,::1] adjacency_matrix, CounterType[:,::1] result_list,
                double[::1] probabilities,
                URNG rng):
    """
    Parameters
    ----------
    
    start_node: int
        start node of the path
        
    end_node : int
        end node of the path
    
    extend_both_directions : boolean
        whether to extend in one or two directions
    
    path_node_set : int array
        this is **not** ordered. really just the set of nodes
        
    counter : int
        count the balance of edges, i.e. whether traversed
        in pos. or neg. direction
        
    initial_edge_number : int
        the initial seed of the path
        this is needed b/c only edges with higher label than the
        initial one can be added to the path
        
    graph : graphstruct
        reference to the graph. essentially the adjacency list
        
    adjacency_matrix : ndarray
        directed adjacency matrix of the graph.
        this is just another representation of ``graph``
        
    result_list : ndarray
         of the form
                no cycles, balanced cycles, unbalanced cycles
         len 0    0           0                    0
         len 1    0           0                    0
         len 2    x           y                    z
         ...
         
         stores how many paths were no cycles at al, balanced or unbalanced cycles
         
    probabilities : ndarray of floats
        probabilities[k] is the probability with which to extend a path
        of length k further
            

    """
    cdef:
        int next_vertex
        int direction
        int edge_number
        int k
        int cycletype
        
    cycletype = cycle_type(start_node, end_node, counter, adjacency_matrix, current_length)
    result_list[current_length-1,cycletype] += 1

    if current_length < max_length:
        # extend at the end 
        if extend_both_directions:
            for k in range(graph.nr_edges[end_node]):
                next_vertex = graph.graph[end_node][k][0]
                direction = graph.graph[end_node][k][1]
                edge_number = graph.graph[end_node][k][2]
                if (edge_number > initial_edge_number
                    and not vertex_in_path(path_node_set, next_vertex, current_length)):
                    if rng.rvs() < probabilities[current_length]:
                        path_node_set[current_length] = next_vertex
                        extend_path(start_node, next_vertex, 
                                    1,
                                    path_node_set,
                                    counter +  direction, initial_edge_number,
                                     graph, max_length, current_length + 1, adjacency_matrix,
                                     result_list, probabilities, rng)

        # then extend at the start if necessary
        for k in range(graph.nr_edges[start_node]):
            next_vertex = graph.graph[start_node][k][0]
            direction = graph.graph[start_node][k][1]
            edge_number = graph.graph[start_node][k][2]
            if (edge_number > initial_edge_number
                and not vertex_in_path(path_node_set, next_vertex, current_length)):
                if rng.rvs() < probabilities[current_length]:
                    path_node_set[current_length] = next_vertex
                    extend_path(next_vertex, end_node, 
                                0,
                                path_node_set,
                                counter -  direction, initial_edge_number,
                                 graph, max_length, current_length + 1, adjacency_matrix,
                                 result_list, probabilities, rng)







def ideal_probabilities_from_branching(adjacency_matrix, max_length, ratio_paths_to_nodes):
    """
    Calculate the probabilities to sample a constant number of expected paths
    of each length heuristically estimating how fast the graph branches
    
    Parameters
    ----------
    
    adjacency_matrix : ndarray
        adjacency matrix of the directed graph
        
    max_length : int
        sample  paths up to length max_length
        
    Returns
    -------
    
    probabilities : ndarray
        numpy array of probabilities
    """
    MAX_LENGTH_FOR_TEST_SAMPLING = 10
    NR_REPETITIONS_FOR_TEST_SAMPLING = 5
    
    
    
    W = adjacency_matrix.copy()
    nr_nodes = W.shape[0]
    for k in range(W.shape[0]):
        W[k,k] = 0
    nr_edges = sp.count_nonzero(W)
    nr_paths_desired = float(nr_nodes)
    

    probabilities = sp.zeros(max_length, dtype=np.float_)
    
    # from some reason, matrix multiplication with integers is very slow
    W_thresh = ( W != 0 ).astype(float)
    E = W_thresh + W_thresh.T
    AA = sp.dot(E, E)
    for k in range(AA.shape[0]):
        AA[k,k] = 0
        
    nr_paths_with_two_edges = float(AA.sum()) / 2.
    extension_probability = float(nr_edges) / float(nr_paths_with_two_edges)
    
    if nr_edges == 0:
        return probabilities
    else:    
        probabilities[0] = -1 # this should never be used
        probabilities[1] = float(nr_paths_desired) / float(nr_edges) 
        probabilities[2:max_length] = extension_probability
            
        test_sampling_length = min(MAX_LENGTH_FOR_TEST_SAMPLING, W.shape[0])
        results = _run_sampling(adjacency_matrix,
                                MAX_LENGTH_FOR_TEST_SAMPLING,
                                probabilities,
                                NR_REPETITIONS_FOR_TEST_SAMPLING)
        sums = results.sum(axis=1)
        
        probabilities[2:max_length] *= extension_ratio_from_nr_paths_sampled(sums)
        probabilities *= ratio_paths_to_nodes**(1. / max_length)
        return probabilities



            
def weak_edge_pes(arr, max_length_py=None, probabilities_py=None, repeat=1, ratio_paths_to_nodes=1):
    """
    Sample weak edge paths uniformly.

    Parameters
    ----------
    
    arr: 2d array
        Directed adjacency matrix.
        
    max_length: int
        Sample paths up to length ``max_length``.
        
    probabilities_py: 1d array
        Probabilities[k] is the probability with which to extend a path
        of length k further.
        
    repeat: int
            Number of times to repeat the building process.
        
        
    Returns
    -------
    
    result_list: array
        List of counts of how many paths, balanced/unbalanced cycles
        of each length were found,

            * result_list[k,0] : no cycles
            * result_list[k,1] : balanced cycles
            * result_list[k,2] : unbalanced cycles

        Here, k is the number of edges in the path.
        Note that for k=0 all entries are always 0.

    """

    arr = np.ascontiguousarray((arr != 0).astype(np.intc), dtype=np.intc)
    max_length_py = int(max_length_py)
    cdef int max_length = (len(arr.shape[0]) if max_length_py is None
                                     else min(max_length_py,arr.shape[0]))
    
    
    # probabilities[k] : prob of extending path with k nodes further
    cdef double[::1] probabilities = (probabilities_py.astype(np.float_) 
                                         if probabilities_py is not None else
                            ideal_probabilities_from_branching(arr, max_length, ratio_paths_to_nodes) )

    return _run_sampling(arr, max_length, probabilities, repeat)



cdef _run_sampling( arr, int max_length,
                   double[::1] probabilities, int repeat):
    """
    Parameters
    ----------
    
    arr : ndarray
        directed adjacency matrix
        
    max_length : int
        sample pahts up to length ``max_length``
        
    probabilities_py: ndarray
        probabilities[k] is the probability with which to extend a path
        of length k further
        
    repeat: int
            nr of times to repeat the builder
        
        
    Returns
    -------
    
    result_list: L ndarray
        list of counts of how many paths, balanced/unbalanced cycles
        of each length were found
    """
    arr = (arr != 0).astype(np.intc)
    py_graph, edges = graph_to_multi_graph_list(arr)
    cdef graphstruct graph = graph_to_c_graph(py_graph)
    cdef URNG rng = URNG(sp.random.randint(max_seed), 0., 1.)

    # of the form
    #        no cycles, balanced cycles, unbalanced cycles
    # len 0    0           0                    0
    # len 1    0           0                    0
    # len 2    x           y                    z
    # ...
    #
    cdef CounterType[:,::1] result_list = sp.zeros( (max_length, 3),
                                                             dtype=np.uint64 )  
    cdef int[::1] path_node_set = sp.zeros(max_length, dtype=np.intc)
    cdef int k
    for k in range(repeat):
        for edge_number, edge in enumerate(edges):
            if rng.rvs() < probabilities[1]:
                path_node_set[0] = edge[0]
                path_node_set[1] = edge[1]
                extend_path(edge[0], edge[1],
                            1, # extend in both directions
                            path_node_set,
                            1, # balance counter
                            edge_number,
                            &graph, max_length,
                            2,  # current length
                            arr, result_list, probabilities, rng)
                            
    free_graphstruct(graph)
    return sp.asarray(result_list)