from .blockmodel import BiologicalBlockmodel
from modelflow import Node, Input
from .bases import NetworkModelMixin


import numpy as np

class LL(NetworkModelMixin, Node):
    """
    Layered recurrent network.

    It has the additional parameters

    :param nr_exc_subpopulations: int
        Number of layers. Only the excitatory neurons are layered.

    :param reciprocity_exc: float in [0, 1]
        Excitatory reciprocity of the network.
        Note that the connectivity and the number of layers constrain the reciprocity.
    """
    nr_exc_subpopulations = Input()
    reciprocity_exc = Input()

    def run(self):
        biological_blockmodel = BiologicalBlockmodel(nr_upper_diagonals=1)
        pars = {k : getattr(self, k) for k in self._input_channel_ids}
        res = biological_blockmodel(**pars)
        self.block_matrix_ = biological_blockmodel.block_matrix_
        return res


def connectivity(nr_layers, connectivity_in_layer, connectivity_forward):
    return (nr_layers*connectivity_in_layer + (nr_layers-1)*connectivity_forward) / nr_layers ** 2


def reciprocity(nr_layers, connectivity_in_layer, connectivity_forward):
    conn = connectivity(nr_layers, connectivity_in_layer, connectivity_forward)
    if not isinstance(conn, np.ndarray):
        if conn == 0:
            return 0
    return (connectivity_in_layer ** 2 / nr_layers
            / connectivity(nr_layers, connectivity_in_layer, connectivity_forward))