"""Exercise 1

Usage:

$ CUDA_VISIBLE_DEVICES=2 python practico_1_train_petfinder.py --dataset_dir ../ --epochs 30 --dropout 0.1 0.1 --hidden_layer_sizes 200 100

To know which GPU to use, you can check it with the command

$ nvidia-smi
"""

import argparse

import os
import mlflow
import pickle
import numpy as np
import pandas as pd
import tensorflow as tf

from sklearn.model_selection import train_test_split
from tensorflow.keras import layers, models

import warnings
warnings.filterwarnings("ignore")

from auxiliary import process_features, load_dataset, build_columns, log_dir_name

TARGET_COL = 'AdoptionSpeed'


def read_args():
    parser = argparse.ArgumentParser(
        description='Training a MLP on the petfinder dataset')
    # Here you have some examples of classifier parameters. You can add
    # more arguments or change these if you need to.
    parser.add_argument('--experiment_name', type=str, default='Base model',
                        help='Name of the experiment, used in mlflow.')
    parser.add_argument('--dataset_dir', default='../petfinder_dataset', type=str,
                        help='Directory with the training and test files.')
    parser.add_argument('--hidden_layer_sizes', nargs='+', default=[100], type=int,
                        help='Number of hidden units of each hidden layer.')
    parser.add_argument('--epochs', default=50, type=int,
                        help='Number of epochs to train.')
    parser.add_argument('--dropout', nargs='+', default=[0.5], type=float,
                        help='Dropout ratio for every layer.')
    parser.add_argument('--batch_size', type=int, default=32,
                        help='Number of instances in each batch.')
    parser.add_argument('--learning_rate', default=1e-3, type=float,
                        help='Learning rate.')
    args = parser.parse_args()

    assert len(args.hidden_layer_sizes) == len(args.dropout)
    return args

def print_args(args):
    print('-------------------------------------------')
    print('PARAMS ------------------------------------')
    print('-------------------------------------------')
    print('--experiment_name   ', args.experiment_name)
    print('--dataset_dir       ', args.dataset_dir)
    print('--epochs            ', args.epochs)
    print('--hidden_layer_sizes', args.hidden_layer_sizes)
    print('--dropout           ', args.dropout)
    print('--batch_size        ', args.batch_size)
    print('--learning_rate     ', args.learning_rate)
    print('-------------------------------------------')


def main():
    args = read_args()
    print_args(args)

    experiment_name = args.experiment_name
    batch_size = args.batch_size
    learning_rate = args.learning_rate
    hidden_layer_sizes = args.hidden_layer_sizes
    dropout = args.dropout
    epochs = args.epochs

    ### Output directory
    dir_name = log_dir_name(args)
    print()
    print(dir_name)
    print()
    output_dir = os.path.join('experiments', experiment_name, dir_name)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    dataset, dev_dataset, test_dataset = load_dataset(args.dataset_dir)
    nlabels = dataset[TARGET_COL].unique().shape[0]

    columns = [
        'Gender', 'Color1', 'Vaccinated', 'Dewormed',
        'Breed1', 
        'Age', 'Fee', 'Quantity']
    one_hot_columns, embedded_columns, numeric_columns = build_columns(dataset, columns)

    # TODO (optional) put these three types of columns in the same dictionary with "column types"
    X_train, y_train = process_features(dataset, one_hot_columns, numeric_columns, embedded_columns)
    direct_features_input_shape = (X_train['direct_features'].shape[1],)
    X_dev, y_dev = process_features(dev_dataset, one_hot_columns, numeric_columns, embedded_columns)
      

    ###########################################################################################################
    ### TODO: Shuffle train dataset - Done
    ###########################################################################################################
    shuffle_len = X_train['direct_features'].shape[0]
    train_ds = tf.data.Dataset.from_tensor_slices((X_train, y_train)).shuffle(shuffle_len).batch(batch_size)
    ###########################################################################################################

    dev_ds = tf.data.Dataset.from_tensor_slices((X_dev, y_dev)).batch(batch_size)
    test_ds = tf.data.Dataset.from_tensor_slices(process_features(
        test_dataset, one_hot_columns, numeric_columns, embedded_columns, test=True)[0]).batch(batch_size)


    ###########################################################################################################
    ### TODO: Build the Keras model - Done
    ###########################################################################################################
    tf.keras.backend.clear_session()

    # Add one input and one embedding for each embedded column
    embedding_layers = []
    inputs = []
    for embedded_col, max_value in embedded_columns.items():
        input_layer = layers.Input(shape=(1,), name=embedded_col)
        inputs.append(input_layer)
        # Define the embedding layer
        embedding_size = int(max_value / 4)
        embedding_layers.append(
            tf.squeeze(layers.Embedding(input_dim=max_value, output_dim=embedding_size)(input_layer), axis=-2))
        print('Adding embedding of size {} for layer {}'.format(embedding_size, embedded_col))

    # Add the direct features already calculated
    direct_features_input = layers.Input(shape=direct_features_input_shape, name='direct_features')
    inputs.append(direct_features_input)
        
    # Concatenate everything together
    features = layers.concatenate(embedding_layers + [direct_features_input])

    denses = []
    dense1 = layers.Dense(hidden_layer_sizes[0], activation='relu')(features)
    denses.append(dense1)
    if len(hidden_layer_sizes) > 1:
        for hidden_layer_size in hidden_layer_sizes[1:]:
            dense = layers.Dense(hidden_layer_size, activation='relu')(denses[-1])
            denses.append(dense)
    output_layer = layers.Dense(nlabels, activation='softmax')(dense1)

    model = models.Model(inputs=inputs, outputs=output_layer)
    ###########################################################################################################
    

    ###########################################################################################################
    ### TODO: Fit the model - Done
    ###########################################################################################################
    mlflow.set_experiment(experiment_name)

    optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
    model.compile(loss='categorical_crossentropy', optimizer=optimizer,
              metrics=['accuracy'])

    logdir = "logs/scalars/" + dir_name
    tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir)

    with mlflow.start_run(nested=True):
        # Log model hiperparameters first
        mlflow.log_param('hidden_layer_size', hidden_layer_sizes)
        mlflow.log_param('dropout', dropout)
        mlflow.log_param('embedded_columns', embedded_columns)
        mlflow.log_param('one_hot_columns', one_hot_columns)
        mlflow.log_param('numeric_columns', numeric_columns)  # Not using these yet
        mlflow.log_param('epochs', epochs)
        mlflow.log_param('batch_size', batch_size)
        mlflow.log_param('learning_rate', learning_rate)
        

        # Train
        history = model.fit(train_ds, epochs=epochs,
                            validation_data=dev_ds,
                            callbacks=[tensorboard_callback])

        #######################################################################################################
        ### TODO: analyze history to see if model converges/overfits
        #######################################################################################################
        output_csv = os.path.join(output_dir, 'history.pickle')
        with open(output_csv, 'bw') as f:
            pickle.dump(history.history, f)
        #######################################################################################################


        #######################################################################################################
        ### TODO: Evaluate the model, calculating the metrics. - Done
        #######################################################################################################
        loss, accuracy = model.evaluate(dev_ds)
        print("*** Dev loss: {} - accuracy: {}".format(loss, accuracy))
        mlflow.log_metric('loss', loss)
        mlflow.log_metric('accuracy', accuracy)       
        predictions = model.predict(test_ds)
        #######################################################################################################


        #######################################################################################################
        ### TODO: Convert predictions to classes - Done
        #######################################################################################################
        prediction_classes = np.argmax(predictions, axis=1)
        #######################################################################################################


        #######################################################################################################
        ### TODO: Save the results for submission - Done
        #######################################################################################################
        output_csv = os.path.join(output_dir, 'submit.csv')
        submissions = pd.DataFrame(prediction_classes, columns=[TARGET_COL], index=test_dataset.PID)
        submissions.to_csv(output_csv)
        #######################################################################################################


   ###########################################################################################################

print('All operations completed')

if __name__ == '__main__':
    main()

[(25, 'warnings.filterwarnings', 'warnings.filterwarnings', (['"""ignore"""'], {}), False, 'import warnings\n'), (33, 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""Training a MLP on the petfinder dataset"""'}), False, 'import argparse\n'), (82, 'auxiliary.log_dir_name', 'log_dir_name', (['args'], {}), False, 'from auxiliary import process_features, load_dataset, build_columns, log_dir_name\n'), (86, 'os.path.join', 'os.path.join', (['"""experiments"""', 'experiment_name', 'dir_name'], {}), False, 'import os\n'), (90, 'auxiliary.load_dataset', 'load_dataset', (['args.dataset_dir'], {}), False, 'from auxiliary import process_features, load_dataset, build_columns, log_dir_name\n'), (97, 'auxiliary.build_columns', 'build_columns', (['dataset', 'columns'], {}), False, 'from auxiliary import process_features, load_dataset, build_columns, log_dir_name\n'), (100, 'auxiliary.process_features', 'process_features', (['dataset', 'one_hot_columns', 'numeric_columns', 'embedded_columns'], {}), False, 'from auxiliary import process_features, load_dataset, build_columns, log_dir_name\n'), (102, 'auxiliary.process_features', 'process_features', (['dev_dataset', 'one_hot_columns', 'numeric_columns', 'embedded_columns'], {}), False, 'from auxiliary import process_features, load_dataset, build_columns, log_dir_name\n'), (120, 'tensorflow.keras.backend.clear_session', 'tf.keras.backend.clear_session', ([], {}), True, 'import tensorflow as tf\n'), (135, 'tensorflow.keras.layers.Input', 'layers.Input', ([], {'shape': 'direct_features_input_shape', 'name': '"""direct_features"""'}), False, 'from tensorflow.keras import layers, models\n'), (139, 'tensorflow.keras.layers.concatenate', 'layers.concatenate', (['(embedding_layers + [direct_features_input])'], {}), False, 'from tensorflow.keras import layers, models\n'), (150, 'tensorflow.keras.models.Model', 'models.Model', ([], {'inputs': 'inputs', 'outputs': 'output_layer'}), False, 'from tensorflow.keras import layers, models\n'), (157, 'mlflow.set_experiment', 'mlflow.set_experiment', (['experiment_name'], {}), False, 'import mlflow\n'), (159, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {'learning_rate': 'learning_rate'}), True, 'import tensorflow as tf\n'), (164, 'tensorflow.keras.callbacks.TensorBoard', 'tf.keras.callbacks.TensorBoard', ([], {'log_dir': 'logdir'}), True, 'import tensorflow as tf\n'), (87, 'os.path.exists', 'os.path.exists', (['output_dir'], {}), False, 'import os\n'), (88, 'os.makedirs', 'os.makedirs', (['output_dir'], {}), False, 'import os\n'), (126, 'tensorflow.keras.layers.Input', 'layers.Input', ([], {'shape': '(1,)', 'name': 'embedded_col'}), False, 'from tensorflow.keras import layers, models\n'), (142, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['hidden_layer_sizes[0]'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers, models\n'), (148, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['nlabels'], {'activation': '"""softmax"""'}), False, 'from tensorflow.keras import layers, models\n'), (166, 'mlflow.start_run', 'mlflow.start_run', ([], {'nested': '(True)'}), False, 'import mlflow\n'), (168, 'mlflow.log_param', 'mlflow.log_param', (['"""hidden_layer_size"""', 'hidden_layer_sizes'], {}), False, 'import mlflow\n'), (169, 'mlflow.log_param', 'mlflow.log_param', (['"""dropout"""', 'dropout'], {}), False, 'import mlflow\n'), (170, 'mlflow.log_param', 'mlflow.log_param', (['"""embedded_columns"""', 'embedded_columns'], {}), False, 'import mlflow\n'), (171, 'mlflow.log_param', 'mlflow.log_param', (['"""one_hot_columns"""', 'one_hot_columns'], {}), False, 'import mlflow\n'), (172, 'mlflow.log_param', 'mlflow.log_param', (['"""numeric_columns"""', 'numeric_columns'], {}), False, 'import mlflow\n'), (173, 'mlflow.log_param', 'mlflow.log_param', (['"""epochs"""', 'epochs'], {}), False, 'import mlflow\n'), (174, 'mlflow.log_param', 'mlflow.log_param', (['"""batch_size"""', 'batch_size'], {}), False, 'import mlflow\n'), (175, 'mlflow.log_param', 'mlflow.log_param', (['"""learning_rate"""', 'learning_rate'], {}), False, 'import mlflow\n'), (186, 'os.path.join', 'os.path.join', (['output_dir', '"""history.pickle"""'], {}), False, 'import os\n'), (197, 'mlflow.log_metric', 'mlflow.log_metric', (['"""loss"""', 'loss'], {}), False, 'import mlflow\n'), (198, 'mlflow.log_metric', 'mlflow.log_metric', (['"""accuracy"""', 'accuracy'], {}), False, 'import mlflow\n'), (206, 'numpy.argmax', 'np.argmax', (['predictions'], {'axis': '(1)'}), True, 'import numpy as np\n'), (213, 'os.path.join', 'os.path.join', (['output_dir', '"""submit.csv"""'], {}), False, 'import os\n'), (214, 'pandas.DataFrame', 'pd.DataFrame', (['prediction_classes'], {'columns': '[TARGET_COL]', 'index': 'test_dataset.PID'}), True, 'import pandas as pd\n'), (112, 'tensorflow.data.Dataset.from_tensor_slices', 'tf.data.Dataset.from_tensor_slices', (['(X_dev, y_dev)'], {}), True, 'import tensorflow as tf\n'), (188, 'pickle.dump', 'pickle.dump', (['history.history', 'f'], {}), False, 'import pickle\n'), (146, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['hidden_layer_size'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers, models\n'), (109, 'tensorflow.data.Dataset.from_tensor_slices', 'tf.data.Dataset.from_tensor_slices', (['(X_train, y_train)'], {}), True, 'import tensorflow as tf\n'), (113, 'auxiliary.process_features', 'process_features', (['test_dataset', 'one_hot_columns', 'numeric_columns', 'embedded_columns'], {'test': '(True)'}), False, 'from auxiliary import process_features, load_dataset, build_columns, log_dir_name\n'), (131, 'tensorflow.keras.layers.Embedding', 'layers.Embedding', ([], {'input_dim': 'max_value', 'output_dim': 'embedding_size'}), False, 'from tensorflow.keras import layers, models\n')]