import os

import numpy as np

import torch

import pickle

from torch.utils.data import Dataset, DataLoader

import json

import matplotlib.pyplot as plt

from glob import glob

from transformers import BartTokenizer, BertTokenizer

from tqdm import tqdm

from fuzzy_match import match

from fuzzy_match import algorithims

# macro

#ZUCO_SENTIMENT_LABELS = json.load(open('./dataset/ZuCo/task1-SR/sentiment_labels/sentiment_labels.json'))

#SST_SENTIMENT_LABELS = json.load(open('./dataset/stanfordsentiment/ternary_dataset.json'))

def normalize_1d(input_tensor):

    # normalize a 1d tensor

    mean = torch.mean(input_tensor)

    std = torch.std(input_tensor)

    input_tensor = (input_tensor - mean)/std

    return input_tensor 

def get_input_sample(sent_obj, tokenizer, eeg_type = 'GD', bands = ['_t1','_t2','_a1','_a2','_b1','_b2','_g1','_g2'], max_len = 56, add_CLS_token = False):

    def get_word_embedding_eeg_tensor(word_obj, eeg_type, bands):

        frequency_features = []

        for band in bands:

            frequency_features.append(word_obj['word_level_EEG'][eeg_type][eeg_type+band])

        word_eeg_embedding = np.concatenate(frequency_features)

        if len(word_eeg_embedding) != 105*len(bands):

            print(f'expect word eeg embedding dim to be {105*len(bands)}, but got {len(word_eeg_embedding)}, return None')

            return None

        # assert len(word_eeg_embedding) == 105*len(bands)

        return_tensor = torch.from_numpy(word_eeg_embedding)

        return normalize_1d(return_tensor)

    def get_sent_eeg(sent_obj, bands):

        sent_eeg_features = []

        for band in bands:

            key = 'mean'+band

            sent_eeg_features.append(sent_obj['sentence_level_EEG'][key])

        sent_eeg_embedding = np.concatenate(sent_eeg_features)

        assert len(sent_eeg_embedding) == 105*len(bands)

        return_tensor = torch.from_numpy(sent_eeg_embedding)

        return normalize_1d(return_tensor)

    if sent_obj is None:

        # print(f'  - skip bad sentence')   

        return None

    input_sample = {}

    # get target label

    target_string = sent_obj['content']

    target_tokenized = tokenizer(target_string, padding='max_length', max_length=max_len, truncation=True, return_tensors='pt', return_attention_mask = True)

    input_sample['target_ids'] = target_tokenized['input_ids'][0]

    # get sentence level EEG features

    sent_level_eeg_tensor = get_sent_eeg(sent_obj, bands)

    if torch.isnan(sent_level_eeg_tensor).any():

        # print('[NaN sent level eeg]: ', target_string)

        return None

    input_sample['sent_level_EEG'] = sent_level_eeg_tensor

    # get sentiment label

    # handle some wierd case

    if 'emp11111ty' in target_string:

        target_string = target_string.replace('emp11111ty','empty')

    if 'film.1' in target_string:

        target_string = target_string.replace('film.1','film.')

    #if target_string in ZUCO_SENTIMENT_LABELS:

    #    input_sample['sentiment_label'] = torch.tensor(ZUCO_SENTIMENT_LABELS[target_string]+1) # 0:Negative, 1:Neutral, 2:Positive

    #else:

    #    input_sample['sentiment_label'] = torch.tensor(-100) # dummy value

    input_sample['sentiment_label'] = torch.tensor(-100) # dummy value

    # get input embeddings

    word_embeddings = []

    """add CLS token embedding at the front"""

    if add_CLS_token:

        word_embeddings.append(torch.ones(105*len(bands)))

    for word in sent_obj['word']:

        # add each word's EEG embedding as Tensors

        word_level_eeg_tensor = get_word_embedding_eeg_tensor(word, eeg_type, bands = bands)

        # check none, for v2 dataset

        if word_level_eeg_tensor is None:

            return None

        # check nan:

        if torch.isnan(word_level_eeg_tensor).any():

            # print()

            # print('[NaN ERROR] problem sent:',sent_obj['content'])

            # print('[NaN ERROR] problem word:',word['content'])

            # print('[NaN ERROR] problem word feature:',word_level_eeg_tensor)

            # print()

            return None

        word_embeddings.append(word_level_eeg_tensor)

    # pad to max_len

    while len(word_embeddings) < max_len:

        word_embeddings.append(torch.zeros(105*len(bands)))

    input_sample['input_embeddings'] = torch.stack(word_embeddings) # max_len * (105*num_bands)

    # mask out padding tokens

    input_sample['input_attn_mask'] = torch.zeros(max_len) # 0 is masked out

    if add_CLS_token:

        input_sample['input_attn_mask'][:len(sent_obj['word'])+1] = torch.ones(len(sent_obj['word'])+1) # 1 is not masked

    else:

        input_sample['input_attn_mask'][:len(sent_obj['word'])] = torch.ones(len(sent_obj['word'])) # 1 is not masked

    # mask out padding tokens reverted: handle different use case: this is for pytorch transformers

    input_sample['input_attn_mask_invert'] = torch.ones(max_len) # 1 is masked out

    if add_CLS_token:

        input_sample['input_attn_mask_invert'][:len(sent_obj['word'])+1] = torch.zeros(len(sent_obj['word'])+1) # 0 is not masked

    else:

        input_sample['input_attn_mask_invert'][:len(sent_obj['word'])] = torch.zeros(len(sent_obj['word'])) # 0 is not masked

    # mask out target padding for computing cross entropy loss

    input_sample['target_mask'] = target_tokenized['attention_mask'][0]

    input_sample['seq_len'] = len(sent_obj['word'])

    # clean 0 length data

    if input_sample['seq_len'] == 0:

        print('discard length zero instance: ', target_string)

        return None

    return input_sample

class ZuCo_dataset(Dataset):

    def __init__(self, input_dataset_dicts, phase, tokenizer, subject = 'ALL', eeg_type = 'GD', bands = ['_t1','_t2','_a1','_a2','_b1','_b2','_g1','_g2'], setting = 'unique_sent', is_add_CLS_token = False):

        self.inputs = []

        self.tokenizer = tokenizer

        if not isinstance(input_dataset_dicts,list):

            input_dataset_dicts = [input_dataset_dicts]

        print(f'[INFO]loading {len(input_dataset_dicts)} task datasets')

        for input_dataset_dict in input_dataset_dicts:

            if subject == 'ALL':

                subjects = list(input_dataset_dict.keys())

                print('[INFO]using subjects: ', subjects)

            else:

                subjects = [subject]

            total_num_sentence = len(input_dataset_dict[subjects[0]])

            train_divider = int(0.8*total_num_sentence)

            dev_divider = train_divider + int(0.1*total_num_sentence)

            print(f'train divider = {train_divider}')

            print(f'dev divider = {dev_divider}')

            if setting == 'unique_sent':

                # take first 80% as trainset, 10% as dev and 10% as test

                if phase == 'train':

                    print('[INFO]initializing a train set...')

                    for key in subjects:

                        for i in range(train_divider):

                            input_sample = get_input_sample(input_dataset_dict[key][i],self.tokenizer,eeg_type,bands = bands, add_CLS_token = is_add_CLS_token)

                            if input_sample is not None:

                                self.inputs.append(input_sample)

                elif phase == 'dev':

                    print('[INFO]initializing a dev set...')

                    for key in subjects:

                        for i in range(train_divider,dev_divider):

                            input_sample = get_input_sample(input_dataset_dict[key][i],self.tokenizer,eeg_type,bands = bands, add_CLS_token = is_add_CLS_token)

                            if input_sample is not None:

                                self.inputs.append(input_sample)

                elif phase == 'test':

                    print('[INFO]initializing a test set...')

                    for key in subjects:

                        for i in range(dev_divider,total_num_sentence):

                            input_sample = get_input_sample(input_dataset_dict[key][i],self.tokenizer,eeg_type,bands = bands, add_CLS_token = is_add_CLS_token)

                            if input_sample is not None:

                                self.inputs.append(input_sample)

            elif setting == 'unique_subj':

                print('WARNING!!! only implemented for SR v1 dataset ')

                # subject ['ZAB', 'ZDM', 'ZGW', 'ZJM', 'ZJN', 'ZJS', 'ZKB', 'ZKH', 'ZKW'] for train

                # subject ['ZMG'] for dev

                # subject ['ZPH'] for test

                if phase == 'train':

                    print(f'[INFO]initializing a train set using {setting} setting...')

                    for i in range(total_num_sentence):

                        for key in ['ZAB', 'ZDM', 'ZGW', 'ZJM', 'ZJN', 'ZJS', 'ZKB', 'ZKH','ZKW']:

                            input_sample = get_input_sample(input_dataset_dict[key][i],self.tokenizer,eeg_type,bands = bands, add_CLS_token = is_add_CLS_token)

                            if input_sample is not None:

                                self.inputs.append(input_sample)

                if phase == 'dev':

                    print(f'[INFO]initializing a dev set using {setting} setting...')

                    for i in range(total_num_sentence):

                        for key in ['ZMG']:

                            input_sample = get_input_sample(input_dataset_dict[key][i],self.tokenizer,eeg_type,bands = bands, add_CLS_token = is_add_CLS_token)

                            if input_sample is not None:

                                self.inputs.append(input_sample)

                if phase == 'test':

                    print(f'[INFO]initializing a test set using {setting} setting...')

                    for i in range(total_num_sentence):

                        for key in ['ZPH']:

                            input_sample = get_input_sample(input_dataset_dict[key][i],self.tokenizer,eeg_type,bands = bands, add_CLS_token = is_add_CLS_token)

                            if input_sample is not None:

                                self.inputs.append(input_sample)

            print('++ adding task to dataset, now we have:', len(self.inputs))

        print('[INFO]input tensor size:', self.inputs[0]['input_embeddings'].size())

        print()

    def __len__(self):

        return len(self.inputs)

    def __getitem__(self, idx):

        input_sample = self.inputs[idx]

        return (

            input_sample['input_embeddings'], 

            input_sample['seq_len'],

            input_sample['input_attn_mask'], 

            input_sample['input_attn_mask_invert'],

            input_sample['target_ids'], 

            input_sample['target_mask'], 

            input_sample['sentiment_label'], 

            input_sample['sent_level_EEG']

        )

        # keys: input_embeddings, input_attn_mask, input_attn_mask_invert, target_ids, target_mask, 

"""for train classifier on stanford sentiment treebank text-sentiment pairs"""

class SST_tenary_dataset(Dataset):

    def __init__(self, ternary_labels_dict, tokenizer, max_len = 56, balance_class = True):

        self.inputs = []

        pos_samples = []

        neg_samples = []

        neu_samples = []

        for key,value in ternary_labels_dict.items():

            tokenized_inputs = tokenizer(key, padding='max_length', max_length=max_len, truncation=True, return_tensors='pt', return_attention_mask = True)

            input_ids = tokenized_inputs['input_ids'][0]

            attn_masks = tokenized_inputs['attention_mask'][0]

            label = torch.tensor(value)

            # count:

            if value == 0:

                neg_samples.append((input_ids,attn_masks,label))

            elif value == 1:

                neu_samples.append((input_ids,attn_masks,label))

            elif value == 2:

                pos_samples.append((input_ids,attn_masks,label))

        print(f'Original distribution:\n\tVery positive: {len(pos_samples)}\n\tNeutral: {len(neu_samples)}\n\tVery negative: {len(neg_samples)}')    

        if balance_class:

            print(f'balance class to {min([len(pos_samples),len(neg_samples),len(neu_samples)])} each...')

            for i in range(min([len(pos_samples),len(neg_samples),len(neu_samples)])):

                self.inputs.append(pos_samples[i])

                self.inputs.append(neg_samples[i])

                self.inputs.append(neu_samples[i])

        else:

            self.inputs = pos_samples + neg_samples + neu_samples

    def __len__(self):

        return len(self.inputs)

    def __getitem__(self, idx):

        input_sample = self.inputs[idx]

        return input_sample

        # keys: input_embeddings, input_attn_mask, input_attn_mask_invert, target_ids, target_mask, 

'''sanity test'''

if __name__ == '__main__':

    check_dataset = 'stanford_sentiment'

    if check_dataset == 'ZuCo':

        whole_dataset_dicts = []

        dataset_path_task1 = '/shared/nas/data/m1/wangz3/SAO_project/SAO/dataset/ZuCo/task1-SR/pickle/task1-SR-dataset-with-tokens_6-25.pickle' 

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

            whole_dataset_dicts.append(pickle.load(handle))

        dataset_path_task2 = '/shared/nas/data/m1/wangz3/SAO_project/SAO/dataset/ZuCo/task2-NR/pickle/task2-NR-dataset-with-tokens_7-10.pickle' 

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

            whole_dataset_dicts.append(pickle.load(handle))

        # dataset_path_task3 = '/shared/nas/data/m1/wangz3/SAO_project/SAO/dataset/ZuCo/task3-TSR/pickle/task3-TSR-dataset-with-tokens_7-10.pickle' 

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

        #     whole_dataset_dicts.append(pickle.load(handle))

        dataset_path_task2_v2 = '/shared/nas/data/m1/wangz3/SAO_project/SAO/dataset/ZuCo/task2-NR-2.0/pickle/task2-NR-2.0-dataset-with-tokens_7-15.pickle' 

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

            whole_dataset_dicts.append(pickle.load(handle))

        print()

        for key in whole_dataset_dicts[0]:

            print(f'task2_v2, sentence num in {key}:',len(whole_dataset_dicts[0][key]))

        print()

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

        dataset_setting = 'unique_sent'

        subject_choice = 'ALL'

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

        eeg_type_choice = 'GD'

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

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

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

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

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

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

        print('trainset size:',len(train_set))

        print('devset size:',len(dev_set))

        print('testset size:',len(test_set))

    elif check_dataset == 'stanford_sentiment':

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

        SST_dataset = SST_tenary_dataset(SST_SENTIMENT_LABELS, tokenizer)

        print('SST dataset size:',len(SST_dataset))

        print(SST_dataset[0])

        print(SST_dataset[1])