"""
Biological Keras layers
=======================

"""


import logging

import scipy as sp
from theano import tensor as T
from keras import backend as K
from keras import initializations
from keras.layers import SimpleRNN, Recurrent, time_distributed_dense
from connectome.model import Network
from connectome.model.blockmodel import BlockMatrix as BlockConnectivityStructure
from connectome.model.blockmodel import block_structure as binary_block_structure
import numpy as np
from .blockstructures import BlockMatrix, BlockVector, BT
from .trainingstrategies import ExcitatoryForwardTrainer
from .constraints import Clip
from .weightinitialization import LLBlockWeightInitializer
from keras.engine import InputSpec
from keras.layers.core import Lambda

block_rnn_logger = logging.getLogger("BlockRNN")
logging.basicConfig(level=logging.DEBUG)
import numpy.linalg as la

# """
# Make block matrix with connectivity entries
#
# Parameters
# ----------
#
# r_ee: float
#    excitatory-excitatory reciprocity
#
# Returns
# -------
#
# Matrix of size nr_exc_populations +1 x nr_exc_populations +1
#
# Each positive entry corresponds to an excitatory connectivity,
# each negative entry to an inhibitory connectivity
# """


def make_block_structure(nr_exc_populations, p_exc, r_ee, p_inh):
    """
    Make block matrix with connectivity entries

    Parameters
    ----------
    nr_exc_populations: int
    p_exc: float
    r_ee: float
    p_inh: float

    Returns
    -------

    Matrix of size nr_exc_populations +1 x nr_exc_populations +1

    Each positive entry corresponds to an excitatory connectivity,
    each negative entry to an inhibitory connectivity
    """
    binary_blocks = binary_block_structure(nr_populations=nr_exc_populations, nr_upper_diagonals=1)
    block_connectivity_structure = BlockConnectivityStructure(connectivity=p_exc,
                                       reciprocity=r_ee,
                                       nr_nodes=-1, blocks=binary_blocks).make_block_connectivities()

    exc_blocks = sp.zeros((nr_exc_populations, nr_exc_populations))

    diag = exc_blocks.diagonal()
    diag.flags.writeable = True
    diag[:] = block_connectivity_structure.diagonal_connectivity

    upper_diag = exc_blocks.diagonal(1)
    upper_diag.flags.writeable = True
    upper_diag[:] = block_connectivity_structure.off_diagonal_connectivity

    total = sp.zeros((nr_exc_populations+1, nr_exc_populations+1))
    total[:nr_exc_populations,:nr_exc_populations] = exc_blocks

    # order is important p_ii has to be overwritten
    total[-1,:] = p_exc
    total[:,-1] = -p_inh
    return total


class MySimpleRnn(SimpleRNN):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.alpha = .4
        self.nr_output_steps = 7

    def step(self, inputs, states):
        prev_output = states[0]
        a = self.alpha
        output = (1 - a) * prev_output + a * T.nnet.relu(K.dot(prev_output, self.U) + inputs)
        return output, [output]

    def set_constraints(self, W_constraint, U_constraint, b_constraint):
        self.constraints = [W_constraint, U_constraint,  b_constraint]

    def get_output(self, train=False):
        output = super().get_output(train)
        if self.return_sequences:
            return output[:, :, -self.nr_output_steps:]
        else:
            return output


def all_positive_init(shape, name="", filling=.5):
    input_dim, input_layer_size = shape
    scale = 1 / (input_dim*input_layer_size)
    initial_value = np.random.uniform(low=0, high=scale, size=shape)
    initial_value[sp.rand(*shape) > filling] = 0
    return K.variable(initial_value, name=name)


class BlockRnn(Recurrent):
    """
    Block RNN model with layer structure.

    Parameters
    ----------

    trainblocks: function
        Receives the block adjacency as input
        and returns an iterable of trainable weights.

    output_all_activities: bool
        If True, output activity of all neurons of all layers,
        including excitatory and inhibitory.

        If False, only output activity of last exc layer.

    alpha: float, 0 < alpha < 1
        In the firing rate model :math:`\\tau \\dot u = - u +  F(Wu)`,
        :math:`\\alpha = \\frac{\Delta}{\\tau}` where :math:`\\Delta`
        is the time discretization.

        The smaller :math:`\\alpha`, the larger the time constant.

    trainer: Trainer
        Implements a training strategy, e.g. an ExcitatoryForwardTrainer.
    """
    INPUT_BLOCK_NR = -2
    OUTPUT_BLOCK_NR = 0

    def __init__(self, *, network: Network,
                 activation=BT.nnet.relu,
                 trainer=ExcitatoryForwardTrainer(),
                 block_weight_init=LLBlockWeightInitializer(),
                 W_in_init=all_positive_init,
                 output_all_activities=False,
                 alpha=.1,
                 **kwargs):
        block_rnn_logger.debug("Init BlockRNN")
        self.network = network
        self.output_all_activities = output_all_activities
        if output_all_activities:
            self.output_dim = self.network.nr_neurons
        else:
            self.output_dim = self.network.blocks[self.OUTPUT_BLOCK_NR].size
        self.trainer = trainer
        self.init_W_in = initializations.get(W_in_init) if isinstance(W_in_init, str) else W_in_init
        self.activation = activation
        self.block_weight_init = block_weight_init
        self.alpha = alpha
        super().__init__(**kwargs)

    def get_config(self):
        base_config = super().get_config()
        config = {"nr_neurons": self.network.nr_neurons,
                  "output_all_activities": self.output_all_activities,
                  "alpha": self.alpha,
                  "network": str(self.network),
                  "trainer": self.trainer.get_config(),
                  "block_init": self.block_weight_init.get_config(),
                  "activation": str(self.activation)}
        base_config.update(config)
        return base_config

    def __repr__(self):
        values = {"clsname": self.__class__.__name__,
                  "network": self.network}
        return "<{clsname} network={network}>".format(**values)

    def build(self, input_shape):
        block_rnn_logger.debug("Build, input_shape={}".format(input_shape))
        self.input_dim = input_shape[2]
        self.input_spec = [InputSpec(shape=input_shape)]
        self._make_input_weights()
        self._make_adjacency()
        self._set_trainable_weights()
        self._set_non_trainable_weights()
        self._set_constraints()

    def _make_input_weights(self):
        input_layer_size = self.network.blocks[self.INPUT_BLOCK_NR].size
        self.W = self.init_W_in((self.input_dim, input_layer_size), name='W')
        self.b = K.zeros((input_layer_size,), name='b')

    def _set_trainable_weights(self):
        self.trainable_weights = self.trainer.get_trainable_weights(self.W, self.b, self.adjacency)

    def _set_non_trainable_weights(self):
        self.non_trainable_weights = self.trainer.get_non_trainable_weights(self.W, self.b, self.adjacency)

    def _set_constraints(self):
        self.constraints = {c: Clip.from_signed_matrix(c.get_value()) for c in self.trainable_weights}

    def preprocess_input(self, x, train=False):
        """
        Returns
        -------

        b + X @ w:
         shape = (n_samples, n_timesteps, exc_population_size).
        """
        block_rnn_logger.debug("Preprocess input")
        input_shape = self.input_spec[0].shape
        input_dim = input_shape[2]
        timesteps = input_shape[1]
        dropout = None
        preprocessed_input = time_distributed_dense(x, self.W, self.b, dropout,
                                      input_dim, self.network.blocks[self.INPUT_BLOCK_NR].size, timesteps)
        return preprocessed_input

    def _make_adjacency(self):
        """
        Make adjacency matrix.
        """
        block_rnn_logger.debug("Make states")
        matrix = []
        for target_ix, target_size in enumerate(self.network.blocks):
            row = []
            for source_ix, source_size in enumerate(self.network.blocks):
                block_view = self.network.get_block(target_ix=target_ix, source_ix=source_ix)
                if block_view.sum():
                    block = block_view.copy()
                    self.block_weight_init.init_weight(block, source_ix, target_ix, len(self.network.blocks))
                    block_theano = K.variable(block.astype("float32"),
                               name="[{}<-{}:shape={}]".format(target_ix, source_ix, block_view.shape))
                else:
                    block_theano = 0
                row.append(block_theano)
            matrix.append(row)

        self.adjacency = BlockMatrix(matrix)
        block_rnn_logger.debug("adjacency {}".format(self.adjacency))

    def step(self, input, states):
        """

        Parameters
        ----------
        input: shape (n_samples, exc_subpopulation_size)
        states: list
            Shape [(n_samples, exc_subpopulation_size)] * nr_exc_subpopulations +  [(n_samples, inh_subpopulation_size)]
            same shape as returned by get_initial_states.

        Returns
        -------

        output, new_states

        output: tensor, shape (n_samples, exc_subpopulation_size).
        new_states: same shape as `states`.
        """
        states = BlockVector(states)
        block_rnn_logger.debug("step states, ids={}".format([id(s) for s in states]))
        block_rnn_logger.debug("step: states={}".format(states))

        input_vector = BlockVector(0 for _ in states)
        input_vector[self.INPUT_BLOCK_NR] = input

        new_states = ((1 - self.alpha) * states
                       + self.alpha * self.activation((self.adjacency.dot(states.T)).T + input_vector)
                      )

        if self.output_all_activities:
            output = K.concatenate(list(new_states))
        else:
            output = new_states[self.OUTPUT_BLOCK_NR]
        return output, list(new_states)

    def get_initial_states(self, X):
        """
        Initial neuron activations. Here: All zero.

        Parameters
        ----------
        X: input data of dimension (n_samples, n_timesteps,  input_dim).

        Returns
        -------

        initial_states: list of zero tensors of shape
        [(n_samples, exc_subpopulation_size)] * nr_exc_subpopulations +  [(n_samples, inh_subpopulation_size)]
        """
        block_rnn_logger.debug("Get initial state")
        block_rnn_logger.debug("X shape={}".format(X.shape))
        block_rnn_logger.debug("input dim={}".format(self.input_dim))

        # the following variable names indicate the dimension
        # build an all-zero tensor of shape (samples, subpopulation_size)
        samples_timesteps_input = K.zeros_like(X)  # (samples, timesteps, input_dim)
        samples_input = K.sum(samples_timesteps_input, axis=1)  # (samples, input_dim)

        # build a list of all-zero tensor of shape [(samples, block.size) for block in <excitatory blocks>]
        initial_states_exc = []
        for block in self.network.blocks[:-1]:
            input_blocksize = K.zeros((self.input_dim, block.size))
            samples_blocksize = K.dot(samples_input, input_blocksize)  # (samples, block.size)
            initial_states_exc.append(samples_blocksize)

        # build an all-zero tensor of shape (samples, inh_subpopulation_size)
        input_inh = K.zeros((self.input_dim, self.network.nr_inh))
        samples_inh = K.dot(samples_input, input_inh)  # (samples, inh_subpopulation_size)

        initial_states = initial_states_exc + [samples_inh]
        return initial_states


def slice_last(nr_time_steps):
    return Lambda(lambda x: x[:,-nr_time_steps:,:])
