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

import pickle

import json

import matplotlib.pyplot as plt

from glob import glob

import time

import copy

from tqdm import tqdm

import re

from transformers import BartTokenizer, BartForConditionalGeneration, BartConfig, BertTokenizer

from data import ZuCo_dataset

from model_decoding import BrainTranslator, BrainTranslatorNaive

from metrics import compute_metrics

from config import get_config

def eval_model(dataloaders, device, tokenizer, criterion, model, output_all_results_path='./results/temp.txt'):

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

    model.eval()  # Set model to evaluate mode

    running_loss = 0.0

    # Iterate over data.

    # sample_count = 0

    target_tokens_list = []

    target_string_list = []

    pred_tokens_list = []

    pred_string_list = []

    with open(output_all_results_path, 'w') as f:

        # count=0

        for input_embeddings, seq_len, input_masks, input_mask_invert, target_ids, target_mask, sentiment_labels, sent_level_EEG in \

        dataloaders['test']:

            # count+=1

            # if count>5:

            #     break

            # load in batch

            input_embeddings_batch = input_embeddings.to(device).float()

            input_masks_batch = input_masks.to(device)

            target_ids_batch = target_ids.to(device)

            input_mask_invert_batch = input_mask_invert.to(device)

            if intput_noise:

                input_embeddings_batch=torch.rand_like(input_embeddings_batch)

            # target_tokens = tokenizer.convert_ids_to_tokens(target_ids_batch, skip_special_tokens = True)

            target_string = tokenizer.batch_decode(target_ids_batch, skip_special_tokens=True)

            # print('target ids tensor:',target_ids_batch[0])

            # print('target ids:',target_ids_batch[0].tolist())

            # print('+' * 100)

            # print('target tokens:',target_tokens)

            # print('target string:', target_string)

            # add to list for later calculate bleu metric

            # target_tokens_list.append([target_tokens])

            target_string_list.extend(target_string)

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

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

            if not teacher_forcing:

                predictions = model.generate(input_embeddings_batch, input_masks_batch, input_mask_invert_batch,

                                             target_ids_batch,

                                             max_length=100,

                                             num_beams=5, do_sample=False, repetition_penalty=5.0,

                                             # num_beams=5,encoder_no_repeat_ngram_size =1,

                                             # do_sample=True, top_k=15,temperature=0.5,num_return_sequences=5,

                                             # early_stopping=True

                                             )

            # predicted_string=predicted_string.squeeze()

            # print(f'predictions:{predictions}')

            # print(f'predicted_string:{predicted_string}')

            #

            # print(f'predicted_string:{predicted_string}')

            else:

                seq2seqLMoutput = model(input_embeddings_batch, input_masks_batch, input_mask_invert_batch,

                                        target_ids_batch)

                logits = seq2seqLMoutput.logits  # bs*seq_len*voc_sz

                probs = logits.softmax(dim=-1)

                values, predictions = probs.topk(1)

                predictions = torch.squeeze(predictions, dim=-1)

                # print(f'predictions:{predictions} predictions shape:{predictions.shape}')

            predicted_string = tokenizer.batch_decode(predictions, skip_special_tokens=True, )

            # print(f'predicted_string:{predicted_string}')

            # start = predicted_string.find("[CLS]") + len("[CLS]")

            # end = predicted_string.find("[SEP]")

            # predicted_string = predicted_string[start:end]

            # predicted_string=merge_consecutive_duplicates(predicted_string,'。')

            # predictions=tokenizer.encode(predicted_string)

            for str_id in range(len(target_string)):

                f.write(f'start################################################\n')

                f.write(f'Predicted: {predicted_string[str_id]}\n')

                f.write(f'True: {target_string[str_id]}\n')

                f.write(f'end################################################\n\n\n')

            # convert to int list

            # predictions = predictions.tolist()

            # truncated_prediction = []

            # for t in predictions:

            #     if t != tokenizer.eos_token_id:

            #         truncated_prediction.append(t)

            #     else:

            #         break

            # pred_tokens = tokenizer.convert_ids_to_tokens(truncated_prediction, skip_special_tokens = True)

            # pred_tokens_list.append(pred_tokens)

            pred_string_list.extend(predicted_string)

            # sample_count += 1

            # print('predicted tokens:',pred_tokens)

            # print('predicted string:',predicted_string)

            # print('-' * 100)

    # print(f'pred_string_list:{pred_string_list}')

    # print(f'target_string_list:{target_string_list}')

    metrics_results=compute_metrics(pred_string_list,target_string_list)

    print(f'teacher_forcing{teacher_forcing} intput_noise{intput_noise}')

    print(metrics_results)

    print(output_all_results_path)

    print(output_all_metrics_results_path)

    with open(output_all_metrics_results_path, "w") as json_file:

        json.dump(metrics_results, json_file, indent=4, ensure_ascii=False)

if __name__ == '__main__':

    home_directory = os.path.expanduser("~")

    ''' get args'''

    args = get_config('eval_decoding')

    ''' load training config'''

    training_config = json.load(open(args['config_path']))

    batch_size = 1

    subject_choice = training_config['subjects']

    print(f'[INFO]subjects: {subject_choice}')

    eeg_type_choice = training_config['eeg_type']

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

    bands_choice = training_config['eeg_bands']

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

    dataset_setting = 'unique_sent'

    task_name = training_config['task_name']

    model_name = training_config['model_name']

    # model_name = 'BrainTranslator'

    # model_name = 'BrainTranslatorNaive'

    # teacher_forcing = True

    # {'wer': 0.7980769276618958, 'rouge1_fmeasure': 23.912235260009766, 'rouge1_precision': 24.66936492919922, 'rouge1_recall': 23.318071365356445, 'rouge2_fmeasure': 6.851282119750977, 'rouge2_precision': 6.962162017822266, 'rouge2_recall': 6.751219272613525, 'rougeL_fmeasure': 22.912235260009766, 'rougeL_precision': 23.61673355102539, 'rougeL_recall': 22.36568832397461, 'rougeLsum_fmeasure': 22.912235260009766, 'rougeLsum_precision': 23.61673355102539, 'rougeLsum_recall': 22.36568832397461, 'bleu-1': 0.23883000016212463, 'bleu-2': 0.13888777792453766, 'bleu-3': 0.0, 'bleu-4': 0.0}

    teacher_forcing = eval(args['tf'])

    intput_noise = eval(args['noise'])

    print(f'teacher_forcing{teacher_forcing} intput_noise{intput_noise}')

    output_all_results_path = (f'./results/{task_name}-{model_name}{"-teacher_forcing" if teacher_forcing else ""}{"-intput_noise" if intput_noise else ""}-all_decoding_results.txt')

    output_all_metrics_results_path = output_all_results_path.replace('txt', 'json')

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

    ''' set up dataloader '''

    whole_dataset_dicts = []

    if 'task1' in task_name:

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

        dataset_path_task1=os.path.join(home_directory,dataset_path_task1)

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

            whole_dataset_dicts.append(pickle.load(handle))

    if 'task2' in task_name:

        dataset_path_task2 = 'datasets/ZuCo/task2-NR/pickle/task2-NR-dataset.pickle'

        dataset_path_task2=os.path.join(home_directory,dataset_path_task2)

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

            whole_dataset_dicts.append(pickle.load(handle))

    if 'task3' in task_name:

        dataset_path_task3 = 'datasets/ZuCo/task3-TSR/pickle/task3-TSR-dataset.pickle'

        dataset_path_task3=os.path.join(home_directory,dataset_path_task3)

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

            whole_dataset_dicts.append(pickle.load(handle))

    if 'taskNRv2' in task_name:

        dataset_path_taskNRv2 = 'datasets/ZuCo/task2-NR-2.0/pickle/task2-NR-2.0-dataset.pickle'

        dataset_path_taskNRv2=os.path.join(home_directory,dataset_path_taskNRv2)

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

            whole_dataset_dicts.append(pickle.load(handle))

    print()

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

    # test dataset

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

    dataset_sizes = {"test_set":len(test_set)}

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

    # dataloaders

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

    dataloaders = {'test':test_dataloader}

    ''' set up model '''

    checkpoint_path = args['checkpoint_path']

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

    if model_name == 'BrainTranslator':

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

    elif model_name == 'BrainTranslatorNaive':

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

    model.load_state_dict(torch.load(checkpoint_path))

    model.to(device)

    criterion = nn.CrossEntropyLoss()

    ''' eval '''

    eval_model(dataloaders, device, tokenizer, criterion, model, output_all_results_path = output_all_results_path)