# -*- coding: utf-8 -*-

"""

Created on Mon Aug 31 10:17:09 2015

@author: mje

"""

import mne

from mne.minimum_norm import (apply_inverse_epochs, read_inverse_operator)

from mne.time_frequency import cwt_morlet

import numpy as np

from my_settings import *

# Using the same inverse operator when inspecting single trials Vs. evoked

snr = 1.0 # Standard assumption for average data but using it for single trial

lambda2 = 1.0 / snr ** 2

method = "dSPM" # use dSPM method (could also be MNE or sLORETA)

freqs = np.arange(8, 13, 1)

n_cycle = freqs / 3.

conditions = ["ent/left", "ctl/left", "ent/right", "ctl/right"]

for subject in subjects_select:

# Load data

labels = mne.read_labels_from_annot(subject, parc='PALS_B12_Brodmann',

regexp="Bro",

subjects_dir=subjects_dir)

labels_sel = [labels[6], labels[7]]

inverse_operator = read_inverse_operator(mne_folder +

"%s-inv.fif" % subject)

src = inverse_operator["src"]

epochs = mne.read_epochs(epochs_folder +

"%s_trial_start-epo.fif" % subject)

# epochs.drop_bad_epochs(reject_params)

# epochs.resample(250, n_jobs=4)

for condition in conditions:

stcs = apply_inverse_epochs(epochs[condition],

inverse_operator,

lambda2,

method,

pick_ori=None)

for label in labels_sel:

label_ts = []

for j in range(len(stcs)):

ts = mne.extract_label_time_course(stcs[j],

labels=label,

src=src,

mode="pca_flip")

ts = np.squeeze(ts)

ts *= np.sign(ts[np.argmax(np.abs(ts))])

label_ts.append(ts)

label_ts = np.asarray(label_ts)

tfr = cwt_morlet(label_ts, epochs.info["sfreq"], freqs,

use_fft=True, n_cycles=n_cycle)

np.save(tf_folder + "%s_%s_%s_%s_%s_sf-tfr" % (subject,

condition[:3],

condition[4:],

label.name, method),

tfr)

np.save(tf_folder + "%s_%s_%s_%s_%s_sf-ts" % (subject,

condition[:3],

condition[4:],

label.name, method),

label_ts)

del stcs

del tfr