import numpy as np

import re

eeg_float_resolution=np.float16

Alpha_ffd_names = ['FFD_a1', 'FFD_a1_diff', 'FFD_a2', 'FFD_a2_diff']

Beta_ffd_names = ['FFD_b1', 'FFD_b1_diff', 'FFD_b2', 'FFD_b2_diff']

Gamma_ffd_names = ['FFD_g1', 'FFD_g1_diff', 'FFD_g2', 'FFD_g2_diff']

Theta_ffd_names = ['FFD_t1', 'FFD_t1_diff', 'FFD_t2', 'FFD_t2_diff']

Alpha_gd_names = ['GD_a1', 'GD_a1_diff', 'GD_a2', 'GD_a2_diff']

Beta_gd_names = ['GD_b1', 'GD_b1_diff', 'GD_b2', 'GD_b2_diff']

Gamma_gd_names = ['GD_g1', 'GD_g1_diff', 'GD_g2', 'GD_g2_diff']

Theta_gd_names = ['GD_t1', 'GD_t1_diff', 'GD_t2', 'GD_t2_diff']

Alpha_gpt_names = ['GPT_a1', 'GPT_a1_diff', 'GPT_a2', 'GPT_a2_diff']

Beta_gpt_names = ['GPT_b1', 'GPT_b1_diff', 'GPT_b2', 'GPT_b2_diff']

Gamma_gpt_names = ['GPT_g1', 'GPT_g1_diff', 'GPT_g2', 'GPT_g2_diff']

Theta_gpt_names = ['GPT_t1', 'GPT_t1_diff', 'GPT_t2', 'GPT_t2_diff']

Alpha_sfd_names = ['SFD_a1', 'SFD_a1_diff', 'SFD_a2', 'SFD_a2_diff']

Beta_sfd_names = ['SFD_b1', 'SFD_b1_diff', 'SFD_b2', 'SFD_b2_diff']

Gamma_sfd_names = ['SFD_g1', 'SFD_g1_diff', 'SFD_g2', 'SFD_g2_diff']

Theta_sfd_names = ['SFD_t1', 'SFD_t1_diff', 'SFD_t2', 'SFD_t2_diff']

Alpha_trt_names = ['TRT_a1', 'TRT_a1_diff', 'TRT_a2', 'TRT_a2_diff']

Beta_trt_names = ['TRT_b1', 'TRT_b1_diff', 'TRT_b2', 'TRT_b2_diff']

Gamma_trt_names = ['TRT_g1', 'TRT_g1_diff', 'TRT_g2', 'TRT_g2_diff']

Theta_trt_names = ['TRT_t1', 'TRT_t1_diff', 'TRT_t2', 'TRT_t2_diff']

# IF YOU CHANGE THOSE YOU MUST ALSO CHANGE CONSTANTS

Alpha_features = Alpha_ffd_names + Alpha_gd_names + Alpha_gpt_names + Alpha_trt_names# + Alpha_sfd_names

Beta_features = Beta_ffd_names + Beta_gd_names + Beta_gpt_names + Beta_trt_names# + Beta_sfd_names

Gamma_features = Gamma_ffd_names + Gamma_gd_names + Gamma_gpt_names + Gamma_trt_names# + Gamma_sfd_names

Theta_features = Theta_ffd_names + Theta_gd_names + Theta_gpt_names + Theta_trt_names# + Theta_sfd_names

# print(Alpha_features)

# GD_EEG_feautres

def extract_all_fixations(data_container, word_data_object, float_resolution = np.float16):

    """

    Extracts all fixations from a word data object

    :param data_container:      (h5py)  Container of the whole data, h5py object

    :param word_data_object:    (h5py)  Container of fixation objects, h5py object

    :param float_resolution:    (type)  Resolution to which data re to be converted, used for data compression

    :return:

        fixations_data  (list)  Data arrays representing each fixation

    """

    word_data = data_container[word_data_object]

    fixations_data = []

    if len(word_data.shape) > 1:

        for fixation_idx in range(word_data.shape[0]):

            fixations_data.append(np.array(data_container[word_data[fixation_idx][0]]).astype(float_resolution))

    return fixations_data

def is_real_word(word):

    """

    Check if the word is a real word

    :param word:    (str)   word string

    :return:

        is_word (bool)  True if it is a real word

    """

    is_word = re.search('[a-zA-Z0-9]', word)

    return is_word

def load_matlab_string(matlab_extracted_object):

    """

    Converts a string loaded from h5py into a python string

    :param matlab_extracted_object:     (h5py)  matlab string object

    :return:

        extracted_string    (str)   translated string

    """

    extracted_string = u''.join(chr(c[0]) for c in matlab_extracted_object)

    return extracted_string

def extract_word_level_data(data_container, word_objects, eeg_float_resolution = np.float16):

    """

    Extracts word level data for a specific sentence

    :param data_container:          (h5py)  Container of the whole data, h5py object

    :param word_objects:            (h5py)  Container of all word data for a specific sentence

    :param eeg_float_resolution:    (type)  Resolution with which to save EEG, used for data compression

    :return:

        word_level_data     (dict)  Contains all word level data indexed by their index number in the sentence,

                                    together with the reading order, indexed by "word_reading_order"

    """

    available_objects = list(word_objects)

    #print(available_objects)

    #print(len(available_objects))

    # print('available_objects:', available_objects)

    if isinstance(available_objects[0], str):

        contentData = word_objects['content']

        #fixations_order_per_word = []

        if "rawEEG" in available_objects:

            rawData = word_objects['rawEEG']

            etData = word_objects['rawET']

            ffdData = word_objects['FFD']

            gdData = word_objects['GD']

            gptData = word_objects['GPT']

            trtData = word_objects['TRT']

            try:

                sfdData = word_objects['SFD']

            except KeyError:

                print("no SFD!")

                sfdData = []

            nFixData = word_objects['nFixations']

            fixPositions = word_objects["fixPositions"]

            Alpha_features_data = [word_objects[feature] for feature in Alpha_features]

            Beta_features_data = [word_objects[feature] for feature in Beta_features]

            Gamma_features_data = [word_objects[feature] for feature in Gamma_features]

            Theta_features_data = [word_objects[feature] for feature in Theta_features]

            #### 

            GD_EEG_features = [word_objects[feature] for feature in ['GD_t1','GD_t2','GD_a1','GD_a2','GD_b1','GD_b2','GD_g1','GD_g2']]

            FFD_EEG_features = [word_objects[feature] for feature in ['FFD_t1','FFD_t2','FFD_a1','FFD_a2','FFD_b1','FFD_b2','FFD_g1','FFD_g2']]

            TRT_EEG_features = [word_objects[feature] for feature in ['TRT_t1','TRT_t2','TRT_a1','TRT_a2','TRT_b1','TRT_b2','TRT_g1','TRT_g2']]

            #### 

            assert len(contentData) == len(etData) == len(rawData), "different amounts of different data!!"

            zipped_data = zip(rawData, etData, contentData, ffdData, gdData, gptData, trtData, sfdData, nFixData, fixPositions)

            word_level_data = {}

            word_idx = 0

            word_tokens_has_fixation = [] 

            word_tokens_with_mask = []

            word_tokens_all = []

            for raw_eegs_obj, ets_obj, word_obj, ffd, gd, gpt, trt, sfd, nFix, fixPos in zipped_data:

                word_string = load_matlab_string(data_container[word_obj[0]])

                if is_real_word(word_string):

                    data_dict = {}

                    data_dict["RAW_EEG"] = extract_all_fixations(data_container, raw_eegs_obj[0], eeg_float_resolution)

                    data_dict["RAW_ET"] = extract_all_fixations(data_container, ets_obj[0], np.float32)

                    data_dict["FFD"] = data_container[ffd[0]][()][0, 0] if len(data_container[ffd[0]][()].shape) == 2 else None

                    data_dict["GD"] = data_container[gd[0]][()][0, 0] if len(data_container[gd[0]][()].shape) == 2 else None

                    data_dict["GPT"] = data_container[gpt[0]][()][0, 0] if len(data_container[gpt[0]][()].shape) == 2 else None

                    data_dict["TRT"] = data_container[trt[0]][()][0, 0] if len(data_container[trt[0]][()].shape) == 2 else None

                    data_dict["SFD"] = data_container[sfd[0]][()][0, 0] if len(data_container[sfd[0]][()].shape) == 2 else None

                    data_dict["nFix"] = data_container[nFix[0]][()][0, 0] if len(data_container[nFix[0]][()].shape) == 2 else None

                    #fixations_order_per_word.append(np.array(data_container[fixPos[0]]))

                    #print([data_container[obj[word_idx][0]][()] for obj in Alpha_features_data])

                    data_dict["ALPHA_EEG"] = np.concatenate([data_container[obj[word_idx][0]][()]

                                                             if len(data_container[obj[word_idx][0]][()].shape) == 2 else []

                                                             for obj in Alpha_features_data], 0)

                    data_dict["BETA_EEG"] = np.concatenate([data_container[obj[word_idx][0]][()]

                                                            if len(data_container[obj[word_idx][0]][()].shape) == 2 else []

                                                            for obj in Beta_features_data], 0)

                    data_dict["GAMMA_EEG"] = np.concatenate([data_container[obj[word_idx][0]][()]

                                                             if len(data_container[obj[word_idx][0]][()].shape) == 2 else []

                                                             for obj in Gamma_features_data], 0)

                    data_dict["THETA_EEG"] = np.concatenate([data_container[obj[word_idx][0]][()]

                                                             if len(data_container[obj[word_idx][0]][()].shape) == 2 else []

                                                             for obj in Theta_features_data], 0)

                    data_dict["word_idx"] = word_idx

                    data_dict["content"] = word_string

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

                    word_tokens_all.append(word_string)

                    if data_dict["nFix"] is not None:

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

                        data_dict["GD_EEG"] = [np.squeeze(data_container[obj[word_idx][0]][()]) if len(data_container[obj[word_idx][0]][()].shape) == 2 else [] for obj in GD_EEG_features]

                        data_dict["FFD_EEG"] = [np.squeeze(data_container[obj[word_idx][0]][()]) if len(data_container[obj[word_idx][0]][()].shape) == 2 else [] for obj in FFD_EEG_features]

                        data_dict["TRT_EEG"] = [np.squeeze(data_container[obj[word_idx][0]][()]) if len(data_container[obj[word_idx][0]][()].shape) == 2 else [] for obj in TRT_EEG_features]

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

                        word_tokens_has_fixation.append(word_string)

                        word_tokens_with_mask.append(word_string)

                    else:

                        word_tokens_with_mask.append('[MASK]')

                    word_level_data[word_idx] = data_dict

                    word_idx += 1

                else:

                    print(word_string + " is not a real word.")

        else:

            # If there are no word-level data it will be word embeddings alone

            word_level_data = {}

            word_idx = 0

            word_tokens_has_fixation = [] 

            word_tokens_with_mask = []

            word_tokens_all = []

            for word_obj in contentData:

                word_string = load_matlab_string(data_container[word_obj[0]])

                if is_real_word(word_string):

                    data_dict = {}

                    data_dict["RAW_EEG"] = []

                    data_dict["ICA_EEG"] = []

                    data_dict["RAW_ET"] = []

                    data_dict["FFD"] = None

                    data_dict["GD"] = None

                    data_dict["GPT"] = None

                    data_dict["TRT"] = None

                    data_dict["SFD"] = None

                    data_dict["nFix"] = None

                    data_dict["ALPHA_EEG"] = []

                    data_dict["BETA_EEG"] = []

                    data_dict["GAMMA_EEG"] = []

                    data_dict["THETA_EEG"] = []

                    data_dict["word_idx"] = word_idx

                    data_dict["content"] = word_string

                    word_level_data[word_idx] = data_dict

                    word_idx += 1

                else:

                    print(word_string + " is not a real word.")

            sentence = " ".join([load_matlab_string(data_container[word_obj[0]]) for word_obj in word_objects['content']])

            #print("Only available objects for the sentence '{}' are {}.".format(sentence, available_objects))

            #word_level_data["word_reading_order"] = extract_word_order_from_fixations(fixations_order_per_word)

    else:

        word_tokens_has_fixation = [] 

        word_tokens_with_mask = []

        word_tokens_all = []

        word_level_data = {}

    return word_level_data, word_tokens_all, word_tokens_has_fixation, word_tokens_with_mask