#!/usr/bin/python

import sys
import numpy as np
import mne
from mne.preprocessing import ICA

argc = len(sys.argv)

# lists of data sources
subjects = ('CC110033', 'CC110037', 'CC110045')
# different origins
archs = ('passive', 'rest', 'task')

id_subject = 0
id_arch = 0
subject = subjects[id_subject]
arch = archs[id_arch]

subjects_dir = 'subjects'

# load raw data from the first subject
raw_data = mne.io.Raw(subject + '/' + arch + '/' + arch + '_raw.fif')

# display the channels. expert can select a channel to 'blacklist' it
if argc > 1 and sys.argv[1] == '-p':
    #raw_data.plot(block=True, lowpass=40)
    raw_data.plot_psd()

# filter the bad channels
calibration_file = 'sss_cal.dat'
cross_talk_file = 'ct_sparse.fif'
preprocessed_data = mne.preprocessing.maxwell_filter(raw_data,
        calibration=calibration_file,
        cross_talk=cross_talk_file)

# display the channels after filtering
if argc > 1 and sys.argv[1] == '-p':
    preprocessed_data.plot(block=True, lowpass=40)
    preprocessed_data.plot_psd()

# keep frequencies from 1Hz to 45Hz (remove noise and signal of non-interest)
low_freq = 1
high_freq = 45
bandpass_data = preprocessed_data.filter(low_freq, high_freq)

# checked that the filtered frequencies are dumped
if argc > 1 and sys.argv[1] == '-p':
    #bandpass_data.plot(block=True, lowpass=40)
    bandpass_data.plot_psd()

# compute ICA
picks_meg = mne.pick_types(bandpass_data.info, meg=True, eeg=False, eog=False,
        stim=False, exclude='bads')

method = 'fastica'
# a good estimate of the number of components is the rank of the Raw object
n_components = bandpass_data.estimate_rank()
decim = 3
random_state = 23
ica = ICA(n_components=n_components, method=method, random_state=random_state)

reject = dict(mag=5e-12, grad=4000e-13)
# goal: remove artifacts, so they should be easy to notice
ica.fit(bandpass_data, picks=picks_meg, decim=decim, reject=reject)

if argc > 1 and sys.argv[1] == '-p':
    ica.plot_components()


# pass on Epoch computing
# 1) extract events from stim channels
events = mne.find_events(bandpass_data, stim_channel='STI101')
# 2) create epochs from Raw + events
epochs = mne.Epochs(bandpass_data, events)
# 3) average the epochs to reduce noise
evoked = epochs.average()

evoked.shift_time(-0.05)
if argc > 1 and sys.argv[1] == '-p':
    evoked.plot()


bandpass_data.save('out.fif')
mne.gui.coregistration(subjects_dir=subjects_dir)