import os

import scipy as sp
from scipy import stats as st
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder
import sqlite3
from .texture_data import TextureData
from .util import add_default_output
from ..modeltest import Task
from ...model import LL, ER, ERFEVER, EXP, SYN, SORN, API   # to be used int he eval part
from sklearn.metrics import confusion_matrix
from functools import lru_cache
from connectome.model import Network

def default_learning_rate(nr_exc_layers):
    """
    For non layered networks, with ADAM a learning rate of 0.1 and 0.01 do not
    make much of a difference.
    """
    return 0.01 if nr_exc_layers == 1 else 0.001


def to_keras(X):
    return sp.ascontiguousarray(sp.transpose(X[sp.newaxis].transpose(0, 2, 1)))


class TextureTask(Task):
    """
    Texture discrimination task.

    Can be run with any network model.

    The reported accuracy is the accuracy of each output time step considered as individual test point.

    Parameters
    ----------

    nb_epoch: int
        Number epochs for training.

    nr_samples: int
        Total number of samples for the task.

    nr_classes: int
        Number of texture classes sampled from natural images.

    length: int
        Length of a texture sequence.

    verbose: int
        How much information to print during training.

    learning_rate: callable
        Learning_rate(nr_exc_layers) should return the learning rate
        for a model with nr_exc_layers excitatory layers.
    """
    EARLY_STOPPING_STEPS = 1000

    def __init__(self, nb_epoch=6000, nr_samples=10000, nr_classes=7, length=500,
                 verbose=0, learning_rate=default_learning_rate, slice_last=250):
        self.nb_epoch = nb_epoch
        self.nr_samples = nr_samples
        self.nr_classes = nr_classes
        self.length = length
        self.verbose = verbose
        self.learning_rate = learning_rate
        self.slice_last = slice_last

    def model_factory(self, network: Network, output_all_activities=False):
        # import here, not at module level to only initialize theano when needed
        from connectome.function.kerasextension import BlockRnn, slice_last
        from connectome.function.kerasextension.weightinitialization import LLBlockWeightInitializer
        from keras.models import Sequential
        from keras.layers import TimeDistributed, Dense
        from connectome.function.kerasextension import Adam
        block_init = LLBlockWeightInitializer(distribution=st.lognorm(s=5))
        rnn = BlockRnn(input_dim=1,
                       network=network,
                       alpha=.2,
                       output_all_activities=output_all_activities,
                       return_sequences=True,
                       block_weight_init=block_init)

        keras_model = Sequential()
        keras_model.add(rnn)
        keras_model.add(slice_last(self.slice_last))
        dense = TimeDistributed(Dense(self.nr_classes, activation="softmax"))
        keras_model.add(dense)
        keras_model.compile(Adam(clipnorm=0.5, lr=self.learning_rate(network.blocks.nr_exc)),
                            'categorical_crossentropy',
                            metrics=["accuracy"])
        return keras_model

    @add_default_output
    def __call__(self, model) -> dict:
        network = model()
        keras_model = self.model_factory(network)
        # import here, not at module level to only initialize theano when needed
        from keras.callbacks import ModelCheckpoint, EarlyStopping
        from kerasvis import DBLogger

        checkpoint_dir = self.make_checkpoint_dir()

        db_logger = DBLogger(comment=repr(model))
        checkpoint = ModelCheckpoint(os.path.join(checkpoint_dir, str(db_logger.id) + ".h5"))


        X, y = self.get_data()
        # test_size=0 only permutes the ordering of the training data
        # X_test and y_test are NOT used
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=.15)


        ### Save training and test data
        data_set_dir = os.path.join(checkpoint_dir, "{}-dataset.npz".format(db_logger.id))

        sp.savez_compressed(os.path.join(data_set_dir), X_train=X_train, X_test=X_test, y_train=y_train, y_test=y_test)


        ### Run the model
        early_stopping = EarlyStopping(monitor="val_loss", patience=self.EARLY_STOPPING_STEPS, mode="min")
        history = keras_model.fit(*self.to_keras_shape(X_train, y_train),
                                  nb_epoch=self.nb_epoch, validation_split=0.1,
                                  verbose=self.verbose, batch_size=128, callbacks=[db_logger, checkpoint, early_stopping])
        evaluation = keras_model.evaluate(*self.to_keras_shape(X_test, y_test))
        return {"validation_accuracy": history.history["val_acc"][-1],
                "accuracy": history.history["acc"][-1],
                "id": db_logger.id,
                "test_evaluation": evaluation,
                "text_evaluation_metrics": keras_model.metrics_names}

    def make_checkpoint_dir(self):
        checkpoint_dir = os.path.join(os.path.expanduser("~"), "model_checkpoints")
        try:
            os.makedirs(checkpoint_dir)
        except FileExistsError:
            pass
        return checkpoint_dir

    def to_keras_shape(self, X_train, y_train):
        X_train_keras = X_train[:, :, None]
        y_train_keras = sp.ones((y_train.shape[0], min(self.slice_last, X_train.shape[1]), self.nr_classes))
        y_train_keras *= y_train[:, None, :]
        return X_train_keras, y_train_keras

    def get_data(self):
        """

        Returns
        -------
        X, y: tuple of arrays
            Data, one-hot-encoded target.
        """
        data = TextureData()
        X, y = data.get_data(self.nr_samples // self.nr_classes,
                             self.nr_classes,
                             self.length,
                             self.length // 10)
        one_hot = OneHotEncoder(self.nr_classes)
        y = one_hot.fit_transform(y[sp.newaxis].T).toarray()
        return X, y


class TextureModelAnalyser:
    def __init__(self, keras_log_db: str, checkpoint_dir: str, id: int):
        self.db = sqlite3.connect(keras_log_db)
        self.checkpoint_dir = checkpoint_dir
        self.id = id
        self.model = self.load_model()

    @property
    def weights_path(self):
        return os.path.join(self.checkpoint_dir, "{}.h5".format(self.id))

    @property
    def dataset_path(self):
        return os.path.join(self.checkpoint_dir, "{}-dataset.npz".format(self.id))

    def load_model(self, output_all_activities=False):
        comment = self.db.execute("select comment from run where id=?", (self.id,)).fetchone()[0]
        self.comment_ = comment

        M = eval(comment)
        net = M()
        self.nr_layers_ = len(net.blocks) - 1

        tt = TextureTask(slice_last=10000000000000)
        model = tt.model_factory(net, output_all_activities=output_all_activities)
        model.load_weights(self.weights_path)
        self.net_ = net
        return model

    def get_test_data(self, first_n_samples=-1):
        """
        Test data

        Returns
        -------
        X, y:
            X: shape: (n_test_samples, sequence_length)
            y: shape: (n_test_samples, n_classes)
                y is one hot encoded
        """
        data = sp.load(self.dataset_path)
        X = data["X_test"]
        y = data["y_test"]
        X = X[:first_n_samples, :]
        y = y[:first_n_samples]
        return X, y

    def get_pool_traces(self, first_n_samples=10):
        X, y = self.get_test_data(first_n_samples)
        pool_traces = self._run_pool(X)
        return pool_traces, y

    def get_pool_traces_single_input(self, x):
        return self._run_pool(x[None,:])

    def _run_pool(self, X):
        import keras.backend as K
        model = self.load_model(output_all_activities=True)
        f = K.function([model.layers[0].input], model.layers[0].output)
        pool_traces = f([X[:, :, None]])
        return pool_traces

    @lru_cache(maxsize=1)
    def predict(self, first_n_samples=-1):
        """
        Parameters
        ----------
        first_n_samples: int
            Take only the first n samples from the test data

        Returns
        -------
        y_pred: array of shape (first_n_samples, nr_time_steps, nr_classes)
            The predicted classes
        """
        X, _ = self.get_test_data(first_n_samples)
        y_pred = self.model.predict(X[:, :, None])
        return y_pred

    def predict_single_trace(self, x):
        y_pred = self.model.predict(x[None, :, None])
        return y_pred

    def class_prediction(self, on_last_n_time_steps=250):
        y_pred = self.predict()[:, -on_last_n_time_steps:, :].sum(1).argmax(1)
        return y_pred

    def confusion_matrix(self, on_last_n_time_steps=250):
        y_pred = self.predict()[:, -on_last_n_time_steps:, :].sum(1).argmax(1)
        _, y_gt = self.get_test_data()
        y_target = y_gt.argmax(1)
        conf_mat = confusion_matrix(y_target.flatten(), y_pred.flatten())
        return conf_mat
