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

"""

Created on Sun Jul 19 15:06:16 2015

@author: yicong

"""

print(__doc__)

import os

import numpy as np

import pandas as pd

from mne.io import RawArray

from mne.channels import read_montage

from mne.epochs import concatenate_epochs

from mne import create_info, find_events, Epochs, concatenate_raws, pick_types

from mne.decoding import CSP

from sklearn.linear_model import LogisticRegression

from glob import glob

from scipy.signal import butter, lfilter, convolve, boxcar

from joblib import Parallel, delayed

#traindir = r"../30 Data/train"

#fnames = [os.path.join(traindir, f) for f in os.listdir(traindir)]

def creat_mne_raw_object(fname,read_events=True):

return raw

subjects = range(1,13)

ids_tot = []

pred_tot = []

# design a butterworth bandpass filter

freqs = [7, 30]

b,a = butter(5,np.array(freqs)/250.0,btype='bandpass')

# CSP parameters

# Number of spatial filter to use

nfilters = 4

# convolution

# window for smoothing features

nwin = 250

# training subsample

subsample = 10

# submission file

submission_file = 'beat_the_benchmark.csv'

cols = ['HandStart','FirstDigitTouch',

'BothStartLoadPhase','LiftOff',

'Replace','BothReleased']

for subject in subjects:

epochs_tot = []

y = []

################ READ DATA ################################################

fnames = glob("../30 Data/train/subj%d_series*_data.csv" % (subject))

# read and concatenate all the files

raw = concatenate_raws([creat_mne_raw_object(fname) for fname in fnames])

# pick eeg signal

picks = pick_types(raw.info,eeg=True)

# Filter data for alpha frequency and beta band

# Note that MNE implement a zero phase (filtfilt) filtering not compatible

# with the rule of future data.

# Here we use left filter compatible with this constraint.

# The function parallelized for speeding up the script

raw._data[picks] = np.array(Parallel(n_jobs=-1)(delayed(lfilter)(b,a,raw._data[i]) for i in picks))

################ CSP Filters training #####################################

# get event posision corresponding to Replace

events = find_events(raw,stim_channel='BothReleased', verbose=False)

# epochs signal for 1.5 second before the movement

epochs = Epochs(raw, events, {'during' : 1}, -2, -0.5, proj=False,

picks=picks, baseline=None, preload=True,

add_eeg_ref=False, verbose=False)

epochs_tot.append(epochs)

y.extend([1]*len(epochs))

# epochs signal for 1.5 second after the movement, this correspond to the

# rest period.

epochs_rest = Epochs(raw, events, {'after' : 1}, 0.5, 2, proj=False,

picks=picks, baseline=None, preload=True,

add_eeg_ref=False, verbose=False)

# Workaround to be able to concatenate epochs with MNE

epochs_rest.times = epochs.times

y.extend([-1]*len(epochs_rest))

epochs_tot.append(epochs_rest)

# Concatenate all epochs

epochs = concatenate_epochs(epochs_tot)

# get data

X = epochs.get_data()

y = np.array(y)

# train CSP

csp = CSP(n_components=nfilters, reg='lws')

csp.fit(X,y)

################ Create Training Features #################################

# apply csp filters and rectify signal

feat = np.dot(csp.filters_[0:nfilters],raw._data[picks])**2

# smoothing by convolution with a rectangle window

feattr = np.array(Parallel(n_jobs=-1)(delayed(convolve)(feat[i],boxcar(nwin),'full') for i in range(nfilters)))

feattr = np.log(feattr[:,0:feat.shape[1]])

# training labels

# they are stored in the 6 last channels of the MNE raw object

labels = raw._data[32:]

################ Create test Features #####################################

# read test data

fnames = glob("../30 Data/test/subj%d_series*_data.csv" % (subject))

raw = concatenate_raws([creat_mne_raw_object(fname, read_events=False) for fname in fnames])

raw._data[picks] = np.array(Parallel(n_jobs=-1)(delayed(lfilter)(b,a,raw._data[i]) for i in picks))

# read ids

ids = np.concatenate([np.array(pd.read_csv(fname)['id']) for fname in fnames])

ids_tot.append(ids)

# apply preprocessing on test data

feat = np.dot(csp.filters_[0:nfilters],raw._data[picks])**2

featte = np.array(Parallel(n_jobs=-1)(delayed(convolve)(feat[i],boxcar(nwin),'full') for i in range(nfilters)))

featte = np.log(featte[:,0:feat.shape[1]])

################ Train classifiers ########################################

lr = LogisticRegression()

pred = np.empty((len(ids),6))

for i in range(6):

print('Train subject %d, class %s' % (subject, cols[i]))

lr.fit(feattr[:,::subsample].T,labels[i,::subsample])

pred[:,i] = lr.predict_proba(featte.T)[:,1]

pred_tot.append(pred)

# create pandas object for sbmission

submission = pd.DataFrame(index=np.concatenate(ids_tot),

columns=cols,

data=np.concatenate(pred_tot))

# write file

submission.to_csv(submission_file,index_label='id',float_format='%.3f')