from modelflow import Input, Node
from .bases import connect_uniform, NetworkModelMixin, Network
import scipy as sp


def _nr_pools(p_exc, pool_size, nr_exc):
    return int(sp.around(sp.log(1 - p_exc) / sp.log(1 - pool_size**2 / nr_exc**2)))


class PoolConnector(Node):
    use_fixed_nr_pools = False
    p_exc = Input()
    pool_size = Input()
    network = Input()

    def run(self):
        self.pools_exc = []
        self.pools_inh = []
        presynaptic_neurons_exc, _ = self._create_random_pool_indices()
        if self.use_fixed_nr_pools:
            nr_pools = self.use_fixed_nr_pools
        else:
            nr_pools = _nr_pools(self.p_exc, self.pool_size, self.network.nr_exc)
        for _ in range(nr_pools):
            postsynaptic_neurons_exc, postsynaptic_neurons_inh = self._create_random_pool_indices()
            self._connect_neuron_sets(presynaptic_neurons_exc, postsynaptic_neurons_exc)
            self._connect_neuron_sets(presynaptic_neurons_exc, postsynaptic_neurons_inh)
            presynaptic_neurons_exc = postsynaptic_neurons_exc
        return self.network

    def _create_random_pool_indices(self):
        pool_size_inh = self.inh_pool_size()
        exc_pool_indices = sp.random.choice(self.network.nr_exc, self.pool_size, replace=False)
        inh_pool_indices = sp.random.choice(self.network.nr_inh, pool_size_inh, replace=False) + self.network.nr_exc
        self.pools_exc.append(exc_pool_indices)
        self.pools_inh.append(inh_pool_indices)
        return exc_pool_indices, inh_pool_indices

    def inh_pool_size(self):
        pool_size_inh = max(int(self.pool_size * (self.network.nr_inh / self.network.nr_exc)), 1)
        return pool_size_inh

    def _connect_neuron_sets(self, pre_neurons, post_neurons):
        """
        Connect pre and post population all to all.

        Parameters
        ----------

        pre_neurons: ndarray
            list of presynaptic neurons

        post_neurons: ndarray
            list of postsynaptic neurons

        """
        W = self.network.adjacency_matrix
        W[post_neurons.reshape(-1, 1), pre_neurons.reshape(1, -1)] = 1


class SYN(NetworkModelMixin, Node):
    """
    Recurrently embedded synfire chain.

    It has the additional parameters

    :param pool_size: int
        Size of the excitatory pools.
    """
    pool_size = Input()

    def run(self):
        network = Network(nr_exc=self.nr_exc, nr_inh=self.nr_inh)
        self.pool_connector_ = PoolConnector(p_exc=self.p_exc,
                                pool_size=self.pool_size,
                                network=network)
        network = self.pool_connector_()
        connect_uniform(network, "I", "I", self.p_inh)
        connect_uniform(network, "I", "E", self.p_inh)
        network.diag_to_zero()
        return network


