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

"""

Created on Wed Jul  8 22:00:08 2015.

@author: rc, alexandre

"""

import numpy as np

import pandas as pd

from mne.io import RawArray

from mne.channels import read_montage

from mne import create_info, concatenate_raws, pick_types

from sklearn.base import BaseEstimator, TransformerMixin

from glob import glob

def getChannelNames():

    """Return Channels names."""

    return ['Fp1', 'Fp2', 'F7', 'F3', 'Fz', 'F4', 'F8', 'FC5', 'FC1', 'FC2',

            'FC6', 'T7', 'C3', 'Cz', 'C4', 'T8', 'TP9', 'CP5', 'CP1', 'CP2',

            'CP6', 'TP10', 'P7', 'P3', 'Pz', 'P4', 'P8', 'PO9', 'O1', 'Oz',

            'O2', 'PO10']

def getEventNames():

    """Return Event name."""

    return ['HandStart', 'FirstDigitTouch', 'BothStartLoadPhase', 'LiftOff',

            'Replace', 'BothReleased']

def load_raw_data(subject, test=False):

    """Load Raw data from files.

    For a given subject, csv files are loaded, converted to MNE raw instance

    and concatenated.

    If test is True, training data are composed of series 1 to 8 and test data

    of series 9 and test. Otherwise, training data are series 1 to 6 and test

    data series 7 and 8.

    """

    fnames_train = glob('../data/train/subj%d_series*_data.csv' % (subject))

    fnames_train.sort()

    if test:

        fnames_test = glob('../data/test/subj%d_series*_data.csv' % (subject))

        fnames_test.sort()

    else:

        fnames_test = fnames_train[-2:]

        fnames_train = fnames_train[:-2]

    # read and concatenate all the files

    raw_train = [creat_mne_raw_object(fname) for fname in fnames_train]

    raw_train = concatenate_raws(raw_train)

    # pick eeg signal

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

    # get training data

    data_train = raw_train._data[picks].T

    labels_train = raw_train._data[32:].T

    raw_test = [creat_mne_raw_object(fname, read_events=not test) for fname in

                fnames_test]

    raw_test = concatenate_raws(raw_test)

    data_test = raw_test._data[picks].T

    # extract labels if validating on series 7&8

    labels_test = None

    if not test:

        labels_test = raw_test._data[32:].T

    return data_train, labels_train, data_test, labels_test

def creat_mne_raw_object(fname, read_events=True):

    """Create a mne raw instance from csv file."""

    # Read EEG file

    data = pd.read_csv(fname)

    # get chanel names

    ch_names = list(data.columns[1:])

    # read EEG standard montage from mne

    montage = read_montage('standard_1005', ch_names)

    ch_type = ['eeg']*len(ch_names)

    data = 1e-6*np.array(data[ch_names]).T

    if read_events:

        # events file

        ev_fname = fname.replace('_data', '_events')

        # read event file

        events = pd.read_csv(ev_fname)

        events_names = events.columns[1:]

        events_data = np.array(events[events_names]).T

        # define channel type, the first is EEG, the last 6 are stimulations

        ch_type.extend(['stim']*6)

        ch_names.extend(events_names)

        # concatenate event file and data

        data = np.concatenate((data, events_data))

    # create and populate MNE info structure

    info = create_info(ch_names, sfreq=500.0, ch_types=ch_type,

                       montage=montage)

    info['filename'] = fname

    # create raw object

    raw = RawArray(data, info, verbose=False)

    return raw

def sliding_window(sig, window=512, subsample=10, estimator=None):

    """Extract a slinding window from signal.

    Raw signal is padded with zeros on the left to avoid use of future data.

    """

    Ne, Ns = sig.shape

    # get the index before padding

    ix = range(0, Ns, subsample)

    # padd data

    padd = np.zeros((Ne, int(window) - 1))

    sig = np.concatenate((padd, sig), axis=1)

    Ne, Ns = sig.shape

    if estimator is None:

        estimator = np.array

    # call this to get the shape

    X = estimator(sig[:, 0:window])

    dims = list(X.shape)

    dims.insert(0, len(ix))

    dims = tuple(dims)

    # allocate array

    X = np.empty(dims, dtype=X.dtype)

    for i in range(len(ix)):

        X[i] = estimator(sig[:, ix[i]:(ix[i] + window)])

    return X

def delay_preds(X, delay=100, skip=2, subsample=1, start=0, jump=None):

    """Delay predictions.

    Create a feature vector by concatenation of present and past sample.

    The concatenation is done by shifting data to the right :

    out = | x1 x2 x3 ...  xn   |

          | 0  x1 x2 ...  xn-1 |

          | 0  0  x1 ...  xn-2 |

    No use of future data.

    """

    if jump is None:

        jump = range(0, delay, skip)

    Ns, Ne = X.shape

    Ns_subsampled = len(range(start, Ns, subsample))

    out = np.zeros((Ns_subsampled, Ne * len(jump)))

    for i, sk in enumerate(jump):

        chunk = X[0:(Ns - sk)][start::subsample]

        out[(Ns_subsampled-chunk.shape[0]):, (i * Ne):((i + 1) * Ne)] = chunk

    return out

def delay_preds_2d(X, delay=100, skip=2, subsample=1, start=0, jump=None):

    """Delay predictions with 2d shape.

    Same thing as delay_pred, but return delayed prediction with a 2d shape.

    """

    if jump is None:

        jump = range(0, delay, skip)

    Ns, Ne = X.shape

    Ns_subsampled = len(range(start, Ns, subsample))

    out = np.zeros((Ns_subsampled, len(jump), Ne))

    for i, sk in enumerate(jump):

        chunk = X[0:(Ns - sk)][start::subsample]

        out[(Ns_subsampled-chunk.shape[0]):, i, :] = chunk

    return out[:, ::-1, :]

class SlidingWindow(BaseEstimator, TransformerMixin):

    """Sliding Window tranformer Mixin."""

    def __init__(self, window=500, subsample=10, estimator=np.array):

        """Init."""

        self.window = window

        self.subsample = subsample

        self.estimator = estimator

    def fit(self, X, y=None):

        """Fit, not used."""

        return self

    def transform(self, X, y=None):

        """Transform."""

        return sliding_window(X.T, window=self.window,

                              subsample=self.subsample,

                              estimator=self.estimator)

    def update_subsample(self, old_sub, new_sub):

        """update subsampling."""

        self.subsample = new_sub

class SubSample(BaseEstimator, TransformerMixin):

    """Subsample tranformer Mixin."""

    def __init__(self, subsample=10):

        """Init."""

        self.subsample = subsample

    def fit(self, X, y=None):

        """Fit, not used."""

        return self

    def transform(self, X, y=None):

        """Transform."""

        return X[::self.subsample]

    def update_subsample(self, old_sub, new_sub):

        """update subsampling."""

        self.subsample = new_sub

class DelayPreds(BaseEstimator, TransformerMixin):

    """Delayed prediction tranformer Mixin."""

    def __init__(self, delay=1000, skip=100, two_dim=False):

        """Init."""

        self.delay = delay

        self.skip = skip

        self.two_dim = two_dim

    def fit(self, X, y=None):

        """Fit, not used."""

        return self

    def transform(self, X, y=None):

        """Transform."""

        if self.two_dim:

            return delay_preds_2d(X, delay=self.delay, skip=self.skip)

        else:

            return delay_preds(X, delay=self.delay, skip=self.skip)

    def update_subsample(self, old_sub, new_sub):

        """update subsampling."""

        ratio = old_sub / new_sub

        self.delay = int(self.delay * ratio)

        self.skip = int(self.skip * ratio)

class NoneTransformer(BaseEstimator, TransformerMixin):

    """Return None Transformer."""

    def __init__(self):

        """Init."""

        pass

    def fit(self, X, y=None):

        """Fit, not used."""

        return self

    def transform(self, X, y=None):

        """Transform."""

        return None