#cython: boundscheck=False
#cython: cdivision=True
#cython: wraparound=False
from cython.parallel cimport prange, threadid
import scipy as sp
import numpy as np


ctypedef long long longtype

class NullRecorder:
    def record(self, t, name, value):
        pass


DEFAULT_SIGMA_NOISE = 0.05

cpdef double[:] default_noise(int nr_neurons):
    """
    Default white noise

    Parameters
    ----------
    nr_neurons: int
        Number of neurons

    Returns
    -------

    Input to the network

    """
    return sp.randn(nr_neurons) * DEFAULT_SIGMA_NOISE

def make_sorn(nr_time_steps=10000, nr_exc=2000, nr_inh=200, p_exc=.2, p_inh=.4,
              eta_equal=0.001, eta_plus=0.001, eta_minus=-0.001,
              history_recorder=NullRecorder(), target_firing_rate=.1, eta_intrinsic=.1,
              pruning_threshold=None, input_generator=default_noise, synapse_add_weight=None, debug=False,
              fixed_nr_add_attempts_overwriting_p=None):
    """
    If eta_equal = eta_plut = eta amd eta_minus = -eta, then
    eta in [0.0006, 0.0015] and
    eta_intrinsic in [0.05, 0.1] work.


    Parameters
    ----------
    nr_time_steps
    nr_exc
    nr_inh
    p_exc
    p_inh

    eta_equal
    eta_plus
    eta_minus




    history_recorder
    target_firing_rate
    eta_intrinsic
    pruning_threshold
    sigma_noise_input
    synapse_add_weight
    debug
    fixed_nr_add_attempts_overwriting_p:
        If this is set so may synapse add attempts are made.

        If this number is >= 1 then the number corresponds to
        the number of add attempts.

        If this number is between 0 and 1, then a synapse is
        added with this probability in each time step.

    Returns
    -------

    """
    if not 0 <= p_exc <= 1:
        raise Exception("make_sorn: 0 < p_exc < 1 required. Got p_exc={}".format(p_exc))

    if not 0 <= p_inh <= 1:
        raise Exception("make_sorn: 0 < p_inh < 1 required. Got p_inh={}".format(p_inh))

    if synapse_add_weight is None:
        synapse_add_weight = 1 / nr_exc
    if pruning_threshold is None:
        pruning_threshold = synapse_add_weight / 10
    if pruning_threshold > synapse_add_weight:
        raise Exception("Pruning threshold ({}) has to be smaller"
                        " than the weight of newly added synapses ({}).".format(pruning_threshold, synapse_add_weight))

    nr_neurons = nr_exc + nr_inh
    adjacency_matrix = sp.concatenate(((sp.rand(nr_neurons, nr_exc) < p_exc).astype(float),
                                      (sp.rand(nr_neurons, nr_inh) < p_inh).astype(float)*-1), axis=1)

    initial_normalize_exc(adjacency_matrix, nr_exc)
    initial_normalize_inh(adjacency_matrix, nr_exc)
    set_diagonal_to_zero(adjacency_matrix)

    current_active_neurons = sp.zeros(nr_neurons, dtype=np.longlong)
    past_active_neurons = sp.zeros(nr_neurons, dtype=np.longlong)
    delta_row_sum_exc = sp.zeros(nr_exc)
    row_sum_exc = adjacency_matrix[:nr_exc, :nr_exc].sum(axis=1)
    firing_threshold = sp.ones(nr_neurons).astype(float)
    target_nr_synapses_exc = nr_exc ** 2 * p_exc

    for time_step in range(nr_time_steps):
        # Dynamical Simulation
        nr_synapses = propagate_normalize_intrinsic_plasticity(current_active_neurons, past_active_neurons,
                                                               adjacency_matrix,
                                                               delta_row_sum_exc, row_sum_exc, firing_threshold,
                                                               target_firing_rate, eta_intrinsic,
                                                               input_generator, nr_exc)

        curr_act_ind_exc = sp.flatnonzero(current_active_neurons[:nr_exc])
        past_act_ind_exc = sp.flatnonzero(past_active_neurons[:nr_exc])
        stdp_pruning(curr_act_ind_exc, past_act_ind_exc,
             adjacency_matrix[:,:nr_exc], delta_row_sum_exc, eta_equal, eta_plus, eta_minus, pruning_threshold)

        if fixed_nr_add_attempts_overwriting_p is not None:
            if fixed_nr_add_attempts_overwriting_p >= 1:
                nr_add_attempts = fixed_nr_add_attempts_overwriting_p
            else:
                nr_add_attempts = 1 if sp.rand() < fixed_nr_add_attempts_overwriting_p else 0
        else:
            nr_add_attempts = int((target_nr_synapses_exc - nr_synapses) / (1 - p_exc))

        history_recorder.record(time_step, "adjacency_matrix_before_add_synapses", adjacency_matrix)
        history_recorder.record(time_step, "nr_synapses_before_add_synapses", nr_synapses)
        if nr_add_attempts > 0:
            nr_synapses_added = add_random_synapses(adjacency_matrix[:nr_exc,:nr_exc], delta_row_sum_exc, nr_add_attempts,
                                synapse_add_weight)
        else:
            nr_synapses_added = 0

        copy_activity(current_active_neurons, past_active_neurons)

        # Recording
        history_recorder.record(time_step, "firing_threshold", firing_threshold)
        history_recorder.record(time_step, "row_sum_exc", row_sum_exc)
        history_recorder.record(time_step, "delta_row_sum_exc", delta_row_sum_exc)
        history_recorder.record(time_step, "nr_synapses_added", nr_synapses_added)
        history_recorder.record(time_step, "current_active_neurons", current_active_neurons)
        history_recorder.record(time_step, "adjacency_matrix", adjacency_matrix)

    if (adjacency_matrix[:,:nr_exc] < 0).any():
        e_sub = adjacency_matrix[:,:nr_exc]
        n_minus = sp.count_nonzero(adjacency_matrix[:,:nr_exc] < 0)
        sum_minus = e_sub[e_sub<0].sum()
        print(e_sub)
        print(nr_time_steps, nr_exc, nr_inh, p_exc, p_inh,
              eta_equal, eta_plus, eta_minus,
              history_recorder, target_firing_rate, eta_intrinsic,
              pruning_threshold, synapse_add_weight, debug,
              fixed_nr_add_attempts_overwriting_p)
        raise Exception("SORN got negative excitatory synapses. number={}, sum={}".format(n_minus,
                        sum_minus))
    if debug:
        return adjacency_matrix, row_sum_exc, delta_row_sum_exc
    else:
        return adjacency_matrix


def initial_normalize_exc(adjacency_matrix, nr_exc):
    sum_exc_in = adjacency_matrix[:,:nr_exc].sum(axis=1, keepdims=True)
    sum_exc_in[sum_exc_in <= 0] = 1   # prevent division by zero. Might be a little misleading to write <=
    adjacency_matrix[:,:nr_exc] /= sum_exc_in

def initial_normalize_inh(adjacency_matrix, nr_exc):
    sum_inh_in = adjacency_matrix[:,nr_exc:].sum(axis=1, keepdims=True) * -1
    sum_inh_in[sum_inh_in <= 0] = 1  # prevent division by zero. Might be a little misleading to write <=
    adjacency_matrix[:,nr_exc:] /= sum_inh_in

def set_diagonal_to_zero(adjacency_matrix):
    diagonal = adjacency_matrix.diagonal()
    diagonal.flags.writeable = True
    diagonal *= 0


cpdef stdp_pruning(long long[:] current_active_neurons_exc, long long[:] past_active_neurons_exc, double[:,:] adjacency_exc_exc,
           double[:] delta_row_sum_exc, double eta_equal, double eta_plus, double eta_minus, double pruning_threshold):
    """
    * STDP
    * pruning
    """
    modify_different_time(adjacency_exc_exc, current_active_neurons_exc, past_active_neurons_exc, delta_row_sum_exc,
                          eta_minus, eta_plus, pruning_threshold)
    modify_equal_time(current_active_neurons_exc, adjacency_exc_exc,
                      delta_row_sum_exc, eta_equal)


cdef void modify_equal_time(long long[:] current_active_neurons_exc, double[:,:] adjacency_exc_exc,
                      double[:] delta_row_sum_exc, double eta_equal):
    cdef:
        int current1_ind, current2_ind, current1, current2
        int len_current_non_zero = current_active_neurons_exc.size

    for current1_ind in range(len_current_non_zero):
        current1 = current_active_neurons_exc[current1_ind]
        for current2_ind in range(current1_ind+1, len_current_non_zero):
            current2 = current_active_neurons_exc[current2_ind]
            modify_equal_time_pair(current1, current2, adjacency_exc_exc, delta_row_sum_exc, eta_equal)
            modify_equal_time_pair(current2, current1, adjacency_exc_exc, delta_row_sum_exc, eta_equal)



cdef void modify_equal_time_pair(long long post, long long pre, double[:,:] adjacency_exc_exc,
                      double[:] delta_row_sum_exc, double eta_equal):
    if adjacency_exc_exc[post, pre]:
        adjacency_exc_exc[post, pre] += eta_equal
        delta_row_sum_exc[post] += eta_equal


cdef void modify_different_time(double[:,:] adjacency_exc_exc, long long[:] current_active_neurons_exc,
                                long long[:] past_active_neurons_exc, double[:] delta_row_sum_exc,
                                double eta_minus, double eta_plus, double pruning_threshold):
    cdef:
        int current, past, current_ind, past_ind
        double synapse
        int len_current_non_zero = current_active_neurons_exc.size
        int len_past_non_zero = past_active_neurons_exc.size

    for current_ind in range(len_current_non_zero):
        current = current_active_neurons_exc[current_ind]
        for past_ind in range(len_past_non_zero):
            past = past_active_neurons_exc[past_ind]
            if current != past:
                modify_causal_direction(current, past, eta_plus, adjacency_exc_exc, delta_row_sum_exc)
    for current_ind in range(len_current_non_zero):
        current = current_active_neurons_exc[current_ind]
        for past_ind in range(len_past_non_zero):
            past = past_active_neurons_exc[past_ind]
            if current != past:
                modify_anti_causal_direction(current, past, eta_minus, pruning_threshold,
                                             adjacency_exc_exc, delta_row_sum_exc)


cdef void modify_anti_causal_direction(int current, int past, double eta_minus, double pruning_threshold,
                                       double[:,:] adjacency_exc_exc, double[:] delta_row_sum_exc):
    cdef double synapse
    if adjacency_exc_exc[past, current]:
        synapse = adjacency_exc_exc[past, current]
        if synapse + eta_minus < pruning_threshold:
            adjacency_exc_exc[past, current] = 0
            delta_row_sum_exc[past] -= synapse
        else:
            adjacency_exc_exc[past, current] += eta_minus
            delta_row_sum_exc[past] += eta_minus


cdef void modify_causal_direction(int current, int past, double eta_plus, double[:,:] adjacency_exc_exc,
                             double[:] delta_row_sum_exc):
    if adjacency_exc_exc[current, past]:
        adjacency_exc_exc[current, past] += eta_plus
        delta_row_sum_exc[current] += eta_plus


cpdef copy_activity(long long[:] current_activity, long long[:] past_activity):
    cdef int k = 0
    for k in range(current_activity.size):
        past_activity[k] = current_activity[k]


cpdef propagate_normalize_intrinsic_plasticity(long long[:] current_activity,
                                               long long[:] past_activity, double[:,:] adjacency,
              double[:] delta_row_sum_exc, double[:] row_sum_exc, double[:] firing_threshold,
              double target_firing_rate, double eta_ip, input_generator, long long nr_exc):
    """
    * Reset delta_row_sum_exc entries to zero
    """
    cdef:
        int nr_neurons = adjacency.shape[0]
        int post, pre, k
        double membrane_potential = 0    # last private variable -> no inplace operations allowed
        double old_actual_row_sum_exc = 0  # last private variable -> no inplace operations allowed
        double old_estimated_row_sum_exc = 0
        long long nr_exc_synapses = 0    # reduction variable: only in place operations in prange allowed
        double[:] xi = input_generator(nr_neurons)

    for post in prange(nr_neurons, nogil=True):
        membrane_potential = 0  # last private variable -> no inplace operations allowed
        old_actual_row_sum_exc = 0

        # calculate membrane potential and synaptic normalization
        if post < nr_exc:
            old_estimated_row_sum_exc = row_sum_exc[post] + delta_row_sum_exc[post]

        for pre in range(nr_neurons):
            if adjacency[post, pre]:
                if pre < nr_exc and post < nr_exc:
                    nr_exc_synapses += 1
                    old_actual_row_sum_exc = old_actual_row_sum_exc + adjacency[post, pre]
                    if past_activity[pre]:
                        # old_actual_row_sum_exc cannot be zero since adjacency[post, pre] != 0
                        if old_estimated_row_sum_exc > 0:
                            membrane_potential = membrane_potential + adjacency[post, pre] / old_estimated_row_sum_exc
                        else:
                            membrane_potential = 0
                    if old_estimated_row_sum_exc > 0:
                        adjacency[post, pre] /= old_estimated_row_sum_exc  # (1)
                    else:
                        adjacency[post, pre] = 0
                else:
                    if past_activity[pre]:
                        membrane_potential = membrane_potential + adjacency[post, pre]
        if post < nr_exc:
            if old_estimated_row_sum_exc > 0:
                row_sum_exc[post] = old_actual_row_sum_exc / old_estimated_row_sum_exc  # do the same division as in (1)
            else:
                row_sum_exc[post] = 0
            delta_row_sum_exc[post] = 0

        activity_and_intrinsic_plasticity(membrane_potential, xi, firing_threshold, post, pre, eta_ip,
                                          target_firing_rate, current_activity)
    return nr_exc_synapses


cdef void activity_and_intrinsic_plasticity(double membrane_potential, double[:] xi, double[:] firing_threshold,
                                       int post, int pre, double eta_ip, double target_firing_rate,
                                       long long[:] current_activity) nogil:
    if membrane_potential + xi[post] > firing_threshold[post]:
        current_activity[post] = 1
        firing_threshold[post] += eta_ip * (1 - target_firing_rate)
    else:
        current_activity[post] = 0
        firing_threshold[post] -= eta_ip * target_firing_rate


def add_random_synapses(double[:,:] adjacency_exc_exc, double[:] delta_row_sum_exc,
                        long long nr_add_attempts, double synapse_weight):
    cdef:
        long nr_exc = adjacency_exc_exc.shape[0]
        long[:] neuron_pairs = sp.random.choice(nr_exc, nr_add_attempts*2)
        long k, pre, post, nr_added

    nr_added = 0
    for k in range(0, nr_add_attempts*2, 2):  # prange is slower
        pre = neuron_pairs[k]
        post = neuron_pairs[k+1]
        if pre != post and adjacency_exc_exc[post, pre] == 0:
            adjacency_exc_exc[post, pre] = synapse_weight
            delta_row_sum_exc[post] += synapse_weight
            nr_added += 1

    return nr_added


def test_threading():
    cdef:
        int n = 100
        long long[:] ids = sp.zeros(n, dtype=np.int64)
        int k

    for k in prange(n, nogil=True):
        ids[k] = threadid()

    return sorted(sp.unique(ids))
