"""Tools for working with epoched data"""

# Authors: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
#          Matti Hamalainen <msh@nmr.mgh.harvard.edu>
#          Daniel Strohmeier <daniel.strohmeier@tu-ilmenau.de>
#          Denis Engemann <denis.engemann@gmail.com>
#          Mainak Jas <mainak@neuro.hut.fi>
#
# License: BSD (3-clause)

import copy as cp
import warnings
import json
import inspect
import os.path as op
from distutils.version import LooseVersion

import numpy as np
import scipy

from .io.write import (start_file, start_block, end_file, end_block,
                       write_int, write_float_matrix, write_float,
                       write_id, write_string, _get_split_size)
from .io.meas_info import read_meas_info, write_meas_info, _merge_info
from .io.open import fiff_open, _get_next_fname
from .io.tree import dir_tree_find
from .io.tag import read_tag
from .io.constants import FIFF
from .io.pick import (pick_types, channel_indices_by_type, channel_type,
                      pick_channels, pick_info)
from .io.proj import setup_proj, ProjMixin, _proj_equal
from .io.base import _BaseRaw, ToDataFrameMixin
from .evoked import EvokedArray, _aspect_rev
from .baseline import rescale
from .channels.channels import (ContainsMixin, UpdateChannelsMixin,
                                SetChannelsMixin, InterpolationMixin)
from .filter import resample, detrend, FilterMixin
from .event import _read_events_fif
from .fixes import in1d
from .viz import (plot_epochs, _drop_log_stats,
                  plot_epochs_psd, plot_epochs_psd_topomap)
from .utils import (check_fname, logger, verbose, _check_type_picks,
                    _time_mask, check_random_state, object_hash)
from .externals.six import iteritems, string_types
from .externals.six.moves import zip


def _save_split(epochs, fname, part_idx, n_parts):
    """Split epochs"""

    # insert index in filename
    path, base = op.split(fname)
    idx = base.find('.')
    if part_idx > 0:
        fname = op.join(path, '%s-%d.%s' % (base[:idx], part_idx,
                                            base[idx + 1:]))

    next_fname = None
    if part_idx < n_parts - 1:
        next_fname = op.join(path, '%s-%d.%s' % (base[:idx], part_idx + 1,
                                                 base[idx + 1:]))
        next_idx = part_idx + 1

    fid = start_file(fname)

    info = epochs.info
    meas_id = info['meas_id']

    start_block(fid, FIFF.FIFFB_MEAS)
    write_id(fid, FIFF.FIFF_BLOCK_ID)
    if info['meas_id'] is not None:
        write_id(fid, FIFF.FIFF_PARENT_BLOCK_ID, info['meas_id'])

    # Write measurement info
    write_meas_info(fid, info)

    # One or more evoked data sets
    start_block(fid, FIFF.FIFFB_PROCESSED_DATA)
    start_block(fid, FIFF.FIFFB_EPOCHS)

    # write events out after getting data to ensure bad events are dropped
    data = epochs.get_data()
    start_block(fid, FIFF.FIFFB_MNE_EVENTS)
    write_int(fid, FIFF.FIFF_MNE_EVENT_LIST, epochs.events.T)
    mapping_ = ';'.join([k + ':' + str(v) for k, v in
                         epochs.event_id.items()])
    write_string(fid, FIFF.FIFF_DESCRIPTION, mapping_)
    end_block(fid, FIFF.FIFFB_MNE_EVENTS)

    # First and last sample
    first = int(epochs.times[0] * info['sfreq'])
    last = first + len(epochs.times) - 1
    write_int(fid, FIFF.FIFF_FIRST_SAMPLE, first)
    write_int(fid, FIFF.FIFF_LAST_SAMPLE, last)

    # save baseline
    if epochs.baseline is not None:
        bmin, bmax = epochs.baseline
        bmin = epochs.times[0] if bmin is None else bmin
        bmax = epochs.times[-1] if bmax is None else bmax
        write_float(fid, FIFF.FIFF_MNE_BASELINE_MIN, bmin)
        write_float(fid, FIFF.FIFF_MNE_BASELINE_MAX, bmax)

    # The epochs itself
    decal = np.empty(info['nchan'])
    for k in range(info['nchan']):
        decal[k] = 1.0 / (info['chs'][k]['cal'] *
                          info['chs'][k].get('scale', 1.0))

    data *= decal[np.newaxis, :, np.newaxis]

    write_float_matrix(fid, FIFF.FIFF_EPOCH, data)

    # undo modifications to data
    data /= decal[np.newaxis, :, np.newaxis]

    write_string(fid, FIFF.FIFFB_MNE_EPOCHS_DROP_LOG,
                 json.dumps(epochs.drop_log))

    write_int(fid, FIFF.FIFFB_MNE_EPOCHS_SELECTION,
              epochs.selection)

    # And now write the next file info in case epochs are split on disk
    if next_fname is not None and n_parts > 1:
        start_block(fid, FIFF.FIFFB_REF)
        write_int(fid, FIFF.FIFF_REF_ROLE, FIFF.FIFFV_ROLE_NEXT_FILE)
        write_string(fid, FIFF.FIFF_REF_FILE_NAME, op.basename(next_fname))
        if meas_id is not None:
            write_id(fid, FIFF.FIFF_REF_FILE_ID, meas_id)
        write_int(fid, FIFF.FIFF_REF_FILE_NUM, next_idx)
        end_block(fid, FIFF.FIFFB_REF)

    end_block(fid, FIFF.FIFFB_EPOCHS)
    end_block(fid, FIFF.FIFFB_PROCESSED_DATA)
    end_block(fid, FIFF.FIFFB_MEAS)
    end_file(fid)


class _BaseEpochs(ProjMixin, ContainsMixin, UpdateChannelsMixin,
  
                SetChannelsMixin, InterpolationMixin, FilterMixin,
                  ToDataFrameMixin):
    """Abstract base class for Epochs-type classes

    This class provides basic functionality and should never be instantiated
    directly. See Epochs below for an e
xplanation of the parameters.
    """
    def __init__(self, info, data, events, event_id, tmin, tmax,
                 baseline=(None, 0), raw=None,
                 picks=None, name='Unknown', reject=None, flat=None,
                 decim=1, reject_tmin=None, reject_tmax=None, detrend=None,
                 add_eeg_ref=True, proj=True, on_missing='error',
                 preload_at_end=False, verbose=None):

        self.verbose = verbose
        self.name = name

        if on_missing not in ['error', 'warning', 'ignore']:
            raise ValueError('on_missing must be one of: error, '
                             'warning, ignore. Got: %s' % on_missing)

        # check out event_id dict
        if event_id is None:  # convert to int to make typing-checks happy
            event_id = dict((str(e), int(e)) for e in np.unique(events[:, 2]))
        elif isinstance(event_id, dict):
            if not all(isinstance(v, int) for v in event_id.values()):
                raise ValueError('Event IDs must be of type integer')
            if not all(isinstance(k, string_types) for k in event_id):
                raise ValueError('Event names must be of type str')
        elif isinstance(event_id, list):
            if not all(isinstance(v, int) for v in event_id):
                raise ValueError('Event IDs must be of type integer')
            event_id = dict(zip((str(i) for i in event_id), event_id))
        elif isinstance(event_id, int):
            event_id = {str(event_id): event_id}
        else:
            raise ValueError('event_id must be dict or int.')
        self.event_id = event_id
        del event_id

        if events is not None:  # RtEpochs can have events=None

            if events.dtype.kind not in ['i', 'u']:
                raise ValueError('events must be an array of type int')
            if events.ndim != 2 or events.shape[1] != 3:
                raise ValueError('events must be 2D with 3 columns')

            for key, val in self.event_id.items():
                if val not in events[:, 2]:
                    msg = ('No matching events found for %s '
                           '(event id %i)' % (key, val))
                    if on_missing == 'error':
                        raise ValueError(msg)
                    elif on_missing == 'warning':
                        logger.warning(msg)
                        warnings.warn(msg)
                    else:  # on_missing == 'ignore':
                        pass

            values = list(self.event_id.values())
            selected = in1d(events[:, 2], values)
            self.selection = np.where(selected)[0]
            self.drop_log = [list() if k in self.selection else ['IGNORED']
                             for k in range(len(events))]
            events = events[selected]
            n_events = len(events)
            if n_events > 1:
          