import scipy as sp

from .network import Network
from modelflow import Node, Input
from .util import diag_to_zero


class NetworkModelMixin:
    """
    Mixin class for network models.

    Every network models derives from this class.

    :param p_inh: float in [0, 1]
        The inhibitory connectivity.
    :param p_exc: float in [0, 1]
        The excitatory connectivity.
    :param nr_neurons: integer >= 1
        The total number of neurons.
    :param inh_ratio: float im [0, 1]
        The fraction of inhibitory neurons.

    The output of this object is a :class:`Network <connectome.model.network.Network>`
    """
    p_inh = Input()
    p_exc = Input()
    nr_neurons = Input()
    inh_ratio = Input()

    @property
    def nr_inh(self):
        """
        Number of inhibitory neurons.
        """
        return int(self.nr_neurons * self.inh_ratio)

    @property
    def nr_exc(self):
        """
        Number of excitatory neurons.
        """
        return self.nr_neurons - self.nr_inh


def connect_uniform(network, pre, post, p):
    """
    Connect uniformly in place.

    Parameters
    ----------
    pre: network index
    post: nework indes
    p: float in [0,1]
    network: Network
    """
    view = network[pre, post]
    sign = 1 if pre == 'E' else -1
    view[:, :] = (sp.rand(*view.shape) < p).astype(float) * sign
    if pre == post:
        diag_to_zero(view)


def sample_weighted_network_from_probabilities(probabilities, nr_exc, nr_inh):
    W = (sp.rand(*probabilities.shape) < probabilities).astype(float)
    W *= probabilities
    W[:, nr_exc:] *= -1
    return Network(W, nr_exc=nr_exc, nr_inh=nr_inh)