from transformers import BertTokenizer, BertForSequenceClassification

from utils.BertArchitecture import BertNER, BioBertNER

from utils.training import train_loop

from utils.dataloader import Dataloader

import torch

from torch.utils.data import DataLoader

from sklearn.model_selection import KFold

from torch.optim import SGD

from torch.optim import Adam

def get_label_descriptions(transfer_learning, type):

    """

    This function returns the correct label descriptions according to the

    current Model Architecture in use.

    Parameters:

    transfer_learning (bool): Whether we train BioBERT for transfer learning or not.

    Returns:

    tuple:

        - label_to_ids (dict): A dictionary mapping labels to their respective IDs.

        - ids_to_label (dict): A dictionary mapping IDs back to their respective labels.

    """

    if not transfer_learning:

        if type == 'Medical Condition':

            type = 'MEDCOND'

        elif type == 'Symptom':

            type = 'SYMPTOM'

        elif type == 'Medication':

            type = 'MEDICATION'

        elif type == 'Vital Statistic':

            type = 'VITALSTAT'

        elif type == 'Measurement Value':

            type = 'MEASVAL'

        elif type == 'Negation Cue':

            type = 'NEGATION'

        elif type == 'Medical Procedure':

            type = 'PROCEDURE'

        else:    

            raise ValueError('Type of annotation needs to be one of the following: Medical Condition, Symptom, Medication, Vital Statistic, Measurement Value, Negation Cue, Medical Procedure')

    else:

        if not type == 'Medical Condition':

            raise ValueError('Type of annotation needs to be Medical Condition when using BioBERT as baseline.')

        type = 'DISEASE'

    label_to_ids = {

        'B-' + type: 0,

        'I-' + type: 1,

        'O': 2

    }

    ids_to_label = {

        0:'B-' + type,

        1:'I-' + type,

        2:'O'

    }

    return label_to_ids, ids_to_label, type

def initialize_model(transfer_learning):

    """

    Initializes the model architecture according to whether we would like to use BioBERT

    or not.

    Parameters:

    transfer_learning (bool): Whether we train BioBERT for transfer learning or not.

    Returns:

    model (BertNER | BioBertNER)

    """

    if not transfer_learning:

        model = BertNER(3) #O, B-, I- -> 3 entities

    else:

        model = BioBertNER(3)

    return model

def train_fold(transfer_learning, train_idx, val_idx, batch_size, learning_rate, optimizer_name, epoch, type):

    """

    Trains a whole fold during hyperparameter tuning.

    Parameters:

    transfer_learning (bool): Whether we train BioBERT for transfer learning or not.

    train_idx (int): Index of the current split in training data.

    val_idx (int): Index of the current split in testing data.

    batch_size (int): Batch size used for training.

    learning_rate (float): Learning rate used for optimizer.

    optimizer_name (string): Name of the optimizer to be used (SGD or Adam).

    epoch (int): Number of epochs used for training.

    Returns:

    tuple:

        - train_res (dict): A dictionary containing the results obtained during training.

        - test_res (dict): A dictionary containing the results obtained during testing.

    """

    model = initialize_model(transfer_learning)

    if optimizer_name == 'SGD':

        optimizer = SGD(model.parameters(), lr=learning_rate, momentum = 0.9)

    else:

        optimizer = Adam(model.parameters(), lr=learning_rate)

    train_subsampler = torch.utils.data.SubsetRandomSampler(train_idx)

    val_subsampler = torch.utils.data.SubsetRandomSampler(val_idx)

    parameters = {

        "model": model,

        "train_dataset": data,

        "eval_dataset" : data,

        "optimizer" : optimizer,

        "batch_size" : batch_size,

        "epochs" : epoch,

        "train_sampler": train_subsampler,

        "eval_sampler": val_subsampler,

        "type" : type

    }

    train_res, test_res = train_loop(**parameters, verbose=False)

    return train_res, test_res

import argparse

parser = argparse.ArgumentParser(

        description='This class is used to optimize the hyperparameters of either the pretrained BioBERT or the base BERT.')

parser.add_argument('-tr', '--transfer_learning', type=bool, default=False,

                    help='Choose whether the BioBERT model should be used as baseline or not.')

parser.add_argument('-t', '--type', type=str, required=True,

                    help='Specify the type of annotation to process. Type of annotation needs to be one of the following: Medical Condition, Symptom, Medication, Vital Statistic, Measurement Value, Negation Cue, Medical Procedure')

args = parser.parse_args()

if args.type not in ['Medical Condition', 'Symptom', 'Medication', 'Vital Statistic', 'Measurement Value', 'Negation Cue', 'Medical Procedure']:

    raise ValueError('Type of annotation needs to be one of the following: Medical Condition, Symptom, Medication, Vital Statistic, Measurement Value, Negation Cue, Medical Procedure')

#-----hyperparameter grids-----#

batch_sizes = [8, 16]  #[8,16,32]

learning_rates = [0.1] #[0.1, 0.01, 0.001, 0.0001]

optimizers = ['SGD']  #['SGD', 'Adam']

epochs = [1] #[5, 10]

max_tokens = 128

label_to_ids, ids_to_label, type = get_label_descriptions(args.transfer_learning, args.type)

dataloader = Dataloader(label_to_ids, ids_to_label, args.transfer_learning, max_tokens, type)

data = dataloader.load_dataset(full=True)

best_f1_score = 0

best_param_grid = {

    "batch_size": 0,

    "learning_rate": 0,

    "epochs" : 0,

    "optimizer" : "",

    "max_tokens" : 0 

}

for batch_size in batch_sizes:

    for learning_rate in learning_rates:

        for optimizer_name in optimizers:

            for epoch in epochs:

                kf = KFold(n_splits=4, shuffle=True, random_state=7)

                test_f1_scores = []

                for fold, (train_idx, val_idx) in enumerate(kf.split(data)):

                    train_res, test_res = train_fold(args.transfer_learning, train_idx, val_idx, batch_size, learning_rate, optimizer_name, epoch, type)

                    test_f1_scores.append(test_res['avg_f1_score'])

                    print(f"Finished fold {fold+1} of 4!")

                local_best_f1 = sum(test_f1_scores) / len(test_f1_scores)

                if local_best_f1 > best_f1_score:

                    best_f1_score = local_best_f1

                    best_param_grid = {

                        "batch_size": batch_size,

                        "learning_rate": learning_rate,

                        "epochs" : epoch,

                        "optimizer" : optimizer_name,

                        "max_tokens" : max_tokens

                    }

                    print(f"Found new best f1 score: {best_f1_score}")

                    print(best_param_grid)

print("-------FINAL RESULTS-------")

print(f"Best f1 score: {best_f1_score}")

print(best_param_grid)