"""
=========================================================
Time frequency : Induced power and inter-trial phase-lock
=========================================================

This script shows how to compute induced power and inter-trial
phase-lock for a list of epochs read in a raw file given
a list of events.

"""
# Authors: Alexandre Gramfort <gramfort@nmr.mgh.harvard.edu>
#
# License: BSD (3-clause)

print __doc__

import numpy as np

import mne
from mne import fiff
from mne.time_frequency import induced_power
from mne.datasets import sample

###############################################################################
# Set parameters
data_path = sample.data_path()
raw_fname = data_path + '/MEG/sample/sample_audvis_raw.fif'
event_fname = data_path + '/MEG/sample/sample_audvis_raw-eve.fif'
event_id, tmin, tmax = 1, -0.2, 0.5

# Setup for reading the raw data
raw = fiff.Raw(raw_fname)
events = mne.read_events(event_fname)

include = []
raw.info['bads'] += ['MEG 2443', 'EEG 053']  # bads + 2 more

# picks MEG gradiometers
picks = fiff.pick_types(raw.info, meg='grad', eeg=False, eog=True,
                        stim=False, include=include, exclude='bads')

epochs = mne.Epochs(raw, events, event_id, tmin, tmax, picks=picks,
                    baseline=(None, 0), reject=dict(grad=4000e-13, eog=150e-6))
data = epochs.get_data()  # as 3D matrix
evoked = epochs.average()  # compute evoked fields

times = 1e3 * epochs.times  # change unit to ms
evoked_data = evoked.data * 1e13  # change unit to fT / cm

# Take only one channel
data = data[:, 97:98, :]
evoked_data = evoked_data[97:98, :]

frequencies = np.arange(7, 30, 3)  # define frequencies of interest
n_cycles = frequencies / float(7)  # different number of cycle per frequency
Fs = raw.info['sfreq']  # sampling in Hz
decim = 3
power, phase_lock = induced_power(data, Fs=Fs, frequencies=frequencies,
                                  n_cycles=n_cycles, n_jobs=1, use_fft=False,
                                  decim=decim, zero_mean=True)

# baseline corrections with ratio
power /= np.mean(power[:, :, times[::decim] < 0], axis=2)[:, :, None]

###############################################################################
# View time-frequency plots
import matplotlib.pyplot as plt
plt.clf()
plt.subplots_adjust(0.1, 0.08, 0.96, 0.94, 0.2, 0.63)
plt.subplot(3, 1, 1)
plt.plot(times, evoked_data.T)
plt.title('Evoked response (%s)' % evoked.ch_names[97])
plt.xlabel('time (ms)')
plt.ylabel('Magnetic Field (fT/cm)')
plt.xlim(times[0], times[-1])
plt.ylim(-150, 300)

plt.subplot(3, 1, 2)
plt.imshow(20 * np.log10(power[0]),
           extent=[times[0], times[-1], frequencies[0], frequencies[-1]],
           aspect='auto', origin='lower')
plt.xlabel('Time (s)')
plt.ylabel('Frequency (Hz)')
plt.title('Induced power (%s)' % evoked.ch_names[97])
plt.colorbar()

plt.subplot(3, 1, 3)
plt.imshow(phase_lock[0],
           extent=[times[0], times[-1], frequencies[0], frequencies[-1]],
           aspect='auto', origin='lower')
plt.xlabel('Time (s)')
plt.ylabel('Frequency (Hz)')
plt.title('Phase-lock (%s)' % evoked.ch_names[97])
plt.colorbar()
plt.show()
