import argparse

from networks import (

    branchy_linear_network,

    deep_linear_network,

    dual_input_model,

    seq_model,

    simple_branchy_linear_network

)

from data_generator import data_generator

from sklearn.metrics import cohen_kappa_score, f1_score

from tensorflow.keras import optimizers

import tensorflow as tf

import tensorflow_addons as tfa

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Process the inputs')

    parser.add_argument(

        '--model', 

        type=str, 

        help='which model would you like to run',

        default='simple_branchy_linear_network'

    )

    parser.add_argument(

        '--epochs', 

        type=int, 

        help='how many epochs',

        default=12

    )

    parser.add_argument(

        '--verbose', 

        type=int, 

        help='0,1,2',

        default=1

    )

    args = parser.parse_args()

    model_ = args.model

    epochs = args.epochs

    verbose = args.verbose

    X_train, X_test, y_train, y_test, y_train_categorical, y_test_categorical, class_weight = data_generator(

        'data/challenge_1_gut_microbiome_data.csv'

    )

    # selecting your model

    if model_ == 'simple_branchy_linear_network':

        model = simple_branchy_linear_network.simple_branchy_linear_network(class_weight)

    elif model_ == 'branchy_linear_network':

        model = branchy_linear_network.branchy_linear_network(class_weight)

    elif model_ == 'seq_model':

        model = seq_model.seq_model()

    elif model_ == 'deep_linear_network':

        model = deep_linear_network.deep_linear_network(class_weight)

    elif model_ == 'dual_input_model':

        model = dual_input_model.dual_input_model()

    print(model.summary())

    # compile the model

    adam = optimizers.Adam(learning_rate=1e-4, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)

    model.compile(

        optimizer=adam, 

        loss=tf.keras.losses.CategoricalCrossentropy(),

        metrics=[

            tfa.metrics.CohenKappa(num_classes=4, weightage='quadratic'),

            tfa.metrics.F1Score(num_classes=4),

            'accuracy'

        ]

    )

    # create call backs

    reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(

        monitor='val_loss',

        factor=0.5,

        patience=5,

        verbose=verbose,

        min_lr=1e-8,

    )

    checkpoint = tf.keras.callbacks.ModelCheckpoint(

        'models/checkpoint',

        monitor='val_cohen_kappa',

        verbose=verbose,

        save_best_only=True,

        mode='max',

        save_weights_only=False

    )

    # fit the model

    model.fit(

        x=X_train,

        y=y_train_categorical,

        batch_size=16,

        epochs=epochs,

        verbose=verbose,

        validation_data=(X_test, y_test_categorical),

        shuffle=True,

        callbacks=[

            reduce_lr,

            checkpoint

        ]

    )

    # evaluate the model

    # load the best model

    model.load_weights('models/checkpoint')

    y_prob = model.predict(X_test) 

    y_classes = y_prob.argmax(axis=-1)

    ck_score = cohen_kappa_score(

        y_test,

        y_classes,

        weights='quadratic'

    )

    f1_score = f1_score(

        y_test,

        y_classes,

        labels=[0,1,2,3],

        average='weighted',

    )

    print('Results:')

    print('cohen kappa score:', ck_score)

    print('f1_score:', f1_score)