import pandas as pd

import numpy as np 

import argparse 

from ast import literal_eval 

import torch 

import torch.nn as nn

from gensim.models import Word2Vec

import logging

from src.bert.bert_model import BERTclassifier

from src.bert.bert_dataset import BERTdataset

from src.bert.bert_train import bert_fit

from src.bert.bert_utils import bert_test_results

from src.rnn.rnn_utils import count_vocab_index, get_emb_matrix

from src.rnn.rnn_dataset import rnndataset

from src.rnn.lstm import LSTMw2vmodel

from src.rnn.gru import GRUw2vmodel

from src.cnn.cnn_dataset import cnndataset

from src.cnn.cnn import character_cnn

from src.hybrid.hybrid_dataset import hybriddataset

from src.hybrid.hybrid import hybrid

from src.hybrid.hybrid_fit import hybrid_fit

from src.hybrid.hybrid_test_results import hybrid_test_results

from src.ovr.mlmodel_data import mlmodel_data

from src.ovr.mlmodel_result import mlmodel_result

from src.ovr.MLmodels import train_classifier

from sklearn.feature_extraction.text import TfidfVectorizer

from src.fit import fit

from src.test_results import test_results

from src.utils import dataloader

def data(args):

    train_diagnosis = pd.read_csv(args.train_path)

    test_diagnosis = pd.read_csv(args.test_path)

    train_diagnosis['ICD9_CODE'] = train_diagnosis['ICD9_CODE'].apply(literal_eval)

    train_diagnosis['ICD9_CATEGORY'] = train_diagnosis['ICD9_CATEGORY'].apply(literal_eval)

    train_diagnosis['ICD10'] = train_diagnosis['ICD10'].apply(literal_eval)

    train_diagnosis['ICD10_CATEGORY'] = train_diagnosis['ICD10_CATEGORY'].apply(literal_eval)

    test_diagnosis['ICD9_CODE'] = test_diagnosis['ICD9_CODE'].apply(literal_eval)

    test_diagnosis['ICD9_CATEGORY'] = test_diagnosis['ICD9_CATEGORY'].apply(literal_eval)

    test_diagnosis['ICD10'] = test_diagnosis['ICD10'].apply(literal_eval)

    test_diagnosis['ICD10_CATEGORY'] = test_diagnosis['ICD10_CATEGORY'].apply(literal_eval)

    return train_diagnosis, test_diagnosis

def run(args):

    train_diagnosis,test_diagnosis = data(args)

    SEED = 2021

    torch.manual_seed(SEED)

    torch.cuda.manual_seed_all(SEED)

    torch.backends.cudnn.deterministic = True

    torch.backends.cudnn.benchmark = False

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    logging.basicConfig(filename='train.log', filemode = 'w', level=logging.DEBUG)

    logging.info("Model Name: %s", args.model_name.upper())

    logging.info("Device: %s", device)

    logging.info("Batch Size: %d", args.batch_size)

    logging.info("Learning Rate: %f", args.learning_rate)

    if args.model_name == "bert":        

        learning_rate = args.learning_rate

        loss_fn = nn.BCELoss()

        opt_fn = torch.optim.Adam

        bert_train_dataset = BERTdataset(train_diagnosis)

        bert_test_dataset = BERTdataset(test_diagnosis)

        bert_train_loader, bert_val_loader, bert_test_loader = dataloader(bert_train_dataset, bert_test_dataset, args.batch_size, args.val_split)

        model = BERTclassifier().to(device)

        bert_fit(args.epochs, model, bert_train_loader, bert_val_loader, args.icd_type, opt_fn, loss_fn, learning_rate, device)

        bert_test_results(model, bert_test_loader, args.icd_type, device)

    elif args.model_name == 'gru':

        learning_rate = args.learning_rate

        loss_fn = nn.BCELoss()

        opt_fn = torch.optim.Adam

        counts, vocab2index = count_vocab_index(train_diagnosis, test_diagnosis)

        rnn_train_dataset = rnndataset(train_diagnosis, vocab2index)

        rnn_test_dataset = rnndataset(train_diagnosis, vocab2index)

        rnn_train_loader, rnn_val_loader, rnn_test_loader = dataloader(rnn_train_dataset, rnn_test_dataset, args.batch_size, args.val_split)

        w2vmodel = Word2Vec.load(args.w2vmodel)

        weights = get_emb_matrix(w2vmodel, counts)

        gruw2vmodel = GRUw2vmodel(weights_matrix = weights, hidden_size = 256, num_layers = 2, device = device).to(device)

        fit(args.epochs, gruw2vmodel, rnn_train_loader, rnn_val_loader, args.icd_type, opt_fn, loss_fn, learning_rate, device)

        test_results(gruw2vmodel, rnn_test_loader, args.icd_type, device)

    elif args.model_name == 'lstm':

        learning_rate = args.learning_rate

        loss_fn = nn.BCELoss()

        opt_fn = torch.optim.Adam

        counts, vocab2index = count_vocab_index(train_diagnosis, test_diagnosis)

        rnn_train_dataset = rnndataset(train_diagnosis, vocab2index)

        rnn_test_dataset = rnndataset(train_diagnosis, vocab2index)

        rnn_train_loader, rnn_val_loader, rnn_test_loader = dataloader(rnn_train_dataset, rnn_test_dataset, args.batch_size, args.val_split)

        w2vmodel = Word2Vec.load(args.w2vmodel)

        weights = get_emb_matrix(w2vmodel, counts)

        lstmw2vmodel = LSTMw2vmodel(weights_matrix = weights, hidden_size = 256, num_layers = 2, device = device).to(device)

        fit(args.epochs, lstmw2vmodel, rnn_train_loader, rnn_val_loader, args.icd_type, opt_fn, loss_fn, learning_rate, device)

        test_results(lstmw2vmodel, rnn_test_loader, args.icd_type, device)

    elif args.model_name == "cnn":

        learning_rate = args.learning_rate

        loss_fn = nn.BCELoss()

        opt_fn = torch.optim.Adam

        cnn_train_dataset = cnndataset(train_diagnosis)

        cnn_test_dataset = cnndataset(test_diagnosis)

        cnn_train_loader, cnn_val_loader, cnn_test_loader = dataloader(cnn_train_dataset, cnn_test_dataset, args.batch_size, args.val_split)

        model = character_cnn(cnn_train_dataset.vocabulary, cnn_train_dataset.sequence_length).to(device)

        fit(args.epochs, model, cnn_train_loader, cnn_val_loader, args.icd_type, opt_fn, loss_fn, learning_rate, device)

        test_results(model, cnn_test_loader, args.icd_type, device)

    elif args.model_name == 'hybrid':

        learning_rate = args.learning_rate

        loss_fn = nn.BCELoss()

        opt_fn = torch.optim.Adam

        counts, vocab2index = count_vocab_index(train_diagnosis, test_diagnosis)

        hybrid_train_dataset = hybriddataset(train_diagnosis, vocab2index)

        hybrid_test_dataset = hybriddataset(train_diagnosis, vocab2index)

        hybrid_train_loader, hybrid_val_loader, hybrid_test_loader = dataloader(hybrid_train_dataset, hybrid_test_dataset, args.batch_size, args.val_split)

        w2vmodel = Word2Vec.load(args.w2vmodel)

        weights = get_emb_matrix(w2vmodel, counts)

        model = hybrid(hybrid_train_dataset.vocabulary, hybrid_train_dataset.sequence_length, weights_matrix = weights, hidden_size = 256, num_layers = 2).to(device)

        hybrid_fit(args.epochs, model, hybrid_train_loader, hybrid_val_loader, args.icd_type, opt_fn, loss_fn, learning_rate, device)

        hybrid_test_results(model, hybrid_test_loader, args.icd_type, device)

    elif args.model_name == 'ovr':

        X_train, y_train = mlmodel_data(train_diagnosis, args.icd_type)

        X_test, y_test = mlmodel_data(test_diagnosis, args.icd_type)

        tfidf_vectorizer = TfidfVectorizer(max_df = 0.8)

        X_train = tfidf_vectorizer.fit_transform(X_train)

        X_test = tfidf_vectorizer.transform(X_test)

        ml_model = train_classifier(X_train, y_train)

        y_predict = ml_model.predict(X_test)

        print('-'*20 + args.icd_type + '-'*20)

        mlmodel_result(y_test, y_predict)

if __name__ == "__main__":

    parser = argparse.ArgumentParser("Automatic Assignment of Medical Codes")

    parser.add_argument("--train_path", type = str, default = './data/train.csv')

    parser.add_argument("--test_path", type = str, default = './data/test.csv')

    parser.add_argument("--model_name", type = str, choices = ['bert', 'hybrid', 'gru', 'lstm', 'cnn', 'ovr'], default = "bert")

    parser.add_argument("--icd_type", type = str, choices = ['icd9cat', 'icd9code', 'icd10cat', 'icd10code'], default = 'icd9cat')

    parser.add_argument("--batch_size", type = int, default = 16)

    parser.add_argument("--val_split", type = float, default = 2/7)

    parser.add_argument("--learning_rate", type = float, default = 2e-5)

    parser.add_argument("--epochs", type = int, default = 4)

    parser.add_argument("--w2vmodel", type = str, default = "w2vmodel.model")

    args = parser.parse_args()

    run(args)