"""
===================================
WIP Prepare multiconditions events
===================================
"""
# Authors: Jean-Remi King <jeanremi.king@gmail.com>
#
# License: BSD (3-clause)


import pandas as pd
import numpy as np

import mne
from mne.datasets import sample

print(__doc__)


class _BaseEvents():
    def __init__(self, events_array, event_id=None):
        # Set default fields in pandas.DataFrame format
        data = [dict(_start=start, _length=length, _value=value)
                for start, length, value in events_array]
        # Set additional fields defined by event_id
        if event_id is not None:
            for i in range(len(data)):
                for key in event_id.keys():
                    data[i][key] = data[i]['_value'] in event_id[key]
        self._data = pd.DataFrame(data)

    def __repr__(self):
        return self._data.__repr__()

    def __getitem__(self, key):
        # For backward compatibility
        if type(key) in [slice, np.ndarray, tuple]:
            # XXX JRK: there's probably a more elegant way to get this with pd
            _data = np.vstack([np.array(self._data['_start']),
                               np.array(self._data['_length']),
                               np.array(self._data['_value'])]).T
            return _data.__getitem__(key)
        elif type(key) in [list, int]:
            return self._data.iloc[key]
        else:
            return self._data[key]

# Setup data
data_path = sample.data_path()
events_fname = data_path + '/MEG/sample/sample_audvis_filt-0-40_raw-eve.fif'
events = mne.read_events(events_fname)
event_id = {'Audio': [1, 2], 'Visual': [3, 4],
            'Left': [1, 3], 'Right': [2, 4]}

events_df = _BaseEvents(events, event_id)
print events_df
print np.where(events_df['Audio'])[0]
print np.where(events_df['Visual' and 'Left'])[0]
print events_df[:, 2]