import pickle
import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.metrics import accuracy_score
from sklearn.metrics import precision_recall_fscore_support
from sklearn.metrics import confusion_matrix
import functions
import model_tokenizer_loaders

# specify GPU
device = torch.device("cuda")

def get_epochs_train_val_losses(lr, dropout, epochs, folder_name):
    """ Read the three files: number of epochs, training losses per epoch, 
    validation losses per epoch generated by fine_tuning_functions.fine_tuning_model(). 
    
    Inputs:
        - lr: float
        learning rate will be fine-tuned with the TPE algorithm
        - dropout: float between 0 and 1
        - epochs: integer
            number of epochs used to train the model 
        - folder_name: str
            name of the folder where are stored the files
    """ 

    try:
        name = folder_name + '/all_epochs_lr{:}_dropout{:}_epochs{:}.txt'.format(lr, dropout, epochs)
        with open(name, "rb") as fp: 
            all_epochs = pickle.load(fp)

        name = folder_name + '/train_losses_lr{:}_dropout{:}_epochs{:}.txt'.format(lr, dropout, epochs)
        with open(name, "rb") as fp: 
            train_losses = pickle.load(fp)

        name = folder_name + '/val_losses_lr{:}_dropout{:}_epochs{:}.txt'.format(lr, dropout, epochs)
        with open(name, "rb") as fp:   
            valid_losses = pickle.load(fp)
    
    except:
        print("Error: the files do not exist")
    
    return(all_epochs, train_losses, valid_losses)

def plot_train_val_losses(all_epochs, train_losses, valid_losses):
    """ Plot the training and validation losses according to the epochs.
    
    Inputs:
        - all_epochs: list
            1st output of get_epochs_train_val_losses()
        - train_losses: list
            2nd output of get_epochs_train_val_losses()
        - valid_losses: list
            3rd output of get_epochs_train_val_losses()
    """ 

    fig, ax = plt.subplots()
    ax.plot(all_epochs, train_losses, "r", label = "Train")
    ax.plot(all_epochs, valid_losses, "b", label = "Val")
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    legend = ax.legend(loc = 'upper right', shadow = False, fontsize = 'medium')
    plt.show()
    
def get_model(model_name, folder_name, lr, dropout, epochs):
    """ Load model that has already been fine-tuned with the associated
    hyperparameters (lr, dropout, epochs) and returns it.
    
    Inputs:
        - model_name: str
            Name of the model to load. It has to be one of these strings:
            "BERT_CLS", "BERT_improved_CLS", "BERT_mean", "RoBERTa_CLS", 
            "RoBERTa_improved_CLS", "RoBERTa_mean", "RobBERT", "mBERT_test",
            "mBERT_all", "XLM-R_test", "XLM-R_all", "mBERT_all_tanh".
        - folder_name: str
            name of the folder where is stored the file
        - lr: float
        learning rate will be fine-tuned with the TPE algorithm
        - dropout: float between 0 and 1
        - epochs: integer
            number of epochs used to train the model  
    """
    
    # Load the fine-tuned model
    bert = model_tokenizer_loaders.load_model(model_name)
    model = model_tokenizer_loaders.BERT_Arch(bert, dropout, model_name) 
    model = model.to(device)
    # Update the weights
    path = folder_name + '/saved_weights_lr{:}_dropout{:}_epochs{:}.pt'.format(lr, dropout, epochs)
    model.load_state_dict(torch.load(path))
    
    return(model)

def get_test_tensors(model_name, test_text, test_labels):
    """ Return the test tensors after loading the tokenizer. 
    We distinguish two cases: unilingual models and multilingual models.
    In the multilingual case, it returns the test sequences and mask
    for the 3 languages separatly.
    
    Inputs:
        - model_name: str
            Name of the model to load. It has to be one of these strings:
            "BERT_CLS", "BERT_improved_CLS", "BERT_mean", "RoBERTa_CLS", 
            "RoBERTa_improved_CLS", "RoBERTa_mean", "RobBERT", "mBERT_test",
            "mBERT_all", "XLM-R_test", "XLM-R_all", "mBERT_all_tanh".
        - test_text: list
            5th output of split_dataset()
        - test_labels: list
            6th output of split_dataset()
    """
     
    tokenizer = model_tokenizer_loaders.load_tokenizer(model_name)

    #If this is a multilingual model
    if "test" in model_name or "all" in model_name:
        test_seq_nl, test_mask_nl, test_seq_en, test_mask_en, test_seq_fr, test_mask_fr, test_y = functions.to_tensor_test(tokenizer, test_text, test_labels)
        
        return(test_seq_nl, test_mask_nl, test_seq_en, test_mask_en, test_seq_fr, test_mask_fr, test_y)
    
    else:
        test_seq, test_mask, test_y = functions.to_tensor_test(tokenizer, test_text, test_labels)
        
        return(test_seq, test_mask, test_y)
    
def print_accuracy_test_unilingual(model, test_seq, test_mask, test_y, device, model_name):
    """ Compute the accuracy on test set of the unilingual models and the accuracy per label
    and print it.
    The model is too big to be able to predict on the whole test set,
    so we predict on 500 observations at a time and take the average.
    
    Inputs:
        - model: Output of BERT_Arch class
            BERT model and its architecture
        - test_seq: Torch tensor
            1st output of to_tensor_test() for unilingual models
        - test_mask: Torch tensor
            2nd output of to_tensor_test() for unilingual models
        - test_y: Torch tensor
            3rd output of to_tensor_test() for unilingual models
        - device: CUDA device
        - model_name: str
            Name of the model to load. It has to be one of these strings:
            "BERT_CLS", "BERT_improved_CLS", "BERT_mean", "RoBERTa_CLS", 
            "RoBERTa_improved_CLS", "RoBERTa_mean", "RobBERT", "mBERT_test",
            "mBERT_all", "XLM-R_test", "XLM-R_all", "mBERT_all_tanh".
    """
    
    acc_list = []
    CM = []
    for i in range(0, len(test_seq), 500):
        with torch.no_grad():
            preds = model(test_seq[i:i+500].to(device), test_mask[i:i+500].to(device), model_name)
            preds = preds.detach().cpu().numpy()

        # Keep the label number with the highest probability
        preds = np.argmax(preds, axis = 1)

        # Add the accuracy of 500 observations to the list
        acc_list.append(accuracy_score(test_y[i:i+500], preds))
        
        # Compute the accuracy per label using the confusion matrix
        cm = confusion_matrix(test_y[i:i+500], preds)
        # Normalize the confusion matrix
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        # The diagonal gives the accuracy per label
        CM.append(cm.diagonal())

    # Take the average to have the complete accuracy
    acc = sum(acc_list) / len(acc_list)
    acc_labels = np.mean(CM, axis = 0)

    print("Accuracy:")
    print(acc)
    print("Per label:")
    print(acc_labels)
    
def print_accuracy_test_multilingual(model, test_seq_nl, test_mask_nl, test_seq_en, test_mask_en, test_seq_fr, test_mask_fr, test_y, device, model_name):
    """ Compute the accuracy on test set of the multilingual models and print it.
    The model is too big to be able to predict on the whole test set,
    so we predict on 500 observations at a time and take the average.
    
    Inputs:
        - model: Output of BERT_Arch class
            BERT model and its architecture
        - test_seq_nl: Torch tensor
            1st output of to_tensor_test() for multlingual models
        - test_mask_nl: Torch tensor
            2nd output of to_tensor_test() for multlingual models
        - test_seq_en: Torch tensor
            3rd output of to_tensor_test() for multlingual models
        - test_mask_en: Torch tensor
            4th output of to_tensor_test() for multlingual models
        - test_seq_fr: Torch tensor
            5th output of to_tensor_test() for multlingual models
        - test_mask_fr: Torch tensor
            6th output of to_tensor_test() for multlingual models
        - test_y: Torch tensor
            7th output of to_tensor_test() for multlingual models
        - device: CUDA device
        - model_name: str
            Name of the model to load. It has to be one of these strings:
            "BERT_CLS", "BERT_improved_CLS", "BERT_mean", "RoBERTa_CLS", 
            "RoBERTa_improved_CLS", "RoBERTa_mean", "RobBERT", "mBERT_test",
            "mBERT_all", "XLM-R_test", "XLM-R_all", "mBERT_all_tanh".
    """
    
    acc_list_nl = []
    acc_list_en = []
    acc_list_fr = []
    CM_nl = []
    CM_en = []
    CM_fr = []
    for i in range(0, len(test_seq_nl), 500):
        with torch.no_grad():
            preds_nl = model(test_seq_nl[i:i+500].to(device), test_mask_nl[i:i+500].to(device), model_name)
            preds_nl = preds_nl.detach().cpu().numpy()

            preds_en = model(test_seq_en[i:i+500].to(device), test_mask_en[i:i+500].to(device), model_name)
            preds_en = preds_en.detach().cpu().numpy()

            preds_fr = model(test_seq_fr[i:i+500].to(device), test_mask_fr[i:i+500].to(device), model_name)
            preds_fr = preds_fr.detach().cpu().numpy()

        # Keep the label number with the highest probability
        preds_nl = np.argmax(preds_nl, axis = 1)
        preds_en = np.argmax(preds_en, axis = 1)
        preds_fr = np.argmax(preds_fr, axis = 1)

        # Add the accuracy of 500 observations to the list
        acc_list_nl.append(accuracy_score(test_y[i:i+500], preds_nl))
        acc_list_en.append(accuracy_score(test_y[i:i+500], preds_en))
        acc_list_fr.append(accuracy_score(test_y[i:i+500], preds_fr))

        # Compute the accuracy per label using the confusion matrix
        cm_nl = confusion_matrix(test_y[i:i+500], preds_nl)
        # Normalize the confusion matrix
        cm_nl = cm_nl.astype('float') / cm_nl.sum(axis=1)[:, np.newaxis]
        # The diagonal gives the accuracy per label
        CM_nl.append(cm_nl.diagonal())
        
        # Same in English
        cm_en = confusion_matrix(test_y[i:i+500], preds_en)
        cm_en = cm_en.astype('float') / cm_en.sum(axis=1)[:, np.newaxis]
        CM_en.append(cm_en.diagonal())

        # Same in French
        cm_fr = confusion_matrix(test_y[i:i+500], preds_fr)
        cm_fr = cm_fr.astype('float') / cm_fr.sum(axis=1)[:, np.newaxis]
        CM_fr.append(cm_fr.diagonal())

    acc_nl = sum(acc_list_nl) / len(acc_list_nl)
    acc_en = sum(acc_list_en) / len(acc_list_en)
    acc_fr = sum(acc_list_fr) / len(acc_list_fr)

    acc = (acc_nl + acc_en + acc_fr) / 3
    
    print("Accuracy Dutch:")
    print(acc_nl)
    print("Per label:")
    print(np.mean(CM_nl, axis = 0))
    print("Accuracy English:")
    print(acc_en)
    print("Per label:")
    print(np.mean(CM_en, axis = 0))
    print("Accuracy French:")
    print(acc_fr)
    print("Per label:")
    print(np.mean(CM_fr, axis = 0))
    print("Total accuracy:")
    print(acc)
    
def print_precision_recall_f1_test_unilingual(model, test_seq, test_mask, test_y, device, model_name):
    """ Caluclate precision, recall and F1 score in the same way 
    as get_accuracy_test() and print it for unilingual models.
    
    Inputs:
        - model: Output of BERT_Arch class
            BERT model and its architecture
        - test_seq: Torch tensor
            1st output of to_tensor_test() for unilingual models
        - test_mask: Torch tensor
            2nd output of to_tensor_test() for unilingual models
        - test_y: Torch tensor
            3rd output of to_tensor_test() for unilingual models
        - device: CUDA device
        - model_name: str
            Name of the model to load. It has to be one of these strings:
            "BERT_CLS", "BERT_improved_CLS", "BERT_mean", "RoBERTa_CLS", 
            "RoBERTa_improved_CLS", "RoBERTa_mean", "RobBERT", "mBERT_test",
            "mBERT_all", "XLM-R_test", "XLM-R_all", "mBERT_all_tanh".
    """

    all_precision = []
    all_recall = []
    all_f1 = []

    for i in range(0, len(test_seq), 500):
        with torch.no_grad():
            preds = model(test_seq[i:i+500].to(device), test_mask[i:i+500].to(device), model_name)
            preds = preds.detach().cpu().numpy()

        preds = np.argmax(preds, axis = 1)
        all_precision.append(precision_recall_fscore_support(test_y[i:i+500], preds, average=None)[0])
        all_recall.append(precision_recall_fscore_support(test_y[i:i+500], preds, average=None)[1])
        all_f1.append(precision_recall_fscore_support(test_y[i:i+500], preds, average=None)[2])

    print("Precision:")
    print(np.mean(all_precision, axis = 0))
    print("Recall:")
    print(np.mean(all_recall, axis = 0))
    print("F1 score:")
    print(np.mean(all_f1, axis = 0))

    print("Total precision:")
    print(np.mean(all_precision, axis = 0).mean())
    print("Total recall:")
    print(np.mean(all_recall, axis = 0).mean())
    print("Total F1 score:")
    print(np.mean(all_f1, axis = 0).mean())
    
def print_precision_recall_f1_test_multilingual(model, test_seq_nl, test_mask_nl, test_seq_en, test_mask_en, test_seq_fr, test_mask_fr, test_y, device, model_name):
    """ Caluclate precision, recall and F1 score in the same way 
    as get_accuracy_test() and print it for multilingual models.
    
    Inputs:
        - model: Output of BERT_Arch class
            BERT model and its architecture
        - test_seq_nl: Torch tensor
            1st output of to_tensor_test() for multlingual models
        - test_mask_nl: Torch tensor
            2nd output of to_tensor_test() for multlingual models
        - test_seq_en: Torch tensor
            3rd output of to_tensor_test() for multlingual models
        - test_mask_en: Torch tensor
            4th output of to_tensor_test() for multlingual models
        - test_seq_fr: Torch tensor
            5th output of to_tensor_test() for multlingual models
        - test_mask_fr: Torch tensor
            6th output of to_tensor_test() for multlingual models
        - test_y: Torch tensor
            7th output of to_tensor_test() for multlingual models
        - device: CUDA device
        - model_name: str
            Name of the model to load. It has to be one of these strings:
            "BERT_CLS", "BERT_improved_CLS", "BERT_mean", "RoBERTa_CLS", 
            "RoBERTa_improved_CLS", "RoBERTa_mean", "RobBERT", "mBERT_test",
            "mBERT_all", "XLM-R_test", "XLM-R_all", "mBERT_all_tanh".
    """

    all_precision_nl = []
    all_recall_nl = []
    all_f1_nl = []

    all_precision_en = []
    all_recall_en = []
    all_f1_en = []

    all_precision_fr = []
    all_recall_fr = []
    all_f1_fr = []

    for i in range(0, len(test_seq_nl), 500):
        with torch.no_grad():
            preds_nl = model(test_seq_nl[i:i+500].to(device), test_mask_nl[i:i+500].to(device), model_name)
            preds_nl = preds_nl.detach().cpu().numpy()

            preds_en = model(test_seq_en[i:i+500].to(device), test_mask_en[i:i+500].to(device), model_name)
            preds_en = preds_en.detach().cpu().numpy()

            preds_fr = model(test_seq_fr[i:i+500].to(device), test_mask_fr[i:i+500].to(device), model_name)
            preds_fr = preds_fr.detach().cpu().numpy()

        preds_nl = np.argmax(preds_nl, axis = 1)
        preds_en = np.argmax(preds_en, axis = 1)
        preds_fr = np.argmax(preds_fr, axis = 1)

    all_precision_nl.append(precision_recall_fscore_support(test_y[i:i+500], preds_nl, average=None)[0])
    all_recall_nl.append(precision_recall_fscore_support(test_y[i:i+500], preds_nl, average=None)[1])
    all_f1_nl.append(precision_recall_fscore_support(test_y[i:i+500], preds_nl, average=None)[2])

    all_precision_en.append(precision_recall_fscore_support(test_y[i:i+500], preds_en, average=None)[0])
    all_recall_en.append(precision_recall_fscore_support(test_y[i:i+500], preds_en, average=None)[1])
    all_f1_en.append(precision_recall_fscore_support(test_y[i:i+500], preds_en, average=None)[2])

    all_precision_fr.append(precision_recall_fscore_support(test_y[i:i+500], preds_fr, average=None)[0])
    all_recall_fr.append(precision_recall_fscore_support(test_y[i:i+500], preds_fr, average=None)[1])
    all_f1_fr.append(precision_recall_fscore_support(test_y[i:i+500], preds_fr, average=None)[2])

    print("----Dutch----")
    print("Precision:")
    print(np.mean(all_precision_nl, axis = 0))
    print("Recall:")
    print(np.mean(all_recall_nl, axis = 0))
    print("F1 score:")
    print(np.mean(all_f1_nl, axis = 0))
    print("Total precision:")
    print(np.mean(all_precision_nl, axis = 0).mean())
    print("Total recall:")
    print(np.mean(all_recall_nl, axis = 0).mean())
    print("Total F1 score:")
    print(np.mean(all_f1_nl, axis = 0).mean())


    print("----English----")
    print("Precision:")
    print(np.mean(all_precision_en, axis = 0))
    print("Recall:")
    print(np.mean(all_recall_en, axis = 0))
    print("F1 score:")
    print(np.mean(all_f1_en, axis = 0))
    print("Total precision:")
    print(np.mean(all_precision_en, axis = 0).mean())
    print("Total recall:")
    print(np.mean(all_recall_en, axis = 0).mean())
    print("Total F1 score:")
    print(np.mean(all_f1_en, axis = 0).mean())

    print("----French----")
    print("Precision:")
    print(np.mean(all_precision_fr, axis = 0))
    print("Recall:")
    print(np.mean(all_recall_fr, axis = 0))
    print("F1 score:")
    print(np.mean(all_f1_fr, axis = 0))
    print("Total precision:")
    print(np.mean(all_precision_fr, axis = 0).mean())
    print("Total recall:")
    print(np.mean(all_recall_fr, axis = 0).mean())
    print("Total F1 score:")
    print(np.mean(all_f1_fr, axis = 0).mean())

    print("----Overall----")
    print("Precision:")
    print((np.mean(all_precision_nl, axis = 0).mean() + np.mean(all_precision_en, axis = 0).mean() + np.mean(all_precision_fr, axis = 0).mean()) / 3)
    print("Recall:")
    print((np.mean(all_recall_nl, axis = 0).mean() + np.mean(all_recall_en, axis = 0).mean() + np.mean(all_recall_fr, axis = 0).mean()) / 3)
    print("F1 score:")
    print((np.mean(all_f1_nl, axis = 0).mean() + np.mean(all_f1_en, axis = 0).mean() + np.mean(all_f1_fr, axis = 0).mean()) / 3)
    
