import os

import numpy as np

import torch

import torch.nn as nn

import torch.optim as optim

from torch.optim import lr_scheduler

from torch.utils.data import Dataset, DataLoader, RandomSampler, SequentialSampler

from torch.nn.utils.rnn import pack_padded_sequence 

import pickle

import json

import matplotlib.pyplot as plt

from glob import glob

import time

import copy

from tqdm import tqdm

from transformers import BartTokenizer, BartForConditionalGeneration, BartConfig, BartForSequenceClassification, BertTokenizer, BertConfig, BertForSequenceClassification, RobertaTokenizer, RobertaForSequenceClassification

from data import ZuCo_dataset

from model_sentiment import BaselineMLPSentence, BaselineLSTM, FineTunePretrainedTwoStep, ZeroShotSentimentDiscovery, JointBrainTranslatorSentimentClassifier

from model_decoding import BrainTranslator, BrainTranslatorNaive

from sklearn.metrics import precision_recall_fscore_support

from sklearn.metrics import accuracy_score

from config import get_config

# Function to calculate the accuracy of our predictions vs labels

def flat_accuracy(preds, labels):

    # preds: numpy array: N * 3 

    # labels: numpy array: N 

    pred_flat = np.argmax(preds, axis=1).flatten()  

    labels_flat = labels.flatten()

    return np.sum(pred_flat == labels_flat) / len(labels_flat)

def flat_accuracy_top_k(preds, labels,k):

    topk_preds = []

    for pred in preds:

        topk = pred.argsort()[-k:][::-1]

        topk_preds.append(list(topk))

    # print(topk_preds)

    topk_preds = list(topk_preds)

    right_count = 0

    # print(len(labels))

    for i in range(len(labels)):

        l = labels[i][0]

        if l in topk_preds[i]:

            right_count+=1

    return right_count/len(labels)

def eval_model(dataloaders, device, model, criterion, optimizer, scheduler, num_epochs=25, tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')):

    def logits2PredString(logits, tokenizer):

        probs = logits[0].softmax(dim = 1)

        # print('probs size:', probs.size())

        values, predictions = probs.topk(1)

        # print('predictions before squeeze:',predictions.size())

        predictions = torch.squeeze(predictions)

        predict_string = tokenizer.decode(predictions)

        return predict_string

    # modified from: https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html

    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())

    best_loss = 100000000000

    best_acc = 0.0

    total_pred_labels = np.array([])

    total_true_labels = np.array([])

    for epoch in range(1):

        print('Epoch {}/{}'.format(epoch, num_epochs - 1))

        print('-' * 10)

        # Each epoch has a training and validation phase

        for phase in ['test']:

            total_accuracy = 0.0

            if phase == 'train':

                model.train()  # Set model to training mode

            else:

                model.eval()   # Set model to evaluate mode

            running_loss = 0.0

            # Iterate over data.

            for input_word_eeg_features, seq_lens, input_masks, input_mask_invert, target_ids, target_mask, sentiment_labels, sent_level_EEG in dataloaders[phase]:

                input_word_eeg_features = input_word_eeg_features.to(device).float()

                input_masks = input_masks.to(device)

                input_mask_invert = input_mask_invert.to(device)

                sent_level_EEG = sent_level_EEG.to(device)

                sentiment_labels = sentiment_labels.to(device)

                target_ids = target_ids.to(device)

                target_mask = target_mask.to(device)

                ## forward ###################

                if isinstance(model, BaselineMLPSentence):

                    logits = model(sent_level_EEG) # before softmax

                    # calculate loss

                    loss = criterion(logits, sentiment_labels)

                elif isinstance(model, BaselineLSTM):

                    x_packed = pack_padded_sequence(input_word_eeg_features, seq_lens, batch_first=True, enforce_sorted=False)

                    logits = model(x_packed)

                    # calculate loss

                    loss = criterion(logits, sentiment_labels)

                elif isinstance(model, BertForSequenceClassification) or isinstance(model, RobertaForSequenceClassification) or isinstance(model, BartForSequenceClassification):

                    output = model(input_ids = target_ids, attention_mask = target_mask, return_dict = True, labels = sentiment_labels)

                    logits = output.logits

                    loss = output.loss

                elif isinstance(model, FineTunePretrainedTwoStep):

                    output = model(input_word_eeg_features, input_masks, input_mask_invert, sentiment_labels)

                    logits = output.logits

                    loss = output.loss

                elif isinstance(model, ZeroShotSentimentDiscovery):    

                    print()

                    print('target string:',tokenizer.decode(target_ids[0]).replace('<pad>','').split('</s>')[0]) 

                    """replace padding ids in target_ids with -100"""

                    target_ids[target_ids == tokenizer.pad_token_id] = -100 

                    output = model(input_word_eeg_features, input_masks, input_mask_invert, target_ids, sentiment_labels)

                    logits = output.logits

                    loss = output.loss

                elif isinstance(model, JointBrainTranslatorSentimentClassifier):

                    print()

                    print('target string:',tokenizer.decode(target_ids[0]).replace('<pad>','').split('</s>')[0]) 

                    """replace padding ids in target_ids with -100"""

                    target_ids[target_ids == tokenizer.pad_token_id] = -100 

                    LM_output, classification_output = model(input_word_eeg_features, input_masks, input_mask_invert, target_ids, sentiment_labels)

                    LM_logits = LM_output.logits

                    print('pred string:', logits2PredString(LM_logits, tokenizer).split('</s></s>')[0].replace('<s>',''))

                    classification_loss = classification_output['loss']

                    logits = classification_output['logits']

                    loss = classification_loss 

                ###############################

                # backward + optimize only if in training phase

                if phase == 'train':

                    # with torch.autograd.detect_anomaly():

                    loss.backward()

                    optimizer.step()

                # calculate accuracy

                preds_cpu = logits.detach().cpu().numpy()

                label_cpu = sentiment_labels.cpu().numpy()

                total_accuracy += flat_accuracy(preds_cpu, label_cpu)

                # add to total pred and label array, for cal F1, precision, recall

                pred_flat = np.argmax(preds_cpu, axis=1).flatten()

                labels_flat = label_cpu.flatten()

                total_pred_labels = np.concatenate((total_pred_labels,pred_flat))

                total_true_labels = np.concatenate((total_true_labels,labels_flat))

                # statistics

                running_loss += loss.item() * sent_level_EEG.size()[0] # batch loss

                # print('[DEBUG]loss:',loss.item())

                # print('#################################')

            if phase == 'train':

                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]

            epoch_acc = total_accuracy / len(dataloaders[phase])

            print('{} Loss: {:.4f}'.format(phase, epoch_loss))

            print('{} Acc: {:.4f}'.format(phase, epoch_acc))

            # deep copy the model

            if phase == 'test' and epoch_loss < best_loss:

                best_loss = epoch_loss

                best_acc = epoch_acc

        print()

    time_elapsed = time.time() - since

    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))

    print('Best test loss: {:4f}'.format(best_loss))

    print('Best test acc: {:4f}'.format(best_acc))

    print()

    print('test sample num:', len(total_pred_labels))

    print('total preds:',total_pred_labels)

    print('total truth:',total_true_labels)

    print('sklearn macro: precision, recall, F1:')

    print(precision_recall_fscore_support(total_true_labels, total_pred_labels, average='macro'))

    print()

    print('sklearn micro: precision, recall, F1:')

    print(precision_recall_fscore_support(total_true_labels, total_pred_labels, average='micro'))

    print()

    print('sklearn accuracy:')

    print(accuracy_score(total_true_labels,total_pred_labels))

    print()

if __name__ == '__main__':

    args = get_config('eval_sentiment')

    ''' config param'''

    num_epochs = 1

    dataset_setting = 'unique_sent'

    '''model name'''

    # model_name = 'BaselineMLP'

    # model_name = 'BaselineLSTM'

    # model_name = 'NaiveFinetuneBert'

    # model_name = 'FinetunedBertOnText'

    # model_name = 'FinetunedRoBertaOnText'

    # model_name = 'FinetunedBartOnText'

    # model_name = 'ZeroShotSentimentDiscovery'

    model_name = args['model_name']

    print(f'[INFO] eval {model_name}')

    if model_name == 'ZeroShotSentimentDiscovery':

        '''load decoder and classifier config'''

        config_decoder = json.load(open(args['decoder_config_path']))

        config_classifier = json.load(open(args['classifier_config_path']))

        '''choose generator'''

        # decoder_name = 'BrainTranslator'

        # decoder_name = 'BrainTranslatorNaive'

        decoder_name = config_decoder['model_name']

        decoder_checkpoint = args['decoder_checkpoint_path']

        print(f'[INFO] using decoder: {decoder_name}')

        '''choose classifier'''

        # pretrain_Bert, pretrain_RoBerta, pretrain_Bart

        classifier_name = config_classifier['model_name']

        classifier_checkpoint = args['classifier_checkpoint_path']

        print(f'[INFO] using classifier: {classifier_name}')

    else:

        checkpoint_path = args['checkpoint_path']

        print('[INFO] loading baseline:', checkpoint_path)

    batch_size = 1

    # subject_choice = 'ALL

    subject_choice = args['subjects']

    print(f'![Debug]using {subject_choice}')

    # eeg_type_choice = 'GD

    eeg_type_choice = args['eeg_type']

    print(f'[INFO]eeg type {eeg_type_choice}')

    # bands_choice = ['_t1'] 

    # bands_choice = ['_t1','_t2','_a1','_a2','_b1','_b2','_g1','_g2'] 

    bands_choice = args['eeg_bands']

    print(f'[INFO]using bands {bands_choice}')

    ''' set random seeds '''

    seed_val = 312

    np.random.seed(seed_val)

    torch.manual_seed(seed_val)

    torch.cuda.manual_seed_all(seed_val)

    ''' set up device '''

    # use cuda

    if torch.cuda.is_available():  

        dev = args['cuda']

    else:  

        dev = "cpu"

    # CUDA_VISIBLE_DEVICES=0,1,2,3  

    device = torch.device(dev)

    print(f'[INFO]using device {dev}')

    ''' load pickle'''

    whole_dataset_dict = []

    dataset_path_task1 = './dataset/ZuCo/task1-SR/pickle/task1-SR-dataset.pickle' 

    with open(dataset_path_task1, 'rb') as handle:

        whole_dataset_dict.append(pickle.load(handle))

    '''set up tokenizer'''

    if model_name in ['BaselineMLP','BaselineLSTM', 'NaiveFinetuneBert', 'FinetunedBertOnText']:

        print('[INFO]using Bert tokenizer')

        tokenizer = BertTokenizer.from_pretrained('bert-base-cased')

    elif model_name == 'FinetunedBartOnText':

        print('[INFO]using Bart tokenizer')

        tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')

    elif model_name == 'FinetunedRoBertaOnText':

        print('[INFO]using RoBerta tokenizer')

        tokenizer =  RobertaTokenizer.from_pretrained('roberta-base')

    elif model_name == 'ZeroShotSentimentDiscovery':

        decoder_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large') # Bart

        tokenizer = decoder_tokenizer

        if classifier_name == 'pretrain_Bert':

            sentiment_tokenizer = BertTokenizer.from_pretrained('bert-base-cased') # Bert

        elif classifier_name == 'pretrain_Bart':

            sentiment_tokenizer = decoder_tokenizer

        elif classifier_name == 'pretrain_RoBerta':

            sentiment_tokenizer = RobertaTokenizer.from_pretrained('roberta-base')

    ''' set up model '''

    if model_name == 'BaselineMLP':

        print('[INFO]Model: BaselineMLP')

        model = BaselineMLPSentence(input_dim = 840, hidden_dim = 128, output_dim = 3)

    elif model_name == 'BaselineLSTM':

        print('[INFO]Model: BaselineLSTM')

        # model = BaselineLSTM(input_dim = 840, hidden_dim = 256, output_dim = 3, num_layers = 1)

        model = BaselineLSTM(input_dim = 840, hidden_dim = 256, output_dim = 3, num_layers = 4)

    elif model_name == 'FinetunedBertOnText':

        print('[INFO]Model: FinetunedBertOnText')

        model = BertForSequenceClassification.from_pretrained('bert-base-cased',num_labels=3)

    elif model_name == 'FinetunedRoBertaOnText':

        print('[INFO]Model: FinetunedRoBertaOnText')

        model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=3)

    elif model_name == 'FinetunedBartOnText':

        print('[INFO]Model: FinetunedBartOnText')

        model = BartForSequenceClassification.from_pretrained('facebook/bart-large', num_labels=3)

    elif model_name == 'ZeroShotSentimentDiscovery':

        print(f'[INFO]Model: ZeroShotSentimentDiscovery, using classifer:{classifier_name}, using generator: {decoder_name}')

        pretrained = BartForConditionalGeneration.from_pretrained('facebook/bart-large')

        if decoder_name == 'BrainTranslator':

            decoder = BrainTranslator(pretrained, in_feature = 105*len(bands_choice), decoder_embedding_size = 1024, additional_encoder_nhead=8, additional_encoder_dim_feedforward = 2048)

        elif decoder_name == 'BrainTranslatorNaive':

            decoder = BrainTranslatorNaive(pretrained, in_feature = 105*len(bands_choice), decoder_embedding_size = 1024, additional_encoder_nhead=8, additional_encoder_dim_feedforward = 2048)

        decoder.load_state_dict(torch.load(decoder_checkpoint))

        if classifier_name == 'pretrain_Bert':

            classifier = BertForSequenceClassification.from_pretrained('bert-base-cased',num_labels=3)

        elif classifier_name == 'pretrain_Bart':

            classifier = BartForSequenceClassification.from_pretrained('facebook/bart-large', num_labels=3)

        elif classifier_name == 'pretrain_RoBerta':

            classifier = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=3)

        classifier.load_state_dict(torch.load(classifier_checkpoint))

        model = ZeroShotSentimentDiscovery(decoder, classifier, decoder_tokenizer, sentiment_tokenizer, device = device)

        model.to(device)

    if model_name != 'ZeroShotSentimentDiscovery':

        # load model and send to device

        model.load_state_dict(torch.load(checkpoint_path))

        model.to(device)

    ''' set up dataloader '''

    # test dataset

    test_set = ZuCo_dataset(whole_dataset_dict, 'test', tokenizer, subject = subject_choice, eeg_type = eeg_type_choice, bands = bands_choice, setting = 'unique_sent')

    dataset_sizes = {'test': len(test_set)}

    # print('[INFO]train_set size: ', len(train_set))

    print('[INFO]test_set size: ', len(test_set))

    test_dataloader = DataLoader(test_set, batch_size = 1, shuffle=False, num_workers=4)

    # dataloaders

    dataloaders = {'test':test_dataloader}

    ''' set up optimizer and scheduler'''

    optimizer_step1 = None

    exp_lr_scheduler_step1 = None

    ''' set up loss function '''

    criterion = nn.CrossEntropyLoss()

    print('=== start training ... ===')

    # return best loss model from step1 training

    model = eval_model(dataloaders, device, model, criterion, optimizer_step1, exp_lr_scheduler_step1, num_epochs=num_epochs, tokenizer = tokenizer)