#!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)

[((5713, 5827), 'tensorflow.keras.callbacks.ModelCheckpoint', 'tf.keras.callbacks.ModelCheckpoint', (['filepath'], {'monitor': '"""loss"""', 'verbos': '(0)', 'save_best_only': '(True)', 'save_freq': '"""epoch"""'}), "(filepath, monitor='loss', verbos=0,\n    save_best_only=True, save_freq='epoch')\n", (5747, 5827), True, 'import tensorflow as tf\n'), ((5966, 5980), 'pandas.DataFrame', 'pd.DataFrame', ([], {}), '()\n', (5978, 5980), True, 'import pandas as pd\n'), ((4856, 4884), 'tensorflow.losses.MeanSquaredError', 'tf.losses.MeanSquaredError', ([], {}), '()\n', (4882, 4884), True, 'import tensorflow as tf\n'), ((4896, 4946), 'tensorflow.optimizers.Adam', 'tf.optimizers.Adam', ([], {'learning_rate': 'lr', 'amsgrad': '(True)'}), '(learning_rate=lr, amsgrad=True)\n', (4914, 4946), True, 'import tensorflow as tf\n'), ((735, 842), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""', 'input_shape': '(1, 50, attrs)'}), "(filters=10, kernel_size=5, padding='same', activation='relu',\n    input_shape=(1, 50, attrs))\n", (748, 842), False, 'from tensorflow.keras import layers\n'), ((856, 879), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (870, 879), False, 'from tensorflow.keras import layers\n'), ((897, 972), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), "(filters=10, kernel_size=5, padding='same', activation='relu')\n", (910, 972), False, 'from tensorflow.keras import layers\n'), ((990, 1013), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (1004, 1013), False, 'from tensorflow.keras import layers\n'), ((1031, 1105), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(1)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), "(filters=1, kernel_size=5, padding='same', activation='relu')\n", (1044, 1105), False, 'from tensorflow.keras import layers\n'), ((1123, 1146), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (1137, 1146), False, 'from tensorflow.keras import layers\n'), ((1164, 1191), 'tensorflow.keras.layers.BatchNormalization', 'layers.BatchNormalization', ([], {}), '()\n', (1189, 1191), False, 'from tensorflow.keras import layers\n'), ((1209, 1225), 'tensorflow.keras.layers.Flatten', 'layers.Flatten', ([], {}), '()\n', (1223, 1225), False, 'from tensorflow.keras import layers\n'), ((1243, 1267), 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(50)', '"""relu"""'], {}), "(50, 'relu')\n", (1255, 1267), False, 'from tensorflow.keras import layers\n'), ((1285, 1300), 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), '(1)\n', (1297, 1300), False, 'from tensorflow.keras import layers\n'), ((1388, 1495), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""', 'input_shape': '(1, 50, attrs)'}), "(filters=10, kernel_size=5, padding='same', activation='relu',\n    input_shape=(1, 50, attrs))\n", (1401, 1495), False, 'from tensorflow.keras import layers\n'), ((1509, 1532), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (1523, 1532), False, 'from tensorflow.keras import layers\n'), ((1550, 1625), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), "(filters=10, kernel_size=5, padding='same', activation='relu')\n", (1563, 1625), False, 'from tensorflow.keras import layers\n'), ((1643, 1666), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (1657, 1666), False, 'from tensorflow.keras import layers\n'), ((1684, 1758), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(1)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), "(filters=1, kernel_size=5, padding='same', activation='relu')\n", (1697, 1758), False, 'from tensorflow.keras import layers\n'), ((1776, 1799), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (1790, 1799), False, 'from tensorflow.keras import layers\n'), ((1817, 1833), 'tensorflow.keras.layers.Flatten', 'layers.Flatten', ([], {}), '()\n', (1831, 1833), False, 'from tensorflow.keras import layers\n'), ((1851, 1875), 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(50)', '"""relu"""'], {}), "(50, 'relu')\n", (1863, 1875), False, 'from tensorflow.keras import layers\n'), ((1893, 1908), 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), '(1)\n', (1905, 1908), False, 'from tensorflow.keras import layers\n'), ((2047, 2154), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""', 'input_shape': '(1, 50, attrs)'}), "(filters=10, kernel_size=5, padding='same', activation='relu',\n    input_shape=(1, 50, attrs))\n", (2060, 2154), False, 'from tensorflow.keras import layers\n'), ((2168, 2191), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (2182, 2191), False, 'from tensorflow.keras import layers\n'), ((2209, 2284), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), "(filters=10, kernel_size=5, padding='same', activation='relu')\n", (2222, 2284), False, 'from tensorflow.keras import layers\n'), ((2302, 2325), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (2316, 2325), False, 'from tensorflow.keras import layers\n'), ((2343, 2417), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(1)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), "(filters=1, kernel_size=5, padding='same', activation='relu')\n", (2356, 2417), False, 'from tensorflow.keras import layers\n'), ((2435, 2458), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (2449, 2458), False, 'from tensorflow.keras import layers\n'), ((2476, 2503), 'tensorflow.keras.layers.BatchNormalization', 'layers.BatchNormalization', ([], {}), '()\n', (2501, 2503), False, 'from tensorflow.keras import layers\n'), ((2521, 2544), 'tensorflow.keras.layers.Reshape', 'layers.Reshape', (['(5, 10)'], {}), '((5, 10))\n', (2535, 2544), False, 'from tensorflow.keras import layers\n'), ((2562, 2601), 'tensorflow.keras.layers.LSTM', 'layers.LSTM', (['(30)'], {'return_sequences': '(False)'}), '(30, return_sequences=False)\n', (2573, 2601), False, 'from tensorflow.keras import layers\n'), ((2619, 2643), 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(50)', '"""relu"""'], {}), "(50, 'relu')\n", (2631, 2643), False, 'from tensorflow.keras import layers\n'), ((2661, 2676), 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), '(1)\n', (2673, 2676), False, 'from tensorflow.keras import layers\n'), ((2764, 2871), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""', 'input_shape': '(1, 50, attrs)'}), "(filters=10, kernel_size=5, padding='same', activation='relu',\n    input_shape=(1, 50, attrs))\n", (2777, 2871), False, 'from tensorflow.keras import layers\n'), ((2885, 2908), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (2899, 2908), False, 'from tensorflow.keras import layers\n'), ((2926, 3001), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), "(filters=10, kernel_size=5, padding='same', activation='relu')\n", (2939, 3001), False, 'from tensorflow.keras import layers\n'), ((3019, 3042), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (3033, 3042), False, 'from tensorflow.keras import layers\n'), ((3060, 3134), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(1)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), "(filters=1, kernel_size=5, padding='same', activation='relu')\n", (3073, 3134), False, 'from tensorflow.keras import layers\n'), ((3152, 3175), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (3166, 3175), False, 'from tensorflow.keras import layers\n'), ((3193, 3216), 'tensorflow.keras.layers.Reshape', 'layers.Reshape', (['(5, 10)'], {}), '((5, 10))\n', (3207, 3216), False, 'from tensorflow.keras import layers\n'), ((3234, 3273), 'tensorflow.keras.layers.LSTM', 'layers.LSTM', (['(30)'], {'return_sequences': '(False)'}), '(30, return_sequences=False)\n', (3245, 3273), False, 'from tensorflow.keras import layers\n'), ((3291, 3315), 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(50)', '"""relu"""'], {}), "(50, 'relu')\n", (3303, 3315), False, 'from tensorflow.keras import layers\n'), ((3333, 3348), 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), '(1)\n', (3345, 3348), False, 'from tensorflow.keras import layers\n'), ((3488, 3595), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""', 'input_shape': '(1, 50, attrs)'}), "(filters=10, kernel_size=5, padding='same', activation='relu',\n    input_shape=(1, 50, attrs))\n", (3501, 3595), False, 'from tensorflow.keras import layers\n'), ((3609, 3632), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (3623, 3632), False, 'from tensorflow.keras import layers\n'), ((3650, 3725), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), "(filters=10, kernel_size=5, padding='same', activation='relu')\n", (3663, 3725), False, 'from tensorflow.keras import layers\n'), ((3743, 3766), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (3757, 3766), False, 'from tensorflow.keras import layers\n'), ((3784, 3858), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(1)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), "(filters=1, kernel_size=5, padding='same', activation='relu')\n", (3797, 3858), False, 'from tensorflow.keras import layers\n'), ((3876, 3899), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (3890, 3899), False, 'from tensorflow.keras import layers\n'), ((3917, 3944), 'tensorflow.keras.layers.BatchNormalization', 'layers.BatchNormalization', ([], {}), '()\n', (3942, 3944), False, 'from tensorflow.keras import layers\n'), ((3962, 3985), 'tensorflow.keras.layers.Reshape', 'layers.Reshape', (['(5, 10)'], {}), '((5, 10))\n', (3976, 3985), False, 'from tensorflow.keras import layers\n'), ((4003, 4041), 'tensorflow.keras.layers.LSTM', 'layers.LSTM', (['(30)'], {'return_sequences': '(True)'}), '(30, return_sequences=True)\n', (4014, 4041), False, 'from tensorflow.keras import layers\n'), ((4059, 4098), 'tensorflow.keras.layers.LSTM', 'layers.LSTM', (['(30)'], {'return_sequences': '(False)'}), '(30, return_sequences=False)\n', (4070, 4098), False, 'from tensorflow.keras import layers\n'), ((4116, 4131), 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), '(1)\n', (4128, 4131), False, 'from tensorflow.keras import layers\n'), ((4219, 4326), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""', 'input_shape': '(1, 50, attrs)'}), "(filters=10, kernel_size=5, padding='same', activation='relu',\n    input_shape=(1, 50, attrs))\n", (4232, 4326), False, 'from tensorflow.keras import layers\n'), ((4340, 4363), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (4354, 4363), False, 'from tensorflow.keras import layers\n'), ((4381, 4456), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(10)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), "(filters=10, kernel_size=5, padding='same', activation='relu')\n", (4394, 4456), False, 'from tensorflow.keras import layers\n'), ((4474, 4497), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (4488, 4497), False, 'from tensorflow.keras import layers\n'), ((4515, 4589), 'tensorflow.keras.layers.Conv1D', 'layers.Conv1D', ([], {'filters': '(1)', 'kernel_size': '(5)', 'padding': '"""same"""', 'activation': '"""relu"""'}), "(filters=1, kernel_size=5, padding='same', activation='relu')\n", (4528, 4589), False, 'from tensorflow.keras import layers\n'), ((4607, 4630), 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['dropout'], {}), '(dropout)\n', (4621, 4630), False, 'from tensorflow.keras import layers\n'), ((4648, 4671), 'tensorflow.keras.layers.Reshape', 'layers.Reshape', (['(5, 10)'], {}), '((5, 10))\n', (4662, 4671), False, 'from tensorflow.keras import layers\n'), ((4689, 4727), 'tensorflow.keras.layers.LSTM', 'layers.LSTM', (['(30)'], {'return_sequences': '(True)'}), '(30, return_sequences=True)\n', (4700, 4727), False, 'from tensorflow.keras import layers\n'), ((4745, 4784), 'tensorflow.keras.layers.LSTM', 'layers.LSTM', (['(30)'], {'return_sequences': '(False)'}), '(30, return_sequences=False)\n', (4756, 4784), False, 'from tensorflow.keras import layers\n'), ((4802, 4817), 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), '(1)\n', (4814, 4817), False, 'from tensorflow.keras import layers\n')]