from collections import namedtuple

import scipy as sp

from .dynamicsyn import DynamicSYN


def make_pool_activations(dynamic_syn: DynamicSYN):
    """

    Parameters
    ----------
    dynamic_syn: DynamicSYN

    Returns
    -------
    pool_activations, t_list

    pool_activations: 2d array of shape (nr_pools+1, nr_unique_spike_times)
        each entry is the number of active neurons in the given pool at the given time
        the inh neurons are excluded from the pools


    t_list: list of length nr_unique_spike_times
         the times of the spike events
    """
    n_time_points = sp.unique(dynamic_syn.spike_monitor_.t).size
    pool_activations = sp.zeros((len(dynamic_syn.pool_sets_), n_time_points))
    t = 0
    last_t = dynamic_syn.spike_monitor_.t[0]
    t_list = [last_t]
    for t_new, neuron in zip(dynamic_syn.spike_monitor_.t, dynamic_syn.spike_monitor_.i):
        if last_t != t_new:
            t += 1
            last_t = t_new
            t_list.append(t_new)
        for pool_nr, pool in enumerate(dynamic_syn.pool_sets_[1:]):  # the first pool are the inh neurons. skip that.
            if neuron in pool:
                pool_activations[pool_nr, t] += 1
    t_list = sp.asarray(t_list)
    return pool_activations, t_list


def estimated_nr_activated_pools(pool_activations, pool_size):
    activation_threshold = .5
    max_activation_sum = 1.2
    pool_activations_cumsum = pool_activations.cumsum(axis=1)
    cum_max = pool_activations_cumsum.max(axis=0)
    last_sane_index = sp.searchsorted(cum_max, pool_size * max_activation_sum)
    sane_pool_activations = pool_activations_cumsum[:,min(last_sane_index, pool_activations_cumsum.shape[1]-1)]
    nr_activated_pools_from_max = sp.nonzero(sane_pool_activations > pool_size * activation_threshold)[0].max() + 1
    nr_activated_pools_from_sum = (sane_pool_activations > pool_size * activation_threshold).sum()
    return ActivatedPools(nr_activated_pools_from_max, nr_activated_pools_from_sum, last_sane_index)


def keep_max_pool_activities(pool_activations):
    """
    Modify in place
    """
    for k in range(pool_activations.shape[1]):
        activities = pool_activations[:,k]
        max_activity = activities.max()
        nr_maximally_active_pools = (activities == max_activity).sum()
        if nr_maximally_active_pools > 1:
            activities[:] = 0
        else:
            activities[activities != max_activity] = 0


ActivatedPools = namedtuple("ActivatedPools", "from_max from_sum last_sane_index")