from brian2 import ms
from scipy import stats as st

from .dynamicsyn import DynamicSYN
from connectome.model import SYN
from .evaluation import make_pool_activations, estimated_nr_activated_pools, keep_max_pool_activities
from ..util import add_default_output
from ...modeltest import Task


class PropagationTask(Task):
    """
    Checks the propagation of activity along a synfire chain.

    Accepts the SYN.

    Parameters
    ----------

    run_time: brian2.Quantity
        Simulation time.

    spike_detection_pool_fraction: float in [0, 1]
        Fraction of neurons in a pool which have to spike
        to detect event as pool activation.
    """
    def __init__(self, run_time=100 * ms,
                 spike_detection_pool_fraction: float=.5,
                 **dyn_syn_pars):
        self.run_time = run_time
        self.spike_detection_pool_fraction = spike_detection_pool_fraction
        self.dyn_syn_pars = dyn_syn_pars

    def dynamic_syn_factory(self, model: SYN):
        default_dyn_syn_pars = dict(
                                 pool_delay=st.uniform(.3, 1),
                                 g_ratio_jitter=0.7,
                                 pool_jitter=st.uniform(0, .5),
                                 inh_jitter=st.uniform(.45, .1))

        default_dyn_syn_pars.update(self.dyn_syn_pars)
        dynamic_syn = DynamicSYN(static_syn=model, **default_dyn_syn_pars)
        return dynamic_syn

    @add_default_output
    def __call__(self, model: SYN) -> dict:
        dynamic_syn = self.dynamic_syn_factory(model)
        dynamic_syn.run(self.run_time)
        pool_activations, _ = make_pool_activations(dynamic_syn)
        pool_activations_original = pool_activations.copy()
        keep_max_pool_activities(pool_activations)
        activated_pools = estimated_nr_activated_pools(pool_activations, model.pool_size)
        nr_pools = len(model.pool_connector_.pools_exc)
        fraction_pools_activated = activated_pools.from_sum / nr_pools
        return {"fraction_pools_activated": fraction_pools_activated,
                "pool_activations": pool_activations_original,
                "nr_pools_spiked": activated_pools.from_sum,
                "nr_pools_spiked_from_max": activated_pools.from_max,
                "last_sane_index": activated_pools.last_sane_index,
                "nr_pools": nr_pools, "pool_size": model.pool_size}
