#!python3

import os

import pandas as pd
import tensorflow as tf
from tensorflow.keras import layers

os.environ["CUDA_VISIBLE_DEVICES"] = "0"


# gpu_devices = tf.config.experimental.list_physical_devices("GPU")
# for device in gpu_devices:
#     tf.config.experimental.set_memory_growth(device, True)


def trainModel(data_in, params_in):
    data_in = data_in.take(2048)
    data_in = data_in.shuffle(24)
    data_in = data_in.batch(1024)

    arch = params_in["Architecture"]
    dropout = params_in["Dropout"]
    lr = params_in["LearningRate"]
    attrs = params_in["Attrs"]
    epochs = params_in["Epochs"]

    if arch == "BaseCNN":
        if params_in["BatchNorm"]:
            model = tf.keras.Sequential([
                layers.Conv1D(filters=10, kernel_size=5, padding="same", activation="relu", input_shape=(1, 50, attrs)),
                layers.Dropout(dropout),
                layers.Conv1D(filters=10, kernel_size=5, padding="same", activation="relu"),
                layers.Dropout(dropout),
                layers.Conv1D(filters=1, kernel_size=5, padding="same", activation="relu"),
                layers.Dropout(dropout),
                layers.BatchNormalization(),
                layers.Flatten(),
                layers.Dense(50, "relu"),
                layers.Dense(1)
            ])
        else:
            model = tf.keras.Sequential([
                layers.Conv1D(filters=10, kernel_size=5, padding="same", activation="relu", input_shape=(1, 50, attrs)),
                layers.Dropout(dropout),
                layers.Conv1D(filters=10, kernel_size=5, padding="same", activation="relu"),
                layers.Dropout(dropout),
                layers.Conv1D(filters=1, kernel_size=5, padding="same", activation="relu"),
                layers.Dropout(dropout),
                layers.Flatten(),
                layers.Dense(50, "relu"),
                layers.Dense(1)
            ])

    elif arch == "CNN-LSTM":
        if params_in["BatchNorm"]:
            model = tf.keras.Sequential([
                layers.Conv1D(filters=10, kernel_size=5, padding="same", activation="relu", input_shape=(1, 50, attrs)),
                layers.Dropout(dropout),
                layers.Conv1D(filters=10, kernel_size=5, padding="same", activation="relu"),
                layers.Dropout(dropout),
                layers.Conv1D(filters=1, kernel_size=5, padding="same", activation="relu"),
                layers.Dropout(dropout),
                layers.BatchNormalization(),
                layers.Reshape((5, 10)),
                layers.LSTM(30, return_sequences=False),
                layers.Dense(50, "relu"),
                layers.Dense(1)
            ])
        else:
            model = tf.keras.Sequential([
                layers.Conv1D(filters=10, kernel_size=5, padding="same", activation="relu", input_shape=(1, 50, attrs)),
                layers.Dropout(dropout),
                layers.Conv1D(filters=10, kernel_size=5, padding="same", activation="relu"),
                layers.Dropout(dropout),
                layers.Conv1D(filters=1, kernel_size=5, padding="same", activation="relu"),
                layers.Dropout(dropout),
                layers.Reshape((5, 10)),
                layers.LSTM(30, return_sequences=False),
                layers.Dense(50, "relu"),
                layers.Dense(1)
            ])

    elif arch == "CNN-2LSTM":
        if params_in["BatchNorm"]:
            model = tf.keras.Sequential([
                layers.Conv1D(filters=10, kernel_size=5, padding="same", activation="relu", input_shape=(1, 50, attrs)),
                layers.Dropout(dropout),
                layers.Conv1D(filters=10, kernel_size=5, padding="same", activation="relu"),
                layers.Dropout(dropout),
                layers.Conv1D(filters=1, kernel_size=5, padding="same", activation="relu"),
                layers.Dropout(dropout),
                layers.BatchNormalization(),
                layers.Reshape((5, 10)),
                layers.LSTM(30, return_sequences=True),
                layers.LSTM(30, return_sequences=False),
                layers.Dense(1)
            ])
        else:
            model = tf.keras.Sequential([
                layers.Conv1D(filters=10, kernel_size=5, padding="same", activation="relu", input_shape=(1, 50, attrs)),
                layers.Dropout(dropout),
                layers.Conv1D(filters=10, kernel_size=5, padding="same", activation="relu"),
                layers.Dropout(dropout),
                layers.Conv1D(filters=1, kernel_size=5, padding="same", activation="relu"),
                layers.Dropout(dropout),
                layers.Reshape((5, 10)),
                layers.LSTM(30, return_sequences=True),
                layers.LSTM(30, return_sequences=False),
                layers.Dense(1)
            ])
    model.compile(loss=tf.losses.MeanSquaredError(), optimizer=tf.optimizers.Adam(learning_rate=lr, amsgrad=True))
    filepath = "./checkpoints/Model_in-" + arch + str(attrs) + ".h5"

    losses = []

    class CustomModelCheckPoint(tf.keras.callbacks.Callback):
        def __init__(self, **kargs):
            super(CustomModelCheckPoint, self).__init__(**kargs)
            self.epoch_loss = {}  # accuracy at given epoch

        def on_epoch_begin(self, epoch, logs={}):
            # Things done on beginning of epoch. 
            return

        def on_epoch_end(self, epoch, logs={}):
            # things done on end of the epoch
            self.epoch_loss[epoch] = logs.get("loss")
            losses.append(self.epoch_loss[epoch])

    if params_in["ResumeTraining"]:
        model.load_weights(filepath)

    checkpoint2 = CustomModelCheckPoint()

    checkpoint = tf.keras.callbacks.ModelCheckpoint(filepath, monitor='loss', verbos=0, save_best_only=True,
                                                    save_freq='epoch')
    model.fit(data_in, epochs=epochs, callbacks=[checkpoint, checkpoint2])

    df_loss = pd.DataFrame()
    df_loss["Epochs"] = list(range(1, epochs + 1))
    df_loss["Loss"] = losses
    df_loss.to_csv("./losses/lossTrend.csv", index=False)

[(136, 'tensorflow.keras.callbacks.ModelCheckpoint', 'tf.keras.callbacks.ModelCheckpoint', (['filepath'], {'monitor': '"""loss"""', 'verbos': '(0)', 'save_best_only': '(True)', 'save_freq': '"""epoch"""'}), True, 'import tensorflow as tf\n'), (140, 'pandas.DataFrame', 'pd.DataFrame', ([], {}), True, 'import pandas as pd\n'), (112, 'tensorflow.losses.MeanSquaredError', 'tf.losses.MeanSquaredError', ([], {}), True, 'import tensorflow as tf\n'), (112, 'tensorflow.optimizers.Adam', 'tf.optimizers.Adam', ([], {'learning_rate': 'lr', 'amsgrad': '(True)'}), True, 'import tensorflow as tf\n'), (31, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""', 'input_shape': '(1, 50, attrs)'}), False, 'from tensorflow.keras import layers\n'), (32, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (33, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (34, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (35, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(1)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (36, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (37, 'tensorflow.keras.layers.BatchNormalization', 'layers.BatchNormalization', ([], {}), False, 'from tensorflow.keras import layers\n'), (38, 'tensorflow.keras.layers.Flatten', 'layers.Flatten', ([], {}), False, 'from tensorflow.keras import layers\n'), (39, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(50)', '"""relu"""'], {}), False, 'from tensorflow.keras import layers\n'), (40, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), False, 'from tensorflow.keras import layers\n'), (44, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""', 'input_shape': '(1, 50, attrs)'}), False, 'from tensorflow.keras import layers\n'), (45, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (46, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (47, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (48, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(1)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (49, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (50, 'tensorflow.keras.layers.Flatten', 'layers.Flatten', ([], {}), False, 'from tensorflow.keras import layers\n'), (51, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(50)', '"""relu"""'], {}), False, 'from tensorflow.keras import layers\n'), (52, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), False, 'from tensorflow.keras import layers\n'), (58, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""', 'input_shape': '(1, 50, attrs)'}), False, 'from tensorflow.keras import layers\n'), (59, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (60, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (61, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (62, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(1)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (63, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (64, 'tensorflow.keras.layers.BatchNormalization', 'layers.BatchNormalization', ([], {}), False, 'from tensorflow.keras import layers\n'), (65, 'tensorflow.keras.layers.Reshape', 'layers.Reshape', (['(5, 10)'], {}), False, 'from tensorflow.keras import layers\n'), (66, 'tensorflow.keras.layers.LSTM', 'layers.LSTM', (['(30)'], {'return_sequences': '(False)'}), False, 'from tensorflow.keras import layers\n'), (67, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(50)', '"""relu"""'], {}), False, 'from tensorflow.keras import layers\n'), (68, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), False, 'from tensorflow.keras import layers\n'), (72, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""', 'input_shape': '(1, 50, attrs)'}), False, 'from tensorflow.keras import layers\n'), (73, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (74, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (75, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (76, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(1)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (77, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (78, 'tensorflow.keras.layers.Reshape', 'layers.Reshape', (['(5, 10)'], {}), False, 'from tensorflow.keras import layers\n'), (79, 'tensorflow.keras.layers.LSTM', 'layers.LSTM', (['(30)'], {'return_sequences': '(False)'}), False, 'from tensorflow.keras import layers\n'), (80, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(50)', '"""relu"""'], {}), False, 'from tensorflow.keras import layers\n'), (81, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), False, 'from tensorflow.keras import layers\n'), (87, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""', 'input_shape': '(1, 50, attrs)'}), False, 'from tensorflow.keras import layers\n'), (88, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (89, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (90, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (91, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(1)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (92, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (93, 'tensorflow.keras.layers.BatchNormalization', 'layers.BatchNormalization', ([], {}), False, 'from tensorflow.keras import layers\n'), (94, 'tensorflow.keras.layers.Reshape', 'layers.Reshape', (['(5, 10)'], {}), False, 'from tensorflow.keras import layers\n'), (95, 'tensorflow.keras.layers.LSTM', 'layers.LSTM', (['(30)'], {'return_sequences': '(True)'}), False, 'from tensorflow.keras import layers\n'), (96, 'tensorflow.keras.layers.LSTM', 'layers.LSTM', (['(30)'], {'return_sequences': '(False)'}), False, 'from tensorflow.keras import layers\n'), (97, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), False, 'from tensorflow.keras import layers\n'), (101, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""', 'input_shape': '(1, 50, attrs)'}), False, 'from tensorflow.keras import layers\n'), (102, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (103, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (104, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (105, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(1)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (106, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), False, 'from tensorflow.keras import layers\n'), (107, 'tensorflow.keras.layers.Reshape', 'layers.Reshape', (['(5, 10)'], {}), False, 'from tensorflow.keras import layers\n'), (108, 'tensorflow.keras.layers.LSTM', 'layers.LSTM', (['(30)'], {'return_sequences': '(True)'}), False, 'from tensorflow.keras import layers\n'), (109, 'tensorflow.keras.layers.LSTM', 'layers.LSTM', (['(30)'], {'return_sequences': '(False)'}), False, 'from tensorflow.keras import layers\n'), (110, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), False, 'from tensorflow.keras import layers\n')]