import argparse

parser = argparse.ArgumentParser(

        description='This class is used to train a transformer-based model on admission notes, labelled with <3 by Patrick.')

parser.add_argument('-o', '--output', type=str, default="",

                    help='Choose where to save the model after training. Saving is optional.')

parser.add_argument('-lr', '--learning_rate', type=float, default=1e-2,

                    help='Choose the learning rate of the model.')

parser.add_argument('-b', '--batch_size', type=int, default=16,

                    help='Choose the batch size of the model.')

parser.add_argument('-e', '--epochs', type=int, default=5,

                    help='Choose the epochs of the model.')

parser.add_argument('-opt', '--optimizer', type=str, default='SGD',

                    help='Choose the optimizer to be used for the model: SDG | Adam')

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('-v', '--verbose', type=bool, default=False,

                    help='Choose whether the model should be evaluated after each epoch or only after the training.')

parser.add_argument('-l', '--input_length', type=int, default=128,

                    help='Choose the maximum length of the model\'s input layer.')

parser.add_argument('-ag', '--data_augmentation', type=bool, default=False,

                    help='Choose whether data-augmentation should be used.')

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')

from utils.dataloader import Dataloader

from utils.BertArchitecture import BertNER

from utils.BertArchitecture import BioBertNER

from utils.metric_tracking import MetricsTracking

from utils.training import train_loop, testing

import torch

from torch.optim import SGD

from torch.optim import Adam

from torch.utils.data import DataLoader

import numpy as np

import pandas as pd

from tqdm import tqdm

#-------MAIN-------#

if not args.transfer_learning:

    print("Training base BERT model...")

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

    if args.type == 'Medical Condition':

        type = 'MEDCOND'

    elif args.type == 'Symptom':

        type = 'SYMPTOM'

    elif args.type == 'Medication':

        type = 'MEDICATION'

    elif args.type == 'Vital Statistic':

        type = 'VITALSTAT'

    elif args.type == 'Measurement Value':

        type = 'MEASVAL'

    elif args.type == 'Negation Cue':

        type = 'NEGATION'

    elif args.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:

    print("Training BERT model based on BioBERT diseases...")

    if not args.type == 'Medical Condition':

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

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

    type = 'DISEASE'

label_to_ids = {

    'B-' + type: 0,

    'I-' + type: 1,

    'O': 2

    }

ids_to_label = {

    0:'B-' + type,

    1:'I-' + type,

    2:'O'

    }

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

train, val, test = dataloader.load_dataset(augment = args.data_augmentation)

if args.optimizer == 'SGD':

    print("Using SGD optimizer...")

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

else:

    print("Using Adam optimizer...")

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

parameters = {

    "model": model,

    "train_dataset": train,

    "eval_dataset" : val,

    "optimizer" : optimizer,

    "batch_size" : args.batch_size,

    "epochs" : args.epochs,

    "type" : type

}

train_loop(**parameters, verbose=args.verbose)

testing(model, test, args.batch_size, type)

#save model if wanted

if args.output:

    torch.save(model.state_dict(), args.output)

    print(f"Model has successfully been saved at {args.output}!")