"""
Initialization for biological RNNs
==================================
"""


import scipy as sp
import numpy.linalg as la
from keras import backend as K
from abc import ABC, abstractmethod
import logging
import scipy.stats
from connectome.model import save_normalize
block_logger = logging.getLogger("Blocker")


class SparseWeightInit:
    def __init__(self, connectivity=.5, weight_std=.4):
        self.connectivity = connectivity
        self.weight_std = weight_std

    def __call__(self, shape, name=None):
        # parameters from misc.classification.esn.ESN
        W = sp.randn(*shape) * self.weight_std
        W -= W.mean()
        W[self.connectivity < sp.rand(*W.shape)] = 0
        return K.variable(W, name=name)


class BlockInitializerBase(ABC):
    """
    Initialize a block network.

    Subclasses have to implement the methods:

    * forward_e_to_e(self, source_size, target_size, connectivity),
    * inpool_e(self, size, connectivity),
    * inpool_ii(self, size, connectivity),
    * e_to_i(self, source_size, target_size, connectivity),
    * i_to_e(self, source_size, target_size, connectivity).
    """

    dtype = "float32"

    def init_weight(self, block, source_block_nr, target_block_nr, nr_blocks) -> sp.ndarray:
        max_block_index = nr_blocks - 1
        if not isinstance(block, sp.ndarray) and block == 0:
            return

        # diagonal blocks
        if source_block_nr == target_block_nr:
            if source_block_nr == max_block_index:
                self.inpool_ii(block)
                assert (block <= 0).all(), "I to I block has to have all entries <= 0"
            else:
                self.inpool_e(block)
                assert (block >= 0).all(), "E to E block has to have all entries >= 0"
            return

        # exc forward blocks, block 0 is output i.e. ... 3 -> 2 -> 1 -> 0
        if source_block_nr < max_block_index and source_block_nr - target_block_nr == 1:
            self.forward_e_to_e(block)
            assert (block >= 0).all(), "E to E block has to have all entries >= 0"
            return

        # exc to inh
        if source_block_nr < max_block_index and target_block_nr == max_block_index:
            self.e_to_i(block)
            assert (block >= 0).all(), "E to I block has to have all entries >= 0"
            return

        # inh to exc
        if source_block_nr == max_block_index and target_block_nr < max_block_index:
            self.i_to_e(block)
            assert (block <= 0).all(), "I to E block has to have all entries <= 0"
            return

        raise Exception("Invalid block. Cannot initialize.")


    @abstractmethod
    def forward_e_to_e(self, block) -> sp.ndarray:
        pass

    @abstractmethod
    def inpool_e(self, block) -> sp.ndarray:
        pass

    @abstractmethod
    def inpool_ii(self, block) -> sp.ndarray:
        pass

    @abstractmethod
    def e_to_i(self, block) -> sp.ndarray:
        pass

    @abstractmethod
    def i_to_e(self, block) -> sp.ndarray:
        pass


def spectral_normalization(block):
    """
    Normalize block.

    Parameters
    ----------
    block: ndarray

    Returns
    -------

    normalized_block: ndarray

    """
    block /= la.norm(block, 2)


def row_normalization(block):
    """
    Normalize block such that rows sum to 1.
    """
    save_normalize(block, 1)


def repr_scipy_stats(rv):
    try:
        dist = str(rv.dist)
        args = str(rv.args)
        kwargs = str(rv.kwds)
    except AttributeError:
        return repr(dist)
    return {"name": dist, "args": args, "kwargs": kwargs}


class LLBlockWeightInitializer(BlockInitializerBase):
    """
    Initializes through the following steps:

    1. create dense matrix with positive entries form ``distribution``,
    2. sparsify matrix randomly according to ``connectivity``,
    3. normalize using ``normalization``,
    4. multiply by corresponding strength,
    5. multiply with -1 if inhibitory.
    """
    def __init__(self, forward_e_to_e=1, inpool_e=1,
                 inpool_ii=1, e_to_i=1, i_to_e=1,
                 distribution=sp.stats.lognorm(s=5),
                 normalization=row_normalization):

        self.strength = {"forward_e_to_e": forward_e_to_e,
                         "inpool_e": inpool_e,
                         "inpool_ii": inpool_ii,
                         "e_to_i": e_to_i,
                         "i_to_e": i_to_e}
        self.distribution = distribution
        self.normalizer = normalization

    def get_config(self):
        config = {"name": self.__class__.__name__,
                  "strength": self.strength,
                  "normalization": repr(self.normalizer),
                  "distribution": repr_scipy_stats(self.distribution)}
        return config

    def _init_block(self, block, strength):
        block[block != 0] = 1
        block *= self.distribution.rvs(block.shape)
        self.normalizer(block)
        block *= strength

    def forward_e_to_e(self, block):
        self._init_block(block, self.strength["forward_e_to_e"])

    def inpool_e(self, block):
        self._init_block(block, self.strength["inpool_e"])

    def e_to_i(self, block):
        self._init_block(block, self.strength["e_to_i"])

    def inpool_ii(self, block):
        self._init_block(block, self.strength["inpool_ii"] * -1)

    def i_to_e(self, block):
        self._init_block(block, self.strength["i_to_e"] * -1)
