from src.constants import *
from keras.models import load_model
from keras.utils.vis_utils import plot_model
from pathlib import Path
import tensorflow as tf
import os
import pickle

# disable tensorflow warnings and file locking
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE"


class Model:
    def __init__(self, params: ModelParameters):
        # model attributes
        self.history = None
        self._model = None
        self.loaded = False

        # default early stopping parameters
        self.best_epoch = 25
        self.best_loss = np.Inf
        self.best_val_loss = np.Inf
        self.best_weights = None

        # bind data
        self.retrain = params.retrain
        self.model_name = params.model_name
        self.model_path = params.model_path
        self.training_dataset = params.train_dataset
        self.validation_dataset = params.validation_dataset
        self.store_history = params.store_history
        self.epochs = params.epochs
        self.alpha = params.alpha
        self.early_stop = params.early_stop
        self.apply_clipping_fix = params.apply_clipping_fix
        self.learning_rate_testing = params.learning_rate_testing
        self.config = params.config
        self.plot_model = params.plot_model

        if self.config is None:
            self.config = [20, 0]

        # if model is not stored or should be retrained, create new model
        if not self.stored() or self.retrain:
            self.build_model()
            self.compile()
            self.fit()
        else:
            # otherwise load model
            self.load()

    def build_model(self):
        # init model as sequential
        self._model = tf.keras.Sequential()
        # check if multi layers were provided
        multi_layer_1 = self.config[1] > 0

        # specify model shape based on model name
        if "multi" in self.model_name:
            shape = (BATCH_SIZE, self.training_dataset.element_spec[0].shape[2])
        else:
            shape = (BATCH_SIZE, 1)

        # first LSTM layer with tanh activation
        self._model.add(
            tf.keras.layers.LSTM(units=self.config[0],
                                 input_shape=shape,
                                 return_sequences=multi_layer_1,
                                 activation='tanh'),
        )

        # second optional LSTM layer with tanh activation
        if self.config[1] > 0:
            self._model.add(
                tf.keras.layers.LSTM(units=self.config[1],
                                     input_shape=(BATCH_SIZE, 1),
                                     return_sequences=False,
                                     activation='tanh'),
            )

        # dense layer with linear activation
        self._model.add(tf.keras.layers.Dense(units=1))

        if self.plot_model:
            plot_model(self._model, to_file='./experiment-results/' + self.model_name, show_shapes=True, show_layer_names=True)

    def compile(self):
        if self.apply_clipping_fix and not self.learning_rate_testing:
            # compile model with adam optimizer with best learning rate and mse loss
            self._model.compile(
                optimizer=tf.keras.optimizers.Adam(learning_rate=self.alpha, clipnorm=1.0, clipvalue=0.5), loss='mse',
                metrics=['mse', 'mae', tf.keras.metrics.RootMeanSquaredError()])
        else:
            # compile model with adam optimizer with default learning rate and mse loss
            self._model.compile(optimizer=tf.keras.optimizers.Adam(), loss='mse', metrics=['mse', 'mae', tf.keras.metrics.RootMeanSquaredError()])

    def fit(self):
        # define early stopping callback
        callbacks = [EarlyStopping(self)] if self.early_stop else []

        # if learning rate testing is enabled, use learning rate scheduler
        if self.learning_rate_testing:
            callbacks.append(tf.keras.callbacks.LearningRateScheduler(
                lambda epoch: 1e-4 * 10 ** (epoch / 100)
            ))

        # fit model
        h = self._model.fit(self.training_dataset,
                            epochs=self.epochs,
                            steps_per_epoch=STEPS_PER_EPOCH,
                            validation_data=self.validation_dataset,
                            validation_steps=VALIDATION_STEPS,
                            verbose=VERBOSE,
                            callbacks=callbacks)
        # store training history
        self.history = h.history

    def save(self, error=''):
        if self.loaded:
            # remove rms from name
            self.model_name = self.model_name.split('-')[0]

        # store model
        name = self.model_path + self.model_name + error
        self._model.save(name + '.h5')

        # store training history
        if self.store_history:
            with open(name + '-history.pkl', 'wb') as f:
                pickle.dump(self.history, f)

        # remove model from memory
        del self._model

    def load(self):
        # load model
        self.loaded = True
        name = self.model_path + self.model_name
        self._model = load_model(name + '.h5')
        # load history
        if self.store_history:
            with open(name + '-history.pkl', 'rb') as f:
                self.history = pickle.load(f)

    def stored(self):
        # find model name
        for file in Path(self.model_path).glob(self.model_name + '*.h5'):
            # remove path and extension
            name = file.name.split('\\')[0].split('.h5')[0]
            self.model_name = name
        # check if model exists
        return Path(self.model_path + self.model_name + '.h5').is_file()

    # predict data
    def predict(self, data):
        return self._model.predict(data)


class EarlyStopping(tf.keras.callbacks.Callback):
    def __init__(self, model_ref=None):
        super(EarlyStopping, self).__init__()
        # model class reference
        self.model_ref = model_ref

    def on_train_begin(self):
        # init default values
        self.model_ref.best_epoch = 25
        self.model_ref.best_loss = np.Inf
        self.model_ref.best_val_loss = np.Inf
        self.model_ref.best_weights = None

    def on_epoch_end(self, epoch, logs=None):
        # get loss values
        loss = logs.get("loss")
        val_loss = logs.get("val_loss")

        # check if loss is less than best loss
        if np.less(val_loss, self.model_ref.best_val_loss):
            # store new best values
            self.model_ref.best_loss = loss
            self.model_ref.best_val_loss = val_loss
            self.model_ref.best_epoch = epoch
            self.model_ref.best_weights = self.model.get_weights()

    def on_train_end(self):
        print("Best epoch was %d with loss %f" % (self.model_ref.best_epoch + 1, self.model_ref.best_val_loss))
        # set best weights
        self.model.set_weights(self.model_ref.best_weights)
