import os
import mne


def load_bad_annotations(filepath, fileending="badSegments.csv"):
    """
    Loads the annotations for bad segments
    :param filepath: The path to the file we want to load
    :param fileending: Depending if the subject, for which we load the annotations, was preprocessed manually,
                       the ending of the filename will be different.
                       The default are file endings of the given preprocessed annotations: "badSegments.csv"
                       For the manual preprocessed annotations, the file endings are: "badannotations.csv"
    :return: The mne annotations
    """
    if os.path.isfile(filepath + "_" + fileending):
        return mne.read_annotations(filepath + "_" + fileending)


def load_preprocessed_data(subject, dataset):
    """
    Load the raw object of the preprocessed file
    :param subject: The subject, for which we want to load the raw object
    :param dataset: The currently viewed dataset
    :param selected_subjects: The manually preprocessed subjects
    :return: The raw object
    """
    folder = "Dataset/" + dataset + "/sub-" + subject + "/ses-" + dataset + "/eeg/"
    filepath = folder + "sub-" + subject + "_task-" + dataset
    raw = mne.io.read_raw_fif(filepath + "_cleaned.fif")
    return raw


def get_keys_for_events(stimulus=None, condition=None):
    """
    For a given stimulus and condition get all the event keys.
    :param stimulus: Either 'face' or 'car' or 'None' for no stimulus
    :param condition: Either 'intact' or 'scrambled' or 'None' for no condition
    :return: A list of keys or 'stimulus' if neither stimulus or condition was given
    """
    if stimulus == 'face':
        if condition == 'intact':
            return ["stimulus:{}".format(k) for k in range(1, 41)]
        elif condition == 'scrambled':
            return ["stimulus:{}".format(k) for k in range(101, 141)]
        else:  # All faces
            return ["stimulus:{}".format(k) for k in list(range(1, 41)) + list(range(101, 141))]
    elif stimulus == 'car':
        if condition == 'intact':
            return ["stimulus:{}".format(k) for k in range(41, 81)]
        elif condition == 'scrambled':
            return ["stimulus:{}".format(k) for k in range(141, 181)]
        else:  # All cars
            return ["stimulus:{}".format(k) for k in list(range(41, 81)) + list(range(141, 181))]
    else:  # If no stimulus is given
        if condition == 'intact':
            return ["stimulus:{}".format(k) for k in range(1, 41) and range(41, 81)]
        elif condition == 'scrambled':
            return ["stimulus:{}".format(k) for k in list(range(101, 141)) + list(range(141, 181))]
        else:  # Every stimulus
            return "stimulus"


def get_epochs(raw, conditions=None, picks=None, tmin=-0.1, tmax=1):
    """
    Returns the epochs for a given dataset
    :param raw: the dataset
    :param conditions: A List of tuples, of the Form [(stimulus, condition), (stimulus,condition)]
        i.e. [('face',None), ('car', 'scrambled')] returns the epochs where the stimulus is face and the stim+condition is car+scrambled
        default is None, i.e. everything
    :param picks: a list. Additional criteria for picking the epochs, e.g. channels, etc.
    :param tmin: onset before the event
    :param tmax: end after the event
    :return:
    """

    events, events_dict = mne.events_from_annotations(raw)
    events_dict_key = {}
    if conditions is None:
        conditions = [(None, None)]

    for condition in conditions:
        wanted_keys = get_keys_for_events(condition[0], condition[1])
        if wanted_keys == "stimulus":
            wanted_keys = [e for e in events_dict.keys() if "stimulus" in e]
        events_dict_key.update(dict((k, events_dict[k]) for k in wanted_keys if k in events_dict))
    epochs = mne.Epochs(raw, events, events_dict_key, tmin=tmin, tmax=tmax, reject_by_annotation=False, picks=picks)
    epochs.drop_bad()

    return epochs, events_dict_key
