# Path: stationary/utils/math_helpers.py
# def slice_dictionary(d, N, slice_index=0, slice_value=0):
# def squared_error(d1, d2):
# def squared_error_dict(d1, d2):
# def multiply_vectors(a, b):
# def dot_product(a, b):
# def normalize(x):
# def normalize_dictionary(x):
# def inc_factorial(x, n):
# def factorial(i):
# def log_inc_factorial(x,n):
# def log_factorial(i):
# def simplex_generator(N, d=2):
# def one_step_generator(d):
# def one_step_indicies_generator(d):
# def kl_divergence(p, q):
# def kl_divergence_dict(p, q):
# def q_divergence(q):
#         def d(x, y):
#         def d(x, y):
#     def d(x, y):
# def shannon_entropy(p):
# def binary_entropy(p):
# 
# Path: stationary/utils/edges.py
# def edge_func_to_edges(edge_func, states):
#     """
#     Convert an edge_func to a list of edges.
#     """
# 
#     edges = []
#     for s1 in states:
#         for s2 in states:
#             edges.append((s1, s2, edge_func(s1, s2)))
#     return edges
# 
# def states_from_edges(edges):
#     """
#     Computes the underlying set of states from the list of edges.
# 
#     Parameters
#     ----------
#     edges: list of tuples
#         Transition probabilities of the form [(source, target, transition_probability
# 
#     Returns
#     -------
#     states, set
#         The set of all the vertices of the edge list
#     """
#     states = set()
#     for (source, target, weight) in edges:
#         states.add(source)
#         states.add(target)
#     return states
# 
# def power_transitions(edges, k=20):
#     """
#     Raises a transition matrix (specified by edges) to the power `k`, returning
#     an edge_func.
#     """
# 
#     mat, all_states, enum = edges_to_matrix(edges)
#     M_k = matrix_power(mat, k)
# 
#     def edge_func(current_state, next_state):
#         return M_k[enum[current_state]][enum[next_state]]
#     return edge_func


from ..utils.math_helpers import (
    simplex_generator, one_step_indicies_generator, logsumexp,
    log_factorial, log_inc_factorial, factorial, inc_factorial)
from ..utils.edges import (
    edge_func_to_edges, states_from_edges, power_transitions)
from numpy import log, exp
from .incentives import *
    for state in simplex_generator(N, n-1):
        t = 1.
        for i in state:
            t *= inc_factorial(alpha, i) / factorial(i)
        t *= factorial(N) / inc_factorial(n * alpha, N)        
        d2[state] = t
    return d2


def log_neutral_stationary(N, alpha, n=3):
    """
    Computes the stationary distribution of the neutral landscape. This process
    is always reversible and there is an explicit formula. This function is the 
    same as neutral_stationary in log-space.

    Parameters
    ----------
    N: int
        Population size / simplex divisor
    alpha:
        Parameter defining the stationary distribution in terms of n and mu
    n: int, 3
        Simplex dimension - 1, Number of types in population

    Returns
    -------
    dictionary, stationary distribution of the process
    """

    d2 = dict()

# Path: stationary/utils/math_helpers.py
# def slice_dictionary(d, N, slice_index=0, slice_value=0):
# def squared_error(d1, d2):
# def squared_error_dict(d1, d2):
# def multiply_vectors(a, b):
# def dot_product(a, b):
# def normalize(x):
# def normalize_dictionary(x):
# def inc_factorial(x, n):
# def factorial(i):
# def log_inc_factorial(x,n):
# def log_factorial(i):
# def simplex_generator(N, d=2):
# def one_step_generator(d):
# def one_step_indicies_generator(d):
# def kl_divergence(p, q):
# def kl_divergence_dict(p, q):
# def q_divergence(q):
#         def d(x, y):
#         def d(x, y):
#     def d(x, y):
# def shannon_entropy(p):
# def binary_entropy(p):
# 
# Path: stationary/utils/edges.py
# def edge_func_to_edges(edge_func, states):
#     """
#     Convert an edge_func to a list of edges.
#     """
# 
#     edges = []
#     for s1 in states:
#         for s2 in states:
#             edges.append((s1, s2, edge_func(s1, s2)))
#     return edges
# 
# def states_from_edges(edges):
#     """
#     Computes the underlying set of states from the list of edges.
# 
#     Parameters
#     ----------
#     edges: list of tuples
#         Transition probabilities of the form [(source, target, transition_probability
# 
#     Returns
#     -------
#     states, set
#         The set of all the vertices of the edge list
#     """
#     states = set()
#     for (source, target, weight) in edges:
#         states.add(source)
#         states.add(target)
#     return states
# 
# def power_transitions(edges, k=20):
#     """
#     Raises a transition matrix (specified by edges) to the power `k`, returning
#     an edge_func.
#     """
# 
#     mat, all_states, enum = edges_to_matrix(edges)
#     M_k = matrix_power(mat, k)
# 
#     def edge_func(current_state, next_state):
#         return M_k[enum[current_state]][enum[next_state]]
#     return edge_func


# Path: stationary/processes/incentive_process.py
from ..utils.math_helpers import (
    simplex_generator, one_step_indicies_generator, logsumexp,
    log_factorial, log_inc_factorial, factorial, inc_factorial)
from ..utils.edges import (
    edge_func_to_edges, states_from_edges, power_transitions)
from numpy import log, exp
from .incentives import *

    for state in simplex_generator(N, n-1):
        t = 1.
        for i in state:
            t *= inc_factorial(alpha, i) / factorial(i)
        t *= factorial(N) / inc_factorial(n * alpha, N)        
        d2[state] = t
    return d2


def log_neutral_stationary(N, alpha, n=3):
    """
    Computes the stationary distribution of the neutral landscape. This process
    is always reversible and there is an explicit formula. This function is the 
    same as neutral_stationary in log-space.

    Parameters
    ----------
    N: int
        Population size / simplex divisor
    alpha:
        Parameter defining the stationary distribution in terms of n and mu
    n: int, 3
        Simplex dimension - 1, Number of types in population

    Returns
    -------
    dictionary, stationary distribution of the process
    """

    d2 = dict()

    for state in simplex_generator(N, n-1):