#!/usr/bin/env python

# ================ IMPORT LIBRARIES ================ #
import sys, os, fnmatch, time
import numpy as np
import pandas as pd

from sklearn.model_selection import train_test_split

sys.path.insert(0, os.path.dirname(os.getcwd()))

from dataset_generator import DataGenerator

import tensorflow as tf
import tensorflow_addons as tfa
from tensorflow import keras
from tensorflow.keras import layers, Input, Sequential
from tensorflow.keras.layers import Bidirectional, LSTM, Dropout, BatchNormalization, Dense, Conv1D, LeakyReLU, AveragePooling1D, Flatten, Reshape, MaxPooling1D
from tensorflow.keras.optimizers import Adam, Adadelta, SGD
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.metrics import RootMeanSquaredError, MeanAbsoluteError

n_timesteps = 501
n_features = 30 
n_outputs = 1

COUNT_MODEL = "FINAL" # This will be appended to the saved model's name. To make sure to not overwrite models, increase this.
MAX_QUEUE_SIZE = 5000
WORKERS = 6

input_shape = (n_timesteps, n_features)

# Input and output folders
PATH_DATA_PROCESSED_DL = sys.argv[1]
PATH_OUTPUT = sys.argv[2]

# ================ INITIAL LOGS ================ #

print("LOGGING: Imported all modules")

# ================ LOAD PREPROCESSED DATA ================ #

# Step 1: Get all the files in the output folder
file_names = os.listdir(PATH_DATA_PROCESSED_DL)

# Step 2: Get the full paths of the files (without extensions)
files = [os.path.splitext(os.path.join(PATH_DATA_PROCESSED_DL, file_name))[0] for file_name in fnmatch.filter(file_names, "*.zarr")]

# Step 3: Load all the metadata
frames = []

for idx, feature_file in enumerate(files):
    df_metadata = pd.read_csv(feature_file.replace("processed_raw_", "processed_metadata_") + ".csv")
    frames.append(df_metadata)

df_metadata = pd.concat(frames) 

# Step 4: Add missing age information based on the age group the subject is in
df_metadata['age_months'].fillna(df_metadata['age_group'], inplace=True)
df_metadata['age_days'].fillna(df_metadata['age_group']*30, inplace=True)
df_metadata['age_years'].fillna(df_metadata['age_group']/12, inplace=True)

# Step 5: List all the unique subject IDs
subject_ids = sorted(list(set(df_metadata["code"].tolist())))

# Step 6: Split the subjects into train, val and test
IDs_train, IDs_temp = train_test_split(subject_ids, test_size=0.3, random_state=42)
IDs_test, IDs_val = train_test_split(IDs_temp, test_size=0.5, random_state=42)

# Step 7: Initialize DataGenerators
train_generator_noise = DataGenerator(list_IDs = IDs_train,
                                      BASE_PATH = PATH_DATA_PROCESSED_DL,
                                      metadata = df_metadata,
                                      n_average = 30,
                                      batch_size = 10,
                                      gaussian_noise=0.01,
                                      iter_per_epoch = 30,
                                      n_timepoints = 501, 
                                      n_channels=30, 
                                      shuffle=True)

val_generator = DataGenerator(list_IDs = IDs_val,
                              BASE_PATH = PATH_DATA_PROCESSED_DL,
                              metadata = df_metadata,
                              n_average = 30,
                              batch_size = 10,
                              iter_per_epoch = 100,
                              n_timepoints = 501,
                              n_channels=30,
                              shuffle=True)

print("LOGGING: Loaded all data and created generators")

# ================ Encoder model ================ #

try:
    def encoder_model():
        """ Returns the Encoder model from Ismail Fawaz et al. (2019). """
        input_layer = keras.layers.Input(input_shape)

        # conv block -1
        conv1 = keras.layers.Conv1D(filters=128,kernel_size=5,strides=1,padding='same')(input_layer)
        conv1 = tfa.layers.InstanceNormalization()(conv1)
        conv1 = keras.layers.PReLU(shared_axes=[1])(conv1)
        conv1 = keras.layers.Dropout(rate=0.2)(conv1)
        conv1 = keras.layers.MaxPooling1D(pool_size=2)(conv1)
        # conv block -2
        conv2 = keras.layers.Conv1D(filters=256,kernel_size=11,strides=1,padding='same')(conv1)
        conv2 = tfa.layers.InstanceNormalization()(conv2)
        conv2 = keras.layers.PReLU(shared_axes=[1])(conv2)
        conv2 = keras.layers.Dropout(rate=0.2)(conv2)
        conv2 = keras.layers.MaxPooling1D(pool_size=2)(conv2)
        # conv block -3
        conv3 = keras.layers.Conv1D(filters=512,kernel_size=21,strides=1,padding='same')(conv2)
        conv3 = tfa.layers.InstanceNormalization()(conv3)
        conv3 = keras.layers.PReLU(shared_axes=[1])(conv3)
        conv3 = keras.layers.Dropout(rate=0.2)(conv3)
        # split for attention
        attention_data = keras.layers.Lambda(lambda x: x[:,:,:256])(conv3)
        attention_softmax = keras.layers.Lambda(lambda x: x[:,:,256:])(conv3)
        # attention mechanism
        attention_softmax = keras.layers.Softmax()(attention_softmax)
        multiply_layer = keras.layers.Multiply()([attention_softmax,attention_data])
        # last layer
        dense_layer = keras.layers.Dense(units=256,activation='sigmoid')(multiply_layer)
        dense_layer = tfa.layers.InstanceNormalization()(dense_layer)
        # output layer
        flatten_layer = keras.layers.Flatten()(dense_layer)
        output_layer = keras.layers.Dense(1)(flatten_layer)

        model = keras.models.Model(inputs=input_layer, outputs=output_layer)

        return model

    model = encoder_model()

    optimizer = Adam(learning_rate=0.00001)    
                
    model.compile(loss='mean_squared_error', 
                optimizer=optimizer, 
                metrics=[RootMeanSquaredError(), MeanAbsoluteError()])

    output_filename = f'Encoder_regressor_{COUNT_MODEL}'
    output_file = os.path.join(PATH_OUTPUT, output_filename)

    checkpointer = ModelCheckpoint(filepath = output_file + ".hdf5", monitor='val_loss', verbose=1, save_best_only=True)

    epochs = 500

    print("LOGGING: Starting Encoder model training")
    # fit network
    history = model.fit(x=train_generator_noise,
                        validation_data=val_generator,
                        epochs=epochs, 
                        verbose=2, 
                        max_queue_size=MAX_QUEUE_SIZE,
                        workers=WORKERS,  
                        callbacks=[checkpointer])
    print("LOGGING: Finished Encoder model training")
except Exception as e:
    print("LOGGING: Failed Encoder model training:")
    print(e)
    pass

tensorflow.keras.callbacks.ModelCheckpoint pandas.concat tensorflow.keras.layers.Lambda tensorflow.keras.models.Model tensorflow.keras.layers.PReLU tensorflow.keras.layers.Dense tensorflow.keras.layers.Conv1D sklearn.model_selection.train_test_split tensorflow.keras.layers.MaxPooling1D tensorflow.keras.optimizers.Adam tensorflow.keras.layers.Multiply tensorflow.keras.metrics.MeanAbsoluteError tensorflow.keras.metrics.RootMeanSquaredError tensorflow.keras.layers.Dropout tensorflow.keras.layers.Flatten tensorflow.keras.layers.Softmax tensorflow.keras.layers.Input
[(44, 'os.listdir', 'os.listdir', (['PATH_DATA_PROCESSED_DL'], {}), False, 'import sys, os, fnmatch, time\n'), (56, 'pandas.concat', 'pd.concat', (['frames'], {}), True, 'import pandas as pd\n'), (67, 'sklearn.model_selection.train_test_split', 'train_test_split', (['subject_ids'], {'test_size': '(0.3)', 'random_state': '(42)'}), False, 'from sklearn.model_selection import train_test_split\n'), (68, 'sklearn.model_selection.train_test_split', 'train_test_split', (['IDs_temp'], {'test_size': '(0.5)', 'random_state': '(42)'}), False, 'from sklearn.model_selection import train_test_split\n'), (71, 'dataset_generator.DataGenerator', 'DataGenerator', ([], {'list_IDs': 'IDs_train', 'BASE_PATH': 'PATH_DATA_PROCESSED_DL', 'metadata': 'df_metadata', 'n_average': '(30)', 'batch_size': '(10)', 'gaussian_noise': '(0.01)', 'iter_per_epoch': '(30)', 'n_timepoints': '(501)', 'n_channels': '(30)', 'shuffle': '(True)'}), False, 'from dataset_generator import DataGenerator\n'), (82, 'dataset_generator.DataGenerator', 'DataGenerator', ([], {'list_IDs': 'IDs_val', 'BASE_PATH': 'PATH_DATA_PROCESSED_DL', 'metadata': 'df_metadata', 'n_average': '(30)', 'batch_size': '(10)', 'iter_per_epoch': '(100)', 'n_timepoints': '(501)', 'n_channels': '(30)', 'shuffle': '(True)'}), False, 'from dataset_generator import DataGenerator\n'), (137, 'tensorflow.keras.optimizers.Adam', 'Adam', ([], {'learning_rate': '(1e-05)'}), False, 'from tensorflow.keras.optimizers import Adam, Adadelta, SGD\n'), (144, 'os.path.join', 'os.path.join', (['PATH_OUTPUT', 'output_filename'], {}), False, 'import sys, os, fnmatch, time\n'), (146, 'tensorflow.keras.callbacks.ModelCheckpoint', 'ModelCheckpoint', ([], {'filepath': "(output_file + '.hdf5')", 'monitor': '"""val_loss"""', 'verbose': '(1)', 'save_best_only': '(True)'}), False, 'from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau\n'), (10, 'os.getcwd', 'os.getcwd', ([], {}), False, 'import sys, os, fnmatch, time\n'), (47, 'fnmatch.filter', 'fnmatch.filter', (['file_names', '"""*.zarr"""'], {}), False, 'import sys, os, fnmatch, time\n'), (99, 'tensorflow.keras.layers.Input', 'keras.layers.Input', (['input_shape'], {}), False, 'from tensorflow import keras\n'), (131, 'tensorflow.keras.models.Model', 'keras.models.Model', ([], {'inputs': 'input_layer', 'outputs': 'output_layer'}), False, 'from tensorflow import keras\n'), (47, 'os.path.join', 'os.path.join', (['PATH_DATA_PROCESSED_DL', 'file_name'], {}), False, 'import sys, os, fnmatch, time\n'), (102, 'tensorflow.keras.layers.Conv1D', 'keras.layers.Conv1D', ([], {'filters': '(128)', 'kernel_size': '(5)', 'strides': '(1)', 'padding': '"""same"""'}), False, 'from tensorflow import keras\n'), (103, 'tensorflow_addons.layers.InstanceNormalization', 'tfa.layers.InstanceNormalization', ([], {}), True, 'import tensorflow_addons as tfa\n'), (104, 'tensorflow.keras.layers.PReLU', 'keras.layers.PReLU', ([], {'shared_axes': '[1]'}), False, 'from tensorflow import keras\n'), (105, 'tensorflow.keras.layers.Dropout', 'keras.layers.Dropout', ([], {'rate': '(0.2)'}), False, 'from tensorflow import keras\n'), (106, 'tensorflow.keras.layers.MaxPooling1D', 'keras.layers.MaxPooling1D', ([], {'pool_size': '(2)'}), False, 'from tensorflow import keras\n'), (108, 'tensorflow.keras.layers.Conv1D', 'keras.layers.Conv1D', ([], {'filters': '(256)', 'kernel_size': '(11)', 'strides': '(1)', 'padding': '"""same"""'}), False, 'from tensorflow import keras\n'), (109, 'tensorflow_addons.layers.InstanceNormalization', 'tfa.layers.InstanceNormalization', ([], {}), True, 'import tensorflow_addons as tfa\n'), (110, 'tensorflow.keras.layers.PReLU', 'keras.layers.PReLU', ([], {'shared_axes': '[1]'}), False, 'from tensorflow import keras\n'), (111, 'tensorflow.keras.layers.Dropout', 'keras.layers.Dropout', ([], {'rate': '(0.2)'}), False, 'from tensorflow import keras\n'), (112, 'tensorflow.keras.layers.MaxPooling1D', 'keras.layers.MaxPooling1D', ([], {'pool_size': '(2)'}), False, 'from tensorflow import keras\n'), (114, 'tensorflow.keras.layers.Conv1D', 'keras.layers.Conv1D', ([], {'filters': '(512)', 'kernel_size': '(21)', 'strides': '(1)', 'padding': '"""same"""'}), False, 'from tensorflow import keras\n'), (115, 'tensorflow_addons.layers.InstanceNormalization', 'tfa.layers.InstanceNormalization', ([], {}), True, 'import tensorflow_addons as tfa\n'), (116, 'tensorflow.keras.layers.PReLU', 'keras.layers.PReLU', ([], {'shared_axes': '[1]'}), False, 'from tensorflow import keras\n'), (117, 'tensorflow.keras.layers.Dropout', 'keras.layers.Dropout', ([], {'rate': '(0.2)'}), False, 'from tensorflow import keras\n'), (119, 'tensorflow.keras.layers.Lambda', 'keras.layers.Lambda', (['(lambda x: x[:, :, :256])'], {}), False, 'from tensorflow import keras\n'), (120, 'tensorflow.keras.layers.Lambda', 'keras.layers.Lambda', (['(lambda x: x[:, :, 256:])'], {}), False, 'from tensorflow import keras\n'), (122, 'tensorflow.keras.layers.Softmax', 'keras.layers.Softmax', ([], {}), False, 'from tensorflow import keras\n'), (123, 'tensorflow.keras.layers.Multiply', 'keras.layers.Multiply', ([], {}), False, 'from tensorflow import keras\n'), (125, 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', ([], {'units': '(256)', 'activation': '"""sigmoid"""'}), False, 'from tensorflow import keras\n'), (126, 'tensorflow_addons.layers.InstanceNormalization', 'tfa.layers.InstanceNormalization', ([], {}), True, 'import tensorflow_addons as tfa\n'), (128, 'tensorflow.keras.layers.Flatten', 'keras.layers.Flatten', ([], {}), False, 'from tensorflow import keras\n'), (129, 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', (['(1)'], {}), False, 'from tensorflow import keras\n'), (141, 'tensorflow.keras.metrics.RootMeanSquaredError', 'RootMeanSquaredError', ([], {}), False, 'from tensorflow.keras.metrics import RootMeanSquaredError, MeanAbsoluteError\n'), (141, 'tensorflow.keras.metrics.MeanAbsoluteError', 'MeanAbsoluteError', ([], {}), False, 'from tensorflow.keras.metrics import RootMeanSquaredError, MeanAbsoluteError\n')]