d

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

@author: J. Abreu

"""

#%% importings
import numpy as np
import mne
import os
from argparse import ArgumentParser

mne.set_log_level('WARNING')
mne.utils.set_config('MNE_USE_CUDA', 'true')
#mne.io.Raw.set_montage(montage='standard_1020')
#%%
parser = ArgumentParser()
parser.add_argument("--dev-sys", "-devsys", type=str, required=False, default='NEUROMAP',
                    help="Indica qual o sistema de captura. Opções 'DSI', 'NEUROMAP'.")

#%%
def read_edf(path='D:/Jeremias/Datasets/tuh_mixed', num_files=1):
    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), preload=True))
        if(len(raw_list) == num_files):
            print(len(raw_list), ' raw files were loaded')
            return raw_list
    else:
        print(len(raw_list), ' raw files were loaded')
        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, set_dig = True, 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)
    print(picks)

    epochs = mne.Epochs(raw, events, event_id=1, picks=picks, preload=True)
    epochs.plot_psd(fmin=1., fmax=40.)
    epochs.plot_psd_topomap()
    evoked = epochs.average()
    #evoked.set_montage(montage, set_dig=True, raise_if_subset=True)
    evoked.plot()
#%%
def ch_order(sys_map):

    if sys_map == 'DSI':
        map_ch = {'EEG FP1-LE':'Fp1',
                 'EEG FP2-LE':'Fp2',
                 'EEG F3-LE':'F3',
                 'EEG F4-LE':'F4',
                 'EEG C3-LE':'C3',
                 'EEG C4-LE':'C4',
                 'EEG A1-LE':'A1',
                 'EEG A2-LE':'A2',
                 'EEG P3-LE':'P3',
                 'EEG P4-LE':'P4',
                 'EEG O1-LE':'O1',
                 'EEG O2-LE':'O2',
                 'EEG F7-LE':'F7',
                 'EEG F8-LE':'F8',
                 'EEG T3-LE':'T3',
                 'EEG T4-LE':'T4',
                 'EEG T5-LE':'T5',
                 'EEG T6-LE':'T6',
                 'EEG FZ-LE':'Fz',
                 'EEG CZ-LE':'Cz',
                 'EEG PZ-LE':'Pz',
                 'EEG OZ-LE':'Oz',
                 'EEG PG1-LE':'Pg1',
                 'EEG PG2-LE':'Pg2',
                 'EEG EKG-LE':'Ekg',
                 'EEG SP2-LE':'Sp2',
                 'EEG SP1-LE':'Sp1',
                 'EEG T1-LE':'T1',
                 'EEG T2-LE':'T2',
                 'EEG C4-LE':'C4', 'EEG CZ-LE':'Cz', 'EEG EKG-LE':'Ekg',
                 'EEG F3-LE':'F3', 'EEG F4-LE':'F4', 'EEG F7-LE':'F7',
                 'EEG F8-LE':'F8', 'EEG FP1-LE':'Fp1', 'EEG FP2-LE':'Fp2',
                 'EEG FZ-LE':'Fz', 'EEG O1-LE':'O1', 'EEG O2-LE':'O2',
                 'EEG OZ-LE':'Oz', 'EEG P3-LE':'P3', 'EEG P4-LE':'P4',
                 'EEG PG1-LE':'Pg1', 'EEG PG2-LE':'Pg2', 'EEG PZ-LE':'Pz',
                 'EEG SP1-LE':'Sp1', 'EEG SP2-LE':'Sp2','EEG T1-LE':'T1',
                 'EEG T2-LE':'T2', 'EEG T3-LE':'T3', 'EEG T4-LE':'T4','EEG T5-LE':'T5',
                 'EEG T6-LE':'T6'}

        ch_order = ['Fp1','Fp2','AF9','AF7','AF5','AF3','AF1','AFz','AF2','AF4',
                    'AF6','AF8','AF10','F9','F7','F5','F3','F1','Fz','F2','F4','F6',
                    'F8','F10','FT9','FT7','FC5','FC3','FC1','FCz','FC2','FC4','FC6',
                    'FT8','FT10','T9','T7','C5','C3','C1','Cz','C2','C4','C6','T8',
                    'T10','TP9','TP7','CP5','CP3','CP1','CPz','CP2','CP4','CP6','TP8',
                    'TP10','P9','P7','P5','P3','P1','Pz','P2','P4','P6','P8','P10',
                    'PO9','PO7','PO5','PO3','PO1','POz','PO2','PO4','PO6','PO8','PO10',
                    'O1','Oz','O2','O9','Iz','O10','T3','T5','T4','T6','M1','M2','A1',
                    'A2']
        return (map_ch, ch_order)
    else:
        print('sys_map is not DSI')
        return None
#%%
raw_list = read_edf(num_files=1)
#%%
raw = raw_list[0]
raw.info['bads'] = ['PHOTIC PH']

picks = mne.pick_types(raw.info, meg=False, eeg=True)
raw_1020 = raw.copy()
raw_1020.drop_channels(['PHOTIC PH','EEG 28-LE', 'EEG 29-LE', 'EEG 30-LE'])

(map_ch, ch_order) = ch_order('DSI')

raw_1020.rename_channels(map_ch)
#raw = raw.reorder_channels(ch_order)
#%%
raw_1020 = raw_1020.set_montage('standard_1020', raise_if_subset=False)

raw_1020.plot_psd_topo()
#%%

montage = mne.channels.make_standard_montage('standard_1020')
raw.set_montage(montage=montage)
raw.set_montage(montage, set_dig=True)
raw.plot_psd_topo()
#viz(raw)
#%%

montage = mne.channels.read_montage('standard_1020')



#raw.set_montage(montage, set_dig=True)