#!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', 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', True, 'import tensorflow as tf\n'), (31, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (32, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (33, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (34, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (35, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (36, 'tensorflow.keras.layers.Dropout', 'layers.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', False, 'from tensorflow.keras import layers\n'), (40, 'tensorflow.keras.layers.Dense', 'layers.Dense', False, 'from tensorflow.keras import layers\n'), (44, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (45, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (46, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (47, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (48, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (49, 'tensorflow.keras.layers.Dropout', 'layers.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', False, 'from tensorflow.keras import layers\n'), (52, 'tensorflow.keras.layers.Dense', 'layers.Dense', False, 'from tensorflow.keras import layers\n'), (58, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (59, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (60, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (61, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (62, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (63, 'tensorflow.keras.layers.Dropout', 'layers.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', False, 'from tensorflow.keras import layers\n'), (66, 'tensorflow.keras.layers.LSTM', 'layers.LSTM', False, 'from tensorflow.keras import layers\n'), (67, 'tensorflow.keras.layers.Dense', 'layers.Dense', False, 'from tensorflow.keras import layers\n'), (68, 'tensorflow.keras.layers.Dense', 'layers.Dense', False, 'from tensorflow.keras import layers\n'), (72, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (73, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (74, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (75, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (76, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (77, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (78, 'tensorflow.keras.layers.Reshape', 'layers.Reshape', False, 'from tensorflow.keras import layers\n'), (79, 'tensorflow.keras.layers.LSTM', 'layers.LSTM', False, 'from tensorflow.keras import layers\n'), (80, 'tensorflow.keras.layers.Dense', 'layers.Dense', False, 'from tensorflow.keras import layers\n'), (81, 'tensorflow.keras.layers.Dense', 'layers.Dense', False, 'from tensorflow.keras import layers\n'), (87, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (88, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (89, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (90, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (91, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (92, 'tensorflow.keras.layers.Dropout', 'layers.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', False, 'from tensorflow.keras import layers\n'), (95, 'tensorflow.keras.layers.LSTM', 'layers.LSTM', False, 'from tensorflow.keras import layers\n'), (96, 'tensorflow.keras.layers.LSTM', 'layers.LSTM', False, 'from tensorflow.keras import layers\n'), (97, 'tensorflow.keras.layers.Dense', 'layers.Dense', False, 'from tensorflow.keras import layers\n'), (101, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (102, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (103, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (104, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (105, 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', False, 'from tensorflow.keras import layers\n'), (106, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', False, 'from tensorflow.keras import layers\n'), (107, 'tensorflow.keras.layers.Reshape', 'layers.Reshape', False, 'from tensorflow.keras import layers\n'), (108, 'tensorflow.keras.layers.LSTM', 'layers.LSTM', False, 'from tensorflow.keras import layers\n'), (109, 'tensorflow.keras.layers.LSTM', 'layers.LSTM', False, 'from tensorflow.keras import layers\n'), (110, 'tensorflow.keras.layers.Dense', 'layers.Dense', False, 'from tensorflow.keras import layers\n')]