import mne
from mne.preprocessing import create_eog_epochs

from mne_bids import BIDSPath, read_raw_bids
from utils.ccs_eeg_utils import read_annotations_core
from utils.file_utils import get_epochs


def load_unprocessed_subject(subject, dataset):
    """
    Load the eeg data of a subject
    :param subject: The subject of which the data will be loaded
    :param dataset: The dataset which will be loaded
    :return: the subject data
    """
    bids_path = BIDSPath(subject=subject, task=dataset, session=dataset, datatype='eeg', suffix='eeg',
                         root='Dataset\\' + dataset)
    raw = read_raw_bids(bids_path)
    # Add annotations
    read_annotations_core(bids_path, raw)
    return raw


def filter_data(raw):
    """
    Filter the data of a single subject with a bandpass filter.
    The lower bound ist 0.5Hz to compensate the slow drifts.
    The upper bound is 50Hz to compensate the high frequencies, including the power line spike at 60Hz
    :param raw: The data to be filtered
    :return: The filtered data
    """
    raw.filter(0.5, 48, fir_design='firwin')
    return raw


def plot_filter_data():
    ds = 'N170'
    for subj in ['014']:
        data = load_unprocessed_subject(subj, ds)
        data.load_data()
        # data.plot(n_channels=len(data.ch_names), block=True, scalings=40e-6)
        filter_data(data)
        fig = mne.viz.plot_raw_psd(data, fmax=80, average=True, show=False)
        fig.savefig("plots/frequency_filtered_subj_" + subj + "_48Hz.png")
        # data.plot(n_channels=len(data.ch_names), block=True, scalings=40e-6)


def plot_filter_data_epoched(subj):
    ds = 'N170'
    data = load_unprocessed_subject(subj, ds)
    data.load_data()
    filter_data(data)
    get_epochs(data)[0].average().plot()


def plot_cleaning():
    ds = 'N170'
    for subj in ['014']:
        data = load_unprocessed_subject(subj, ds)
        data.load_data()
        filter_data(data)
        folder = "Dataset\\" + ds + "\\sub-" + subj + "\\ses-" + ds + "\\eeg\\"
        filepath = folder + "sub-" + subj + "_task-" + ds
        print(filepath)
        ann = mne.read_annotations(filepath + "_" + "badannotations.csv")
        data.annotations.append(ann.onset, ann.duration, ann.description)
        data.plot(n_channels=len(data.ch_names), block=True, scalings=40e-6)


def plot_ica():
    ds = 'N170'
    for subj in ['014']:
        data = load_unprocessed_subject(subj, ds)
        folder = "Dataset\\" + ds + "\\sub-" + subj + "\\ses-" + ds + "\\eeg\\"
        filepath = folder + "sub-" + subj + "_task-" + ds
        ann = mne.read_annotations(filepath + "_" + "badannotations.csv")
        data.annotations.append(ann.onset, ann.duration, ann.description)
        data.load_data()
        data.set_channel_types({'HEOG_left': 'eog', 'HEOG_right': 'eog', 'VEOG_lower': 'eog'})
        data.set_montage('standard_1020', match_case=False)
        ica_raw = data.copy()
        ica_raw.filter(l_freq=1, h_freq=None)

        # Then run ICA
        ica = mne.preprocessing.ICA(method="fastica",
                                    random_state=123)  # Use a random state for reproducable results #TODO Old Random state 123 or new one?
        ica.fit(ica_raw, verbose=True)
        ica_raw.load_data()
        # ica.plot_components(inst=ica_raw, ch_type='eeg', contours=0, topomap_args={'extrapolate': 'head'},
        #                    psd_args={'fmin': 0, 'fmax': 80})
        ica.plot_sources(ica_raw)
        ica.plot_properties(inst=ica_raw, dB=False, topomap_args={'extrapolate': 'head', 'contours': 0},
                            psd_args={'fmin': 3, 'fmax': 50}, picks=['eeg'])


def plot_joint_eog_plots():
    ds = 'N170'
    for subj in ['014']:
        data = load_unprocessed_subject(subj, ds)
        data.load_data()
        data.set_channel_types({'HEOG_left': 'eog', 'HEOG_right': 'eog', 'VEOG_lower': 'eog'})
        data.set_montage('standard_1020', match_case=False)
        eog_evoked = create_eog_epochs(data).average()
        eog_evoked.apply_baseline(baseline=(None, -0.2))
        eog_evoked.plot_joint()


# plot_ica()
# plot_joint_eog_plots()
plot_filter_data_epoched('003')