import json
import random
import unittest

import scipy as sp
import theano
import theano.tensor as T
from keras.layers import Dense, SimpleRNN
from keras.models import Sequential
from keras.optimizers import SGD

from connectome.function.kerasextension import BlockRnn
from connectome.model import ER, LL
from util.testing import run_slow_test


class TestBlockRnn(unittest.TestCase):
    def setUp(self):
        self.er = ER(nr_neurons=50, inh_ratio=.5, p_exc=.5, p_inh=.5)()
        self.ll = LL(nr_neurons=50, inh_ratio=.5, p_exc=.5, p_inh=.5, nr_exc_subpopulations=3,
                     reciprocity_exc=.2)()

    def test_preprocess_input(self):
        input_dim = random.randint(1, 5)
        n_samples = random.randint(10, 20)
        n_timesteps = random.randint(10, 20)

        rnn = BlockRnn(network=self.er, input_dim=input_dim, return_sequences=True)
        Sequential().add(rnn)
        input = rnn.get_input_at(0)
        preprocessed_input = rnn.preprocess_input(input)
        f = theano.function([input], preprocessed_input, allow_input_downcast=True)
        input_shape = (n_samples, n_timesteps, input_dim)
        X = sp.zeros(input_shape)
        preprocessed_input_calculated = f(X)

        self.assertEqual(preprocessed_input_calculated.shape, (n_samples, n_timesteps, rnn.network.nr_exc))

    def test_step_output(self):
        n_samples = 13
        n_timesteps = 23
        input_dim = 1
        X = theano.shared(sp.zeros((n_samples, n_timesteps,  input_dim)))


        rnn = BlockRnn(input_dim=input_dim, network=self.er, return_sequences=True)
        Sequential().add(rnn)
        initial_states = rnn.get_initial_states(X)
        preprocessed_input = rnn.preprocess_input(X)
        reshuffled_preprocessed_input = reshuffle_preprocessed_input(preprocessed_input)
        output, new_states = rnn.step(reshuffled_preprocessed_input[0], initial_states)
        self.assertEqual(output.eval().shape, (n_samples, rnn.network.nr_exc))

    def test_step_state_update(self):
        n_samples = 13
        n_timesteps = 23
        input_dim = 1
        X = theano.shared(sp.zeros((n_samples, n_timesteps,  input_dim)))

        rnn = BlockRnn(input_dim=input_dim, network=self.ll, return_sequences=True)
        Sequential().add(rnn)  # for initialization
        initial_states = rnn.get_initial_states(X)
        preprocessed_input = rnn.preprocess_input(X)
        reshuffled_preprocessed_input = reshuffle_preprocessed_input(preprocessed_input)
        output, new_states = rnn.step(reshuffled_preprocessed_input[0], initial_states)
        for k, (new_state, block) in enumerate(zip(new_states, rnn.network.blocks)):
            try:
                self.assertEqual(new_state.eval().shape, (n_samples, block.size))
            except Exception:
                print("Problem in block", k, flush=True)
                raise

    def test_dot(self):
        rnn = BlockRnn(input_dim=1, network=self.ll, return_sequences=True)
        rnn._make_adjacency()
        mat = rnn.adjacency[-2,-2]
        inp = T.vector()
        f = theano.function([inp], mat.dot(inp), allow_input_downcast=True)
        res = f(sp.zeros(rnn.network.blocks[-2].size))
        self.assertEqual(res.shape, (rnn.network.blocks[-2].size,))

    def test_output(self):
        n_samples = 13
        n_timesteps = 23
        input_dim = 2

        rnn = BlockRnn(input_dim=input_dim, network=self.ll,  return_sequences=True)
        Sequential().add(rnn)
        input = rnn.get_input_at(0)
        output = rnn.get_output_at(0)
        print("input", input, "output", output)
        f = theano.function([input], output, allow_input_downcast=True)

        res = f(sp.zeros((n_samples, n_timesteps,  input_dim)))
        self.assertEqual(res.shape, (n_samples, n_timesteps,  rnn.network.blocks[0].size))

    def test_initial_states(self):
        n_samples = 13
        n_timesteps = 23
        input_dim = 1

        ll = LL(nr_neurons=50, inh_ratio=.5, p_exc=.3, p_inh=.5, nr_exc_subpopulations=3, reciprocity_exc=.2)()
        rnn = BlockRnn(input_dim=input_dim, network=ll,  return_sequences=True)
        X = theano.shared(sp.zeros((n_samples, n_timesteps,  input_dim)))
        initial_states = rnn.get_initial_states(X)
        self.assertEqual([s.eval().shape for s in initial_states],
                         [(n_samples, block.size) for block in ll.blocks])

    def test_compile(self):
        input_dim = 4
        n_samples = 23
        n_timesteps = 12
        nb_epoch = 1
        batch_size = 2
        X = sp.rand(*(n_samples, n_timesteps, input_dim)).astype("float32")
        y = sp.zeros(n_samples).astype(int)

        rnn = BlockRnn(input_dim=input_dim, network=self.er, return_sequences=False,
                       trainable=True)
        model = Sequential()
        model.add(rnn)
        model.add(Dense(1))
        model.compile("sgd", "binary_crossentropy")
        history = model.fit(X, y, nb_epoch=nb_epoch, batch_size=batch_size, verbose=0)

        self.assertEqual(nb_epoch, len(history.history["loss"]))

    def test_fit_weight_not_change(self):
        input_dim = 1
        n_samples = 100
        n_timesteps = 30
        nb_epoch = 3
        batch_size = 30
        X = sp.rand(*(n_samples, n_timesteps, input_dim)).astype("float32") * 5
        y = (sp.rand(n_samples) < .5).astype(int)

        rnn = BlockRnn(input_dim=input_dim, network=self.ll, return_sequences=False,
                       trainable=True)
        print(rnn.network.blocks)
        model = Sequential()
        model.add(rnn)
        model.add(Dense(1, trainable=False))
        model.compile(SGD(lr=10), "binary_crossentropy")

        e_to_e_foward_weight_before_training = rnn.adjacency[0, 1].get_value()
        e_to_e_inpool_weight_before_training = rnn.adjacency[0, 0].get_value()

        model.fit(X, y, nb_epoch=nb_epoch, batch_size=batch_size)

        e_to_e_foward_weight_after_training = rnn.adjacency[0, 1].get_value()
        e_to_e_inpool_weight_after_training = rnn.adjacency[0, 0].get_value()

        self.assertTrue((e_to_e_inpool_weight_before_training == e_to_e_inpool_weight_after_training).all())

    def test_compile_simple_rnn_validation(self):
        """
        To check consistency with Keras.
        Not an actual test.
        """
        input_dim = 4
        n_samples = 23
        n_timesteps = 12
        nb_epoch = 3
        batch_size = 2
        X = sp.rand(*(n_samples, n_timesteps, input_dim)).astype("float32")
        y = sp.zeros(n_samples).astype(int)

        rnn = SimpleRNN(55, input_dim=input_dim, trainable=False)

        model = Sequential()
        model.add(rnn)
        model.add(Dense(1))
        model.compile("sgd", "binary_crossentropy")
        history = model.fit(X, y, nb_epoch=nb_epoch, batch_size=batch_size, verbose=0)
        self.assertEqual(nb_epoch, len(history.history["loss"]))

    def test_jsonify(self):
        blockrnn = BlockRnn(network=self.er)
        config = blockrnn.get_config()
        try:
            json.dumps(config)
        except Exception as e:
            self.fail("BlockRNN not json encodable: {}".format(e))


def reshuffle_preprocessed_input(preprocessed_input):
    """
    This function is copied from K.rnn.
    Before the call to the theano scan, a reshuffle is done

    Parameters
    ----------
    preprocessed_input

    Returns
    -------

    """
    ndim = preprocessed_input.ndim
    assert ndim >= 3, 'Input should be at least 3D.'
    axes = [1, 0] + list(range(2, ndim))
    preprocessed_input = preprocessed_input.dimshuffle(axes)
    return preprocessed_input


if __name__ == "__main__":
    unittest.main()
