import sys

import torch

import numpy as np

import consts

from datasets import ContrastiveDataset, SimpleDataset

from torch.utils.data import DataLoader

from transformers import AutoModel

from data import process_raw_data

from SNN_model import BERT_Arch, SiameseNeuralNetwork

from SNN_training import train_siamese_network

from projection import construct_train_matrix, extract_prototypes, project_to_dissimilarity_space

from SVM_model import ensemble_of_classifiers

from sklearn.model_selection import train_test_split

from transformers import BertTokenizerFast

def predict(projected_test, classifiers_list, categories_order):

    pred_y = []

    for classifier in classifiers_list:

         pred_y.append(classifier.predict_proba(projected_test)[:,1])    # predict_proba returns probabiltiy for class==0 and for class==1, so we take only the probabilities of class==1

    pred_y = np.vstack(pred_y) # (num_classifiers, num_samples_test)

    highest_predictions = categories_order[np.argmax(pred_y, axis=0)]

    print(pred_y)

    print(highest_predictions)

    return highest_predictions

if __name__ == '__main__':

    data_path = sys.argv[1]

    # ------------- Data --------------------------

    data, test_unseen_categories= process_raw_data(data_path)

    train_text, temp_text, train_labels, temp_labels = train_test_split(data['description'], data['labels'],

                                                                        random_state=42,

                                                                        test_size=0.3,

                                                                        stratify=data['labels'])

    val_text, test_text, val_labels, test_labels = train_test_split(temp_text, temp_labels,

                                                                    random_state=42,

                                                                    test_size=0.5,

                                                                    stratify=temp_labels)

    unseen_train_text, unseen_test_text, unseen_train_labels, unseen_test_labels = train_test_split(

        test_unseen_categories['description'], test_unseen_categories['labels'],

        random_state=42,

        test_size=0.2,

        stratify=test_unseen_categories['labels'])

    # Tokinization

    # Load the BERT tokenizer

    model_name = consts.model_name

    tokenizer = BertTokenizerFast.from_pretrained(model_name)

    # tokenize and encode sequences in the sets set

    texts = [train_text, val_text, test_text, unseen_train_text, unseen_test_text]

    tokens_texts = []

    for text in texts:

        tokens_texts.append(

            tokenizer.batch_encode_plus(text.tolist(), max_length=consts.MAX_SENTENCE_LENGTh, padding='max_length',

                                        truncation=True))

    train_tokinized, val_tokinized, test_tokinized, unseen_train_tokinized, unseen_test_tokinized = tokens_texts

    def convert_to_tensors(data, labels):

        seq = torch.tensor(data['input_ids'])

        mask = torch.tensor(data['attention_mask'])

        y = torch.tensor(labels.tolist())

        return seq, mask, y

    train_seq, train_mask, train_y = convert_to_tensors(train_tokinized, train_labels)

    val_seq, val_mask, val_y = convert_to_tensors(val_tokinized, val_labels)

    test_seq, test_mask, test_y = convert_to_tensors(test_tokinized, test_labels)

    unseen_train_seq, unseen_train_mask, unseen_train_y = convert_to_tensors(unseen_train_tokinized, unseen_train_labels)

    unseen_test_seq, unseen_test_mask, unseen_test_y = convert_to_tensors(unseen_test_tokinized, unseen_test_labels)

    train_set = ContrastiveDataset(train_seq, train_mask, train_y)

    val_set = ContrastiveDataset(val_seq, val_mask, val_y)

    test_set = ContrastiveDataset(test_seq, test_mask, test_y)

    train_set_simple = SimpleDataset(train_seq, train_mask, train_y)

    test_set_simple = SimpleDataset(test_seq, test_mask, test_y)

    unseen_train_set_simple = SimpleDataset(unseen_train_seq, unseen_train_mask, unseen_train_y)

    unseen_test_set_simple = SimpleDataset(unseen_test_seq, unseen_test_mask, unseen_test_y)

    trainLoader = DataLoader(train_set, batch_size=32, shuffle=True, drop_last=False, num_workers=0)

    valLoader = DataLoader(val_set, batch_size=32, shuffle=True, drop_last=False, num_workers=0)

    testLoader = DataLoader(test_set, batch_size=10, shuffle=False, drop_last=False, num_workers=0)

    trainLoader_simple = DataLoader(train_set_simple, batch_size=32, shuffle=False, drop_last=False, num_workers=0)

    testLoader_simple = DataLoader(test_set_simple, batch_size=64, shuffle=False, drop_last=False, num_workers=0)

    unseen_trainLoader_simple = DataLoader(unseen_train_set_simple, batch_size=64, shuffle=False, drop_last=False,

                                           num_workers=0)

    unseen_testLoader_simple = DataLoader(unseen_test_set_simple, batch_size=64, shuffle=False, drop_last=False,

                                          num_workers=0)

    # -------------- Parametrs --------------------

    model_name = consts.model_name

    # ------------- Train SNN --------------------------

    # specify GPU

    # Get cpu or gpu device for training.

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

    print("Using {} device".format(device))

    # import BERT-base pretrained model

    bert = AutoModel.from_pretrained(

        model_name)  # ('bert-base-uncased') 'microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext'

    # freeze all the parameters

    for param in bert.parameters():

        param.requires_grad = False

    # pass the pre-trained BERT to our define architecture

    bert_arch = BERT_Arch(bert)

    SNN_model = SiameseNeuralNetwork(bert_arch).to(device)

    num_epochs = 30

    train_loss_history, val_loss_history, similarities_list = train_siamese_network(SNN_model,

                                                                                    dataloaders={"train": trainLoader,

                                                                                                 "val": valLoader},

                                                                                    num_epochs=num_epochs,

                                                                                    device=device)

    non_matching_similarity, matching_similarity, val_non_matching_similarity, val_matching_similarity = similarities_list

    # ----------------- Prototypes Selection ------------------------

    train_matrix = construct_train_matrix(SNN_model, trainLoader_simple)

    prototypes_list = extract_prototypes(100, trainLoader_simple, train_labels, train_matrix)

    # ---------------- Data Projection ------------------------------

    projected_train = project_to_dissimilarity_space(trainLoader_simple, SNN_model, prototypes_list)

    # ----------------- SVM Ensemble -----------------------------------

    classifiers, categories_order = ensemble_of_classifiers(projected_train, train_labels)

    # ------------------ Test: Seen categories ---------------------

    projected_test = project_to_dissimilarity_space(testLoader_simple, SNN_model, prototypes_list)

    preds = predict(projected_test, classifiers, categories_order)

    # ------------------ Test: Unseen categories ---------------------

    unseen_train_matrix = construct_train_matrix(SNN_model, unseen_trainLoader_simple)

    unseen_prototypes_list = extract_prototypes(100, unseen_trainLoader_simple, unseen_train_labels,

                                                unseen_train_matrix)

    unseen_projected_train = project_to_dissimilarity_space(unseen_trainLoader_simple, SNN_model,

                                                            unseen_prototypes_list)

    unseen_classifiers, unseen_categories_order = ensemble_of_classifiers(unseen_projected_train, unseen_train_labels)

    unseen_projected_test = project_to_dissimilarity_space(unseen_testLoader_simple, SNN_model, unseen_prototypes_list)

    unseen_preds = predict(unseen_projected_test, unseen_classifiers, unseen_categories_order)