mne

# -*- coding: utf-8 -*-
"""
Created on Mon Mar 16 11:06:18 2020

@author: J. Abreu

"""

#%% importings
import numpy as np
import mne
import os

mne.set_log_level('WARNING')
mne.utils.set_config('MNE_USE_CUDA', 'true')
#mne.io.Raw.set_montage(montage='standard_1020')
#%%
def read_edf(path='D:/Jeremias/Datasets/tuh_mixed'):
    raw_list = []

    for file in os.listdir(path):
        if os.path.isfile(os.path.join(path, file)) and os.path.join(path, file)[-3:] == 'edf':
            raw_list.append(mne.io.read_raw_edf(os.path.join(path, file)))
    return raw_list
#%%
def sum_events(raw_list):
    if type(raw_list) == list:
        sum = 0
        for raw in raw_list:
            sum += len(raw.info['events'])
        print(sum)
    elif type(raw_list) == mne.io.edf.edf.RawEDF:
        print(len(raw_list.info['events']))
    else:
        print('type of file not accepted! only list of io.edf.edf.RawEDF ' \
              + 'or a single io.edf.edf.RawEDF file')

#%%
def get_annotation_file(raw=None):
    if raw is None:
        print('no file name')
        return

    try:
        path = raw._filenames[0][:-9] + '.txt'
        with open(path, 'r') as f:
            txt = f.read()

        print(txt)
    except:
        raise FileNotFoundError('verify if the file name is correct')
#%% set standard montage
#montage = mne.channels.make_standard_montage('stardad_1020')
#raw.set_montage(montage, raise_if_subset=False, verbose=False)

#%%
def viz(raw):
    montage = mne.channels.make_standard_montage('standard_1020')
    raw.set_montage(montage, raise_if_subset=False, verbose=False)

    print('number of events: ', len(raw.info['events']))

    events = mne.make_fixed_length_events(raw, first_samp=True)
    picks = mne.pick_types(raw.info, meg=False, eeg=True)

    baseline = (None, 0)

    epochs = mne.Epochs(raw, events, event_id=1, picks=picks, preload=True)
    evoked = epochs.average()
    evoked.plot()