"""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, 0, 25, 33), 'warnings.filterwarnings', 'warnings.filterwarnings', ({(25, 24, 25, 32): '"""ignore"""'}, {}), "('ignore')", False, 'import warnings\n'), ((33, 13, 34, 62), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((82, 15, 82, 33), 'auxiliary.log_dir_name', 'log_dir_name', ({(82, 28, 82, 32): 'args'}, {}), '(args)', False, 'from auxiliary import process_features, load_dataset, build_columns, log_dir_name\n'), ((86, 17, 86, 71), 'os.path.join', 'os.path.join', ({(86, 30, 86, 43): '"""experiments"""', (86, 45, 86, 60): 'experiment_name', (86, 62, 86, 70): 'dir_name'}, {}), "('experiments', experiment_name, dir_name)", False, 'import os\n'), ((90, 41, 90, 71), 'auxiliary.load_dataset', 'load_dataset', ({(90, 54, 90, 70): 'args.dataset_dir'}, {}), '(args.dataset_dir)', False, 'from auxiliary import process_features, load_dataset, build_columns, log_dir_name\n'), ((97, 57, 97, 88), 'auxiliary.build_columns', 'build_columns', ({(97, 71, 97, 78): 'dataset', (97, 80, 97, 87): 'columns'}, {}), '(dataset, columns)', False, 'from auxiliary import process_features, load_dataset, build_columns, log_dir_name\n'), ((100, 23, 100, 100), 'auxiliary.process_features', 'process_features', ({(100, 40, 100, 47): 'dataset', (100, 49, 100, 64): 'one_hot_columns', (100, 66, 100, 81): 'numeric_columns', (100, 83, 100, 99): 'embedded_columns'}, {}), '(dataset, one_hot_columns, numeric_columns, embedded_columns)', False, 'from auxiliary import process_features, load_dataset, build_columns, log_dir_name\n'), ((102, 19, 102, 100), 'auxiliary.process_features', 'process_features', ({(102, 36, 102, 47): 'dev_dataset', (102, 49, 102, 64): 'one_hot_columns', (102, 66, 102, 81): 'numeric_columns', (102, 83, 102, 99): 'embedded_columns'}, {}), '(dev_dataset, one_hot_columns, numeric_columns,\n    embedded_columns)', False, 'from auxiliary import process_features, load_dataset, build_columns, log_dir_name\n'), ((120, 4, 120, 36), 'tensorflow.keras.backend.clear_session', 'tf.keras.backend.clear_session', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((135, 28, 135, 99), 'tensorflow.keras.layers.Input', 'layers.Input', (), '', False, 'from tensorflow.keras import layers, models\n'), ((139, 15, 139, 77), 'tensorflow.keras.layers.concatenate', 'layers.concatenate', ({(139, 34, 139, 76): 'embedding_layers + [direct_features_input]'}, {}), '(embedding_layers + [direct_features_input])', False, 'from tensorflow.keras import layers, models\n'), ((150, 12, 150, 61), 'tensorflow.keras.models.Model', 'models.Model', (), '', False, 'from tensorflow.keras import layers, models\n'), ((157, 4, 157, 42), 'mlflow.set_experiment', 'mlflow.set_experiment', ({(157, 26, 157, 41): 'experiment_name'}, {}), '(experiment_name)', False, 'import mlflow\n'), ((159, 16, 159, 69), 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', (), '', True, 'import tensorflow as tf\n'), ((164, 27, 164, 73), 'tensorflow.keras.callbacks.TensorBoard', 'tf.keras.callbacks.TensorBoard', (), '', True, 'import tensorflow as tf\n'), ((87, 11, 87, 37), 'os.path.exists', 'os.path.exists', ({(87, 26, 87, 36): 'output_dir'}, {}), '(output_dir)', False, 'import os\n'), ((88, 8, 88, 31), 'os.makedirs', 'os.makedirs', ({(88, 20, 88, 30): 'output_dir'}, {}), '(output_dir)', False, 'import os\n'), ((126, 22, 126, 65), 'tensorflow.keras.layers.Input', 'layers.Input', (), '', False, 'from tensorflow.keras import layers, models\n'), ((142, 13, 142, 67), 'tensorflow.keras.layers.Dense', 'layers.Dense', (), '', False, 'from tensorflow.keras import layers, models\n'), ((148, 19, 148, 62), 'tensorflow.keras.layers.Dense', 'layers.Dense', (), '', False, 'from tensorflow.keras import layers, models\n'), ((166, 9, 166, 38), 'mlflow.start_run', 'mlflow.start_run', (), '', False, 'import mlflow\n'), ((168, 8, 168, 65), 'mlflow.log_param', 'mlflow.log_param', ({(168, 25, 168, 44): '"""hidden_layer_size"""', (168, 46, 168, 64): 'hidden_layer_sizes'}, {}), "('hidden_layer_size', hidden_layer_sizes)", False, 'import mlflow\n'), ((169, 8, 169, 44), 'mlflow.log_param', 'mlflow.log_param', ({(169, 25, 169, 34): '"""dropout"""', (169, 36, 169, 43): 'dropout'}, {}), "('dropout', dropout)", False, 'import mlflow\n'), ((170, 8, 170, 62), 'mlflow.log_param', 'mlflow.log_param', ({(170, 25, 170, 43): '"""embedded_columns"""', (170, 45, 170, 61): 'embedded_columns'}, {}), "('embedded_columns', embedded_columns)", False, 'import mlflow\n'), ((171, 8, 171, 60), 'mlflow.log_param', 'mlflow.log_param', ({(171, 25, 171, 42): '"""one_hot_columns"""', (171, 44, 171, 59): 'one_hot_columns'}, {}), "('one_hot_columns', one_hot_columns)", False, 'import mlflow\n'), ((172, 8, 172, 60), 'mlflow.log_param', 'mlflow.log_param', ({(172, 25, 172, 42): '"""numeric_columns"""', (172, 44, 172, 59): 'numeric_columns'}, {}), "('numeric_columns', numeric_columns)", False, 'import mlflow\n'), ((173, 8, 173, 42), 'mlflow.log_param', 'mlflow.log_param', ({(173, 25, 173, 33): '"""epochs"""', (173, 35, 173, 41): 'epochs'}, {}), "('epochs', epochs)", False, 'import mlflow\n'), ((174, 8, 174, 50), 'mlflow.log_param', 'mlflow.log_param', ({(174, 25, 174, 37): '"""batch_size"""', (174, 39, 174, 49): 'batch_size'}, {}), "('batch_size', batch_size)", False, 'import mlflow\n'), ((175, 8, 175, 56), 'mlflow.log_param', 'mlflow.log_param', ({(175, 25, 175, 40): '"""learning_rate"""', (175, 42, 175, 55): 'learning_rate'}, {}), "('learning_rate', learning_rate)", False, 'import mlflow\n'), ((186, 21, 186, 63), 'os.path.join', 'os.path.join', ({(186, 34, 186, 44): 'output_dir', (186, 46, 186, 62): '"""history.pickle"""'}, {}), "(output_dir, 'history.pickle')", False, 'import os\n'), ((197, 8, 197, 39), 'mlflow.log_metric', 'mlflow.log_metric', ({(197, 26, 197, 32): '"""loss"""', (197, 34, 197, 38): 'loss'}, {}), "('loss', loss)", False, 'import mlflow\n'), ((198, 8, 198, 47), 'mlflow.log_metric', 'mlflow.log_metric', ({(198, 26, 198, 36): '"""accuracy"""', (198, 38, 198, 46): 'accuracy'}, {}), "('accuracy', accuracy)", False, 'import mlflow\n'), ((206, 29, 206, 59), 'numpy.argmax', 'np.argmax', (), '', True, 'import numpy as np\n'), ((213, 21, 213, 59), 'os.path.join', 'os.path.join', ({(213, 34, 213, 44): 'output_dir', (213, 46, 213, 58): '"""submit.csv"""'}, {}), "(output_dir, 'submit.csv')", False, 'import os\n'), ((214, 22, 214, 100), 'pandas.DataFrame', 'pd.DataFrame', (), '', True, 'import pandas as pd\n'), ((112, 13, 112, 63), 'tensorflow.data.Dataset.from_tensor_slices', 'tf.data.Dataset.from_tensor_slices', ({(112, 48, 112, 62): '(X_dev, y_dev)'}, {}), '((X_dev, y_dev))', True, 'import tensorflow as tf\n'), ((188, 12, 188, 43), 'pickle.dump', 'pickle.dump', ({(188, 24, 188, 39): 'history.history', (188, 41, 188, 42): 'f'}, {}), '(history.history, f)', False, 'import pickle\n'), ((146, 20, 146, 70), 'tensorflow.keras.layers.Dense', 'layers.Dense', (), '', False, 'from tensorflow.keras import layers, models\n'), ((109, 15, 109, 69), 'tensorflow.data.Dataset.from_tensor_slices', 'tf.data.Dataset.from_tensor_slices', ({(109, 50, 109, 68): '(X_train, y_train)'}, {}), '((X_train, y_train))', True, 'import tensorflow as tf\n'), ((113, 49, 114, 84), 'auxiliary.process_features', 'process_features', (), '', False, 'from auxiliary import process_features, load_dataset, build_columns, log_dir_name\n'), ((131, 23, 131, 87), 'tensorflow.keras.layers.Embedding', 'layers.Embedding', (), '', False, 'from tensorflow.keras import layers, models\n')]