# Authors: Alexandre Gramfort <gramfort@nmr.mgh.harvard.edu>
#          Matti Hamalainen <msh@nmr.mgh.harvard.edu>
#          Martin Luessi <mluessi@nmr.mgh.harvard.edu>
#          Denis Engemann <d.engemann@fz-juelich.de>
#
# License: BSD (3-clause)

from ..externals.six import string_types
from math import floor, ceil
import copy
from copy import deepcopy
import warnings
import os
import os.path as op

import numpy as np
from scipy.signal import hilbert
from scipy import linalg

from .constants import FIFF
from .open import fiff_open
from .meas_info import read_meas_info, write_meas_info
from .tree import dir_tree_find
from .tag import read_tag
from .pick import pick_types, channel_type
from .proj import (setup_proj, activate_proj, proj_equal, ProjMixin,
                   _has_eeg_average_ref_proj, make_eeg_average_ref_proj)
from .channels import ContainsMixin, DropChannelsMixin
from .compensator import get_current_comp, set_current_comp, make_compensator

from ..filter import (low_pass_filter, high_pass_filter, band_pass_filter,
                      notch_filter, band_stop_filter, resample)
from ..parallel import parallel_func
from ..utils import (_check_fname, estimate_rank, _check_pandas_installed,
                     logger, verbose)
from ..viz import plot_raw, plot_raw_psds, _mutable_defaults

class Raw(ProjMixin, ContainsMixin, DropChannelsMixin):
    """Raw data

    Parameters
    ----------
    fnames : list, or string
        A list of the raw files to treat as a Raw instance, or a single
        raw file.
    allow_maxshield : bool, (default False)
        allow_maxshield if True, allow loading of data that has been
        processed with Maxshield. Maxshield-processed data should generally
        not be loaded directly, but should be processed using SSS first.
    preload : bool or str (default False)
        Preload data into memory for data manipulation and faster indexing.
        If True, the data will be preloaded into memory (fast, requires
        large amount of memory). If preload is a string, preload is the
        file name of a memory-mapped file which is used to store the data
        on the hard drive (slower, requires less memory).
    proj : bool
        Apply the signal space projection (SSP) operators present in
        the file to the data. Note: Once the projectors have been
        applied, they can no longer be removed. It is usually not
        recommended to apply the projectors at this point as they are
        applied automatically later on (e.g. when computing inverse
        solutions).
    compensation : None | int
        If None the compensation in the data is not modified.
        If set to n, e.g. 3, apply gradient compensation of grade n as
        for CTF systems.
    add_eeg_ref : bool
        If True, add average EEG reference projector (if it's not already
        present).
    verbose : bool, str, int, or None
        If not None, override default verbose level (see mne.verbose).

    Attributes
    ----------
    info : dict
        Measurement info.
    ch_names : list of string
        List of channels' names.
    n_times : int
        Total number of time points in the raw file.
    verbose : bool, str, int, or None
        See above.
    """
    @verbose
    def __init__(self, fnames, allow_maxshield=False, preload=False,
                 proj=False, compensation=None, add_eeg_ref=True,
                 verbose=None):

        if not isinstance(fnames, list):
            fnames = [op.abspath(fnames)] if not op.isabs(fnames) else [fnames]
        else:
            fnames = [op.abspath(f) if not op.isabs(f) else f for f in fnames]

        raws = [self._read_raw_file(fname, allow_maxshield, preload,
                                    compensation) for fname in fnames]
        _check_raw_compatibility(raws)

        # combine information from each raw file to construct self
        self.first_samp = raws[0].first_samp  # meta first sample
        self._first_samps = np.array([r.first_samp for r in raws])
        self._last_samps = np.array([r.last_samp for r in raws])
        self._raw_lengths = np.array([r.n_times for r in raws])
        self.last_samp = self.first_samp + sum(self._raw_lengths) - 1
        self.cals = raws[0].cals
        self.rawdirs = [r.rawdir for r in raws]
        self.comp = copy.deepcopy(raws[0].comp)
        self._orig_comp_grade = raws[0]._orig_comp_grade
        self.fids = [r.fid for r in raws]
        self.info = copy.deepcopy(raws[0].info)
        self.verbose = verbose
        self.info['filenames'] = fnames
        self.orig_format = raws[0].orig_format
        self.proj = False
        self._add_eeg_ref(add_eeg_ref)

        if preload:
            self._preload_data(preload)
        else:
            self._preloaded = False

        self._projector = None
        # setup the SSP projector
        self.proj = proj
        if proj:
            self.apply_proj()

    def __del__(self):
        # remove file for memmap
        if hasattr(self, '_data') and hasattr(self._data, 'filename'):
            # First, close the file out; happens automatically on del
            filename = self._data.filename
            del self._data
            # Now file can be removed
            os.remove(filename)

    def __enter__(self):
        """ Entering with block """
        return self

    def __exit__(self, exception_type, exception_val, trace):
        """ Exiting with block """
        try:
            self.close()
        except:
            return exception_type, exception_val, trace

    def _add_eeg_ref(self, add_eeg_ref):
        """Helper to add an average EEG reference"""
        if add_eeg_ref:
            eegs = pick_types(self.info, meg=False, eeg=True, ref_meg=False)
            projs = self.info['projs']
            if len(eegs) > 0 and not _has_eeg_average_ref_proj(projs):
                eeg_ref = make_eeg_average_ref_proj(self.info, activate=False)
                projs.append(eeg_ref)

    def _preload_data(self, preload):
        """This function actually preloads the data"""
        if isinstance(preload, string_types):
            # we will use a memmap: preload is a filename
            data_buffer = preload
        else:
            data_buffer = None

        self._data, self._times = self._read_segment(data_buffer=data_buffer)
        self._preloaded = True
        # close files once data are preloaded
        self.close()

    @verbose
    def _read_raw_file(self, fname, allow_maxshield, preload, compensation,
                       verbose=None):
        """Read in header information from a raw file"""
        logger.info('Opening raw data file %s...' % fname)

        #   Read in the whole file if preload is on and .fif.gz (saves time)
        ext = os.path.splitext(fname)[1].lower()
        whole_file = preload if '.gz' in ext else False
        fid, tree, _ = fiff_open(fname, preload=whole_file)

        #   Read the measurement info
        info, meas = read_meas_info(fid, tree)

        #   Locate the data of interest
        raw_node = dir_tree_find(meas, FIFF.FIFFB_RAW_DATA)
        if len(raw_node) == 0:
            raw_node = dir_tree_find(meas, FIFF.FIFFB_CONTINUOUS_DATA)
            if allow_maxshield:
                raw_node = dir_tree_find(meas, FIFF.FIFFB_SMSH_RAW_DATA)
                if len(raw_node) == 0:
                    raise ValueError('No raw data in %s' % fname)
            else:
                if len(raw_node) == 0:
                    raise ValueError('No raw data in %s' % fname)

        if len(raw_node) == 1:
            raw_node = raw_node[0]

        #   Set up the output structure
        info['filename'] = fname

        #   Process the directory
        directory = raw_node['directory']
        nent = raw_node['nent']
        nchan = int(info['nchan'])
        first = 0
        first_samp = 0
        first_skip = 0

        #   Get first sample tag if it is there
        if directory[first].kind == FIFF.FIFF_FIRST_SAMPLE:
            tag = read_tag(fid, directory[first].pos)
            first_samp = int(tag.data)
            first += 1

        #   Omit initial skip
        if directory[first].kind == FIFF.FIFF_DATA_SKIP:
            # This first skip can be applied only after we know the buffer size
            tag = read_tag(fid, directory[first].pos)
            first_skip = int(tag.data)
            first += 1

        #  Get first sample tag if it is there
        if directory[first].kind == FIFF.FIFF_FIRST_SAMPLE:
            tag = read_tag(fid, directory[first].pos)
            first_samp += int(tag.data)
            first += 1

        raw = _RawShell()
        raw.first_samp = first_samp

        #   Go through the remaining tags in the directory
        rawdir = list()
        nskip = 0
        orig_format = None
        for k in range(first, nent):
            ent = directory[k]
            if ent.kind == FIFF.FIFF_DATA_SKIP:
                tag = read_tag(fid, ent.pos)
                nskip = int(tag.data)
            elif ent.kind == FIFF.FIFF_DATA_BUFFER:
                #   Figure out the number of samples in this buffer
                if ent.type == FIFF.FIFFT_DAU_PACK16:
                    nsamp = ent.size // (2 * nchan)
                elif ent.type == FIFF.FIFFT_SHORT:
                    nsamp = ent.size // (2 * nchan)
                elif ent.type == FIFF.FIFFT_FLOAT:
                    nsamp = ent.size // (4 * nchan)
                elif ent.type == FIFF.FIFFT_DOUBLE:
                    nsamp = ent.size // (8 * nchan)
                elif ent.type == FIFF.FIFFT_INT:
                    nsamp = ent.size // (4 * nchan)
                elif ent.type == FIFF.FIFFT_COMPLEX_FLOAT:
                    nsamp = ent.size // (8 * nchan)
                elif ent.type == FIFF.FIFFT_COMPLEX_DOUBLE:
                    nsamp = ent.size // (16 * nchan)
                else:
                    fid.close()
                    raise ValueError('Cannot handle data buffers of type %d' %
                                     ent.type)
                if orig_format is None:
                    if ent.type == FIFF.FIFFT_DAU_PACK16:
                        orig_format = 'short'
                    elif ent.type == FIFF.FIFFT_SHORT:
                        orig_format = 'short'
                    elif ent.type == FIFF.FIFFT_FLOAT:
                        orig_format = 'single'
                    elif ent.type == FIFF.FIFFT_DOUBLE:
                        orig_format = 'double'
                    elif ent.type == FIFF.FIFFT_INT:
                        orig_format = 'int'
                    elif ent.type == FIFF.FIFFT_COMPLEX_FLOAT:
                        orig_format = 'single'
                    elif ent.type == FIFF.FIFFT_COMPLEX_DOUBLE:
                        orig_format = 'double'

                #  Do we have an initial skip pending?
                if first_skip > 0:
                    first_samp += nsamp * first_skip
                    raw.first_samp = first_samp
                    first_skip = 0

                #  Do we have a skip pending?
                if nskip > 0:
                    rawdir.append(dict(ent=None, first=first_samp,
                                       last=first_samp + nskip * nsamp - 1,
                                       nsamp=nskip * nsamp))
                    first_samp += nskip * nsamp
                    nskip = 0

                #  Add a data buffer
                rawdir.append(dict(ent=ent, first=first_samp,
                                   last=first_samp + nsamp - 1,
                                   nsamp=nsamp))
                first_samp += nsamp

        raw.last_samp = first_samp - 1
        raw.orig_format = orig_format

        #   Add the calibration factors
        cals = np.zeros(info['nchan'])
        for k in range(info['nchan']):
            cals[k] = info['chs'][k]['range'] * info['chs'][k]['cal']

        raw.cals = cals
        raw.rawdir = rawdir
        raw.comp = None
        raw._orig_comp_grade = None

        #   Set up the CTF compensator
        current_comp = get_current_comp(info)
        if current_comp is not None:
            logger.info('Current compensation grade : %d' % current_comp)

        if compensation is not None:
            raw.comp = make_compensator(info, current_comp, compensation)
            if raw.comp is not None:
                logger.info('Appropriate compensator added to change to '
                            'grade %d.' % (compensation))
                raw._orig_comp_grade = current_comp
                set_current_comp(info, compensation)

        logger.info('    Range : %d ... %d =  %9.3f ... %9.3f secs' % (
                    raw.first_samp, raw.last_samp,
                    float(raw.first_samp) / info['sfreq'],
                    float(raw.last_samp) / info['sfreq']))

        # store the original buffer size
        info['buffer_size_sec'] = (np.median([r['nsamp'] for r in rawdir])
                                   / info['sfreq'])

        raw.fid = fid
        raw.info = info
        raw.verbose = verbose

        logger.info('Ready.')

        return raw

    def _parse_get_set_params(self, item):
        # make sure item is a tuple
        if not isinstance(item, tuple):  # only channel selection passed
            item = (item, slice(None, None, None))

        if len(item) != 2:  # should be channels and time instants
            raise RuntimeError("Unable to access raw data (need both channels "
                               "and time)")

        if isinstance(item[0], slice):
            start = item[0].start if item[0].start is not None else 0
            nchan = self.info['nchan']
            stop = item[0].stop if item[0].stop is not None else nchan
            step = item[0].step if item[0].step is not None else 1
            sel = list(range(start, stop, step))
        else:
            sel = item[0]

        if isinstance(item[1], slice):
            time_slice = item[1]
            start, stop, step = (time_slice.start, time_slice.stop,
                                 time_slice.step)
        else:
            item1 = item[1]
            # Let's do automated type conversion to integer here
            if np.array(item[1]).dtype.kind == 'i':
                item1 = int(item1)
            if isinstance(item1, int):
                start, stop, step = item1, item1 + 1, 1
            else:
                raise ValueError('Must pass int or slice to __getitem__')

        if start is None:
            start = 0
        if (step is not None) and (step is not 1):
            raise ValueError('step needs to be 1 : %d given' % step)

        if isinstance(sel, int):
            sel = np.array([sel])

        if sel is not None and len(sel) == 0:
            raise ValueError("Empty channel list")

        return sel, start, stop

    def __getitem__(self, item):
        """getting raw data content with python slicing"""
        sel, start, stop = self._parse_get_set_params(item)
        if self._preloaded:
            data, times = self._data[sel, start:stop], self._times[start:stop]
        else:
            data, times = self._read_segment(start=start, stop=stop, sel=sel,
                                             projector=self._projector,
                                             verbose=self.verbose)
        return data, times

    def __setitem__(self, item, value):
        """setting raw data content with python slicing"""
        if not self._preloaded:
            raise RuntimeError('Modifying data of Raw is only supported '
                               'when preloading is used. Use preload=True '
                               '(or string) in the constructor.')
        sel, start, stop = self._parse_get_set_params(item)
        # set the data
        self._data[sel, start:stop] = value

    def anonymize(self):
        """Anonymize data

        This function will remove info['subject_info'] if it exists."""
        self.info._anonymize()

    @verbose
    def apply_function(self, fun, picks, dtype, n_jobs, verbose=None, *args,
                       **kwargs):
        """ Apply a function to a subset of channels.

        The function "fun" is applied to the channels defined in "picks". The
        data of the Raw object is modified inplace. If the function returns
        a different data type (e.g. numpy.complex) it must be specified using
        the dtype parameter, which causes the data type used for representing
        the raw data to change.

        The Raw object has to be constructed using preload=True (or string).

        Note: If n_jobs > 1, more memory is required as "len(picks) * n_times"
              additional time points need to be temporaily stored in memory.

        Note: If the data type changes (dtype != None), more memory is required
              since the original and the converted data needs to be stored in
              memory.

        Parameters
        ----------
        fun : function
            A function to be applied to the channels. The first argument of
            fun has to be a timeseries (numpy.ndarray). The function must
            return an numpy.ndarray with the same size as the input.
        picks : list of int
            Indices of channels to apply the function to.
        dtype : numpy.dtype
            Data type to use for raw data after applying the function. If None
            the data type is not modified.
        n_jobs: int
            Number of jobs to run in parallel.
        verbose : bool, str, int, or None
            If not None, override default verbose level (see mne.verbose).
            Defaults to self.verbose.
        *args :
            Additional positional arguments to pass to fun (first pos. argument
            of fun is the timeseries of a channel).
        **kwargs :
            Keyword arguments to pass to fun.
        """
        if not self._preloaded:
            raise RuntimeError('Raw data needs to be preloaded. Use '
                               'preload=True (or string) in the constructor.')

        if not callable(fun):
            raise ValueError('fun needs to be a function')

        data_in = self._data
        if dtype is not None and dtype != self._data.dtype:
            self._data = self._data.astype(dtype)

        if n_jobs == 1:
            # modify data inplace to save memory
            for idx in picks:
                self._data[idx, :] = fun(data_in[idx, :], *args, **kwargs)
        else:
            # use parallel function
            parallel, p_fun, _ = parallel_func(fun, n_jobs)
            data_picks_new = parallel(p_fun(data_in[p], *args, **kwargs)
                                      for p in picks)
            for pp, p in enumerate(picks):
                self._data[p, :] = data_picks_new[pp]

    @verbose
    def apply_hilbert(self, picks, envelope=False, n_jobs=1, verbose=None):
        """ Compute analytic signal or envelope for a subset of channels.

        If envelope=False, the analytic signal for the channels defined in
        "picks" is computed and the data of the Raw object is converted to
        a complex representation (the analytic signal is complex valued).

        If envelope=True, the absolute value of the analytic signal for the
        channels defined in "picks" is computed, resulting in the envelope
        signal.

        Note: DO NOT use envelope=True if you intend to compute an inverse
              solution from the raw data. If you want to compute the
              envelope in source space, use envelope=False and compute the
              envelope after the inverse solution has been obtained.

        Note: If envelope=False, more memory is required since the original
              raw data as well as the analytic signal have temporarily to
              be stored in memory.

        Note: If n_jobs > 1 and envelope=True, more memory is required as
              "len(picks) * n_times" additional time points need to be
              temporaily stored in memory.

        Parameters
        ----------
        picks : list of int
            Indices of channels to apply the function to.
        envelope : bool (default: False)
            Compute the envelope signal of each channel.
        n_jobs: int
            Number of jobs to run in parallel.
        verbose : bool, str, int, or None
            If not None, override default verbose level (see mne.verbose).
            Defaults to self.verbose.

        Notes
        -----
        The analytic signal "x_a(t)" of "x(t)" is::

            x_a = F^{-1}(F(x) 2U) = x + i y

        where "F" is the Fourier transform, "U" the unit step function,
        and "y" the Hilbert transform of "x". One usage of the analytic
        signal is the computation of the envelope signal, which is given by
        "e(t) = abs(x_a(t))". Due to the linearity of Hilbert transform and the
        MNE inverse solution, the enevlope in source space can be obtained
        by computing the analytic signal in sensor space, applying the MNE
        inverse, and computing the envelope in source space.
        """
        if envelope:
            self.apply_function(_envelope, picks, None, n_jobs)
        else:
            self.apply_function(hilbert, picks, np.complex64, n_jobs)

    @verbose
    def filter(self, l_freq, h_freq, picks=None, filter_length='10s',
               l_trans_bandwidth=0.5, h_trans_bandwidth=0.5, n_jobs=1,
               method='fft', iir_params=None, verbose=None):
        """Filter a subset of channels.

        Applies a zero-phase low-pass, high-pass, band-pass, or band-stop
        filter to the channels selected by "picks". The data of the Raw
        object is modified inplace.

        The Raw object has to be constructed using preload=True (or string).

        l_freq and h_freq are the frequencies below which and above which,
        respectively, to filter out of the data. Thus the uses are:
            l_freq < h_freq: band-pass filter
            l_freq > h_freq: band-stop filter
            l_freq is not None, h_freq is None: low-pass filter
            l_freq is None, h_freq is not None: high-pass filter

        Note: If n_jobs > 1, more memory is required as "len(picks) * n_times"
              additional time points need to be temporarily stored in memory.

        Note: self.info['lowpass'] and self.info['highpass'] are only updated
              with picks=None.

        Parameters
        ----------
        l_freq : float | None
            Low cut-off frequency in Hz. If None the data are only low-passed.
        h_freq : float | None
            High cut-off frequency in Hz. If None the data are only
            high-passed.
        picks : list of int | None
            Indices of channels to filter. If None only the data (MEG/EEG)
            channels will be filtered.
        filter_length : str (Default: '10s') | int | None
            Length of the filter to use. If None or "len(x) < filter_length",
            the filter length used is len(x). Otherwise, if int, overlap-add
            filtering with a filter of the specified length in samples) is
            used (faster for long signals). If str, a human-readable time in
            units of "s" or "ms" (e.g., "10s" or "5500ms") will be converted
            to the shortest power-of-two length at least that duration.
        l_trans_bandwidth : float
            Width of the transition band at the low cut-off frequency in Hz.
        h_trans_bandwidth : float
            Width of the transition band at the high cut-off frequency in Hz.
        n_jobs : int | str
            Number of jobs to run in parallel. Can be 'cuda' if scikits.cuda
            is installed properly, CUDA is initialized, and method='fft'.
        method : str
            'fft' will use overlap-add FIR filtering, 'iir' will use IIR
            forward-backward filtering (via filtfilt).
        iir_params : dict | None
            Dictionary of parameters to use for IIR filtering.
            See mne.filter.construct_iir_filter for details. If iir_params
            is None and method="iir", 4th order Butterworth will be used.
        verbose : bool, str, int, or None
            If not None, override default verbose level (see mne.verbose).
            Defaults to self.verbose.
        """
        if verbose is None:
            verbose = self.verbose
        fs = float(self.info['sfreq'])
        if l_freq == 0:
            l_freq = None
        if h_freq is not None and h_freq > (fs / 2.):
            h_freq = None
        if l_freq is not None and not isinstance(l_freq, float):
            l_freq = float(l_freq)
        if h_freq is not None and not isinstance(h_freq, float):
            h_freq = float(h_freq)

        if not self._preloaded:
            raise RuntimeError('Raw data needs to be preloaded to filter. Use '
                               'preload=True (or string) in the constructor.')
        if picks is None:
            if 'ICA ' in ','.join(self.ch_names):
                pick_parameters = dict(misc=True, ref_meg=False)
            else:
                pick_parameters = dict(meg=True, eeg=True, ref_meg=False)
            picks = pick_types(self.info, exclude=[], **pick_parameters)
            # let's be safe.
            if len(picks) < 1:
                raise RuntimeError('Could not find any valid channels for '
                                   'your Raw object. Please contact the '
                                   'MNE-Python developers.')

            # update info if filter is applied to all data channels,
            # and it's not a band-stop filter
            if h_freq is not None and (l_freq is None or l_freq < h_freq) and \
                    h_freq < self.info['lowpass']:
                self.info['lowpass'] = h_freq
            if l_freq is not None and (h_freq is None or l_freq < h_freq) and \
                    l_freq > self.info['highpass']:
                self.info['highpass'] = l_freq
        if l_freq is None and h_freq is not None:
            logger.info('Low-pass filtering at %0.2g Hz' % h_freq)
            low_pass_filter(self._data, fs, h_freq,
                            filter_length=filter_length,
                            trans_bandwidth=l_trans_bandwidth, method=method,
                            iir_params=iir_params, picks=picks, n_jobs=n_jobs,
                            copy=False)
        if l_freq is not None and h_freq is None:
            logger.info('High-pass filtering at %0.2g Hz' % l_freq)
            high_pass_filter(self._data, fs, l_freq,
                             filter_length=filter_length,
                             trans_bandwidth=h_trans_bandwidth, method=method,
                             iir_params=iir_params, picks=picks, n_jobs=n_jobs,
                             copy=False)
        if l_freq is not None and h_freq is not None:
            if l_freq < h_freq:
                logger.info('Band-pass filtering from %0.2g - %0.2g Hz'
                            % (l_freq, h_freq))
                self._data = band_pass_filter(self._data, fs, l_freq, h_freq,
                    filter_length=filter_length,
                    l_trans_bandwidth=l_trans_bandwidth,
                    h_trans_bandwidth=h_trans_bandwidth,
                    method=method, iir_params=iir_params, picks=picks,
                    n_jobs=n_jobs, copy=False)
            else:
                logger.info('Band-stop filtering from %0.2g - %0.2g Hz'
                            % (h_freq, l_freq))
                self._data = band_stop_filter(self._data, fs, h_freq, l_freq,
                    filter_length=filter_length,
                    l_trans_bandwidth=h_trans_bandwidth,
                    h_trans_bandwidth=l_trans_bandwidth, method=method,
                    iir_params=iir_params, picks=picks, n_jobs=n_jobs,
                    copy=False)

    @verbose
    def notch_filter(self, freqs, picks=None, filter_length='10s',
                     notch_widths=None, trans_bandwidth=1.0, n_jobs=1,
                     method='fft', iir_params=None,
                     mt_bandwidth=None, p_value=0.05, verbose=None):
        """Notch filter a subset of channels.

        Applies a zero-phase notch filter to the channels selected by
        "picks". The data of the Raw object is modified inplace.

        The Raw object has to be constructed using preload=True (or string).

        Note: If n_jobs > 1, more memory is required as "len(picks) * n_times"
              additional time points need to be temporaily stored in memory.

        Parameters
        ----------
        freqs : float | array of float | None
            Specific frequencies to filter out from data, e.g.,
            np.arange(60, 241, 60) in the US or np.arange(50, 251, 50) in
            Europe. None can only be used with the mode 'spectrum_fit',
            where an F test is used to find sinusoidal components.
        picks : list of int | None
            Indices of channels to filter. If None only the data (MEG/EEG)
            channels will be filtered.
        filter_length : str (Default: '10s') | int | None
            Length of the filter to use. If None or "len(x) < filter_length",
            the filter length used is len(x). Otherwise, if int, overlap-add
            filtering with a filter of the specified length in samples) is
            used (faster for long signals). If str, a human-readable time in
            units of "s" or "ms" (e.g., "10s" or "5500ms") will be converted
            to the shortest power-of-two length at least that duration.
        notch_widths : float | array of float | None
            Width of each stop band (centred at each freq in freqs) in Hz.
            If None, freqs / 200 is used.
        trans_bandwidth : float
            Width of the transition band in Hz.
        n_jobs : int | str
            Number of jobs to run in parallel. Can be 'cuda' if scikits.cuda
            is installed properly, CUDA is initialized, and method='fft'.
        method : str
            'fft' will use overlap-add FIR filtering, 'iir' will use IIR
            forward-backward filtering (via filtfilt). 'spectrum_fit' will
            use multi-taper estimation of sinusoidal components.
        iir_params : dict | None
            Dictionary of parameters to use for IIR filtering.
            See mne.filter.construct_iir_filter for details. If iir_params
            is None and method="iir", 4th order Butterworth will be used.
        mt_bandwidth : float | None
            The bandwidth of the multitaper windowing function in Hz.
            Only used in 'spectrum_fit' mode.
        p_value : float
            p-value to use in F-test thresholding to determine significant
            sinusoidal components to remove when method='spectrum_fit' and
            freqs=None. Note that this will be Bonferroni corrected for the
            number of frequencies, so large p-values may be justified.
        verbose : bool, str, int, or None
            If not None, override default verbose level (see mne.verbose).
            Defaults to self.verbose.

        Notes
        -----
        For details, see mne.filter.notch_filter.
        """
        if verbose is None:
            verbose = self.verbose
        fs = float(self.info['sfreq'])
        if picks is None:
            if 'ICA ' in ','.join(self.ch_names):
                pick_parameters = dict(misc=True)
            else:
                pick_parameters = dict(meg=True, eeg=True)
            picks = pick_types(self.info, exclude=[], **pick_parameters)
            # let's be safe.
            if len(picks) < 1:
                raise RuntimeError('Could not find any valid channels for '
                                   'your Raw object. Please contact the '
                                   'MNE-Python developers.')
        if not self._preloaded:
            raise RuntimeError('Raw data needs to be preloaded to filter. Use '
                               'preload=True (or string) in the constructor.')

        self._data = notch_filter(self._data, fs, freqs,
                                  filter_length=filter_length,
                                  notch_widths=notch_widths,
                                  trans_bandwidth=trans_bandwidth,
                                  method=method, iir_params=iir_params,
                                  mt_bandwidth=mt_bandwidth, p_value=p_value,
                                  picks=picks, n_jobs=n_jobs, copy=False)

    @verbose
    def resample(self, sfreq, npad=100, window='boxcar',
                 stim_picks=None, n_jobs=1, verbose=None):
        """Resample data channels.

        Resamples all channels. The data of the Raw object is modified inplace.

        The Raw object has to be constructed using preload=True (or string).

        WARNING: The intended purpose of this function is primarily to speed
        up computations (e.g., projection calculation) when precise timing
        of events is not required, as downsampling raw data effectively
        jitters trigger timings. It is generally recommended not to epoch
        downsampled data, but instead epoch and then downsample, as epoching
        downsampled data jitters triggers.

        Parameters
        ----------
        sfreq : float
            New sample rate to use.
        npad : int
            Amount to pad the start and end of the data.
        window : string or tuple
            Window to use in resampling. See scipy.signal.resample.
        stim_picks : array of int | None
            Stim channels. These channels are simply subsampled or
            supersampled (without applying any filtering). This reduces
            resampling artifacts in stim channels, but may lead to missing
            triggers. If None, stim channels are automatically chosen using
            mne.fiff.pick_types(raw.info, meg=False, stim=True, exclude=[]).
        n_jobs : int | str
            Number of jobs to run in parallel. Can be 'cuda' if scikits.cuda
            is installed properly and CUDA is initialized.
        verbose : bool, str, int, or None
            If not None, override default verbose level (see mne.verbose).
            Defaults to self.verbose.

        Notes
        -----
        For some data, it may be more accurate to use npad=0 to reduce
        artifacts. This is dataset dependent -- check your data!
        """
        if not self._preloaded:
            raise RuntimeError('Can only resample preloaded data')
        sfreq = float(sfreq)
        o_sfreq = float(self.info['sfreq'])

        offsets = np.concatenate(([0], np.cumsum(self._raw_lengths)))
        new_data = list()
        # set up stim channel processing
        if stim_picks is None:
            stim_picks = pick_types(self.info, meg=False, ref_meg=False,
                                    stim=True, exclude=[])
        stim_picks = np.asanyarray(stim_picks)
        ratio = sfreq / o_sfreq
        for ri in range(len(self._raw_lengths)):
            data_chunk = self._data[:, offsets[ri]:offsets[ri + 1]]
            new_data.append(resample(data_chunk, sfreq, o_sfreq, npad,
                                     n_jobs=n_jobs))
            new_ntimes = new_data[ri].shape[1]

            # Now deal with the stim channels. In empirical testing, it was
            # faster to resample all channels (above) and then replace the
            # stim channels than it was to only resample the proper subset
            # of channels and then use np.insert() to restore the stims

            # figure out which points in old data to subsample
            # protect against out-of-bounds, which can happen (having
            # one sample more than expected) due to padding
            stim_inds = np.minimum(np.floor(np.arange(new_ntimes)
                                            / ratio).astype(int),
                                   data_chunk.shape[1] - 1)
            for sp in stim_picks:
                new_data[ri][sp] = data_chunk[[sp]][:, stim_inds]

            self._first_samps[ri] = int(self._first_samps[ri] * ratio)
            self._last_samps[ri] = self._first_samps[ri] + new_ntimes - 1
            self._raw_lengths[ri] = new_ntimes

        # adjust affected variables
        self._data = np.concatenate(new_data, axis=1)
        self.first_samp = self._first_samps[0]
        self.last_samp = self.first_samp + self._data.shape[1] - 1
        self.info['sfreq'] = sfreq
        self._times = (np.arange(self.n_times, dtype=np.float64)
                       / self.info['sfreq'])

    def crop(self, tmin=0.0, tmax=None, copy=True):
        """Crop raw data file.

        Limit the data from the raw file to go between specific times. Note
        that the new tmin is assumed to be t=0 for all subsequently called
        functions (e.g., time_as_index, or Epochs). New first_samp and
        last_samp are set accordingly. And data are modified in-place when
        called with copy=False.

        Parameters
        ----------
        tmin : float
            New start time (must be >= 0).
        tmax : float | None
            New end time of the data (cannot exceed data duration).
        copy : bool
            If False Raw is cropped in place.

        Returns
        -------
        raw : instance of Raw
            The cropped raw object.
        """
        raw = self.copy() if copy is True else self
        max_time = (raw.n_times - 1) / raw.info['sfreq']
        if tmax is None:
            tmax = max_time

        if tmin > tmax:
            raise ValueError('tmin must be less than tmax')
        if tmin < 0.0:
            raise ValueError('tmin must be >= 0')
        elif tmax > max_time:
            raise ValueError('tmax must be less than or equal to the max raw '
                             'time (%0.4f sec)' % max_time)

        smin = raw.time_as_index(tmin)[0]
        smax = raw.time_as_index(tmax)[0]
        cumul_lens = np.concatenate(([0], np.array(raw._raw_lengths,
                                     dtype='int')))
        cumul_lens = np.cumsum(cumul_lens)
        keepers = np.logical_and(np.less(smin, cumul_lens[1:]),
                                 np.greater_equal(smax, cumul_lens[:-1]))
        keepers = np.where(keepers)[0]
        raw._first_samps = np.atleast_1d(raw._first_samps[keepers])
        # Adjust first_samp of first used file!
        raw._first_samps[0] += smin - cumul_lens[keepers[0]]
        raw._last_samps = np.atleast_1d(raw._last_samps[keepers])
        raw._last_samps[-1] -= cumul_lens[keepers[-1] + 1] - 1 - smax
        raw._raw_lengths = raw._last_samps - raw._first_samps + 1
        raw.fids = [f for fi, f in enumerate(raw.fids) if fi in keepers]
        raw.rawdirs = [r for ri, r in enumerate(raw.rawdirs)
                       if ri in keepers]
        raw.first_samp = raw._first_samps[0]
        raw.last_samp = raw.first_samp + (smax - smin)
        if raw._preloaded:
            raw._data = raw._data[:, smin:smax + 1]
            raw._times = np.arange(raw.n_times) / raw.info['sfreq']
        return raw

    @verbose
    def save(self, fname, picks=None, tmin=0, tmax=None, buffer_size_sec=10,
             drop_small_buffer=False, proj=False, format='single',
             overwrite=False, verbose=None):
        """Save raw data to file

        Parameters
        ----------
        fname : string
            File name of the new dataset. This has to be a new filename
            unless data have been preloaded.
        picks : list of int
            Indices of channels to include.
        tmin : float
            Time in seconds of first sample to save.
        tmax : float
            Time in seconds of last sample to save.
        buffer_size_sec : float | None
            Size of data chunks in seconds. If None, the buffer size of
            the original file is used.
        drop_small_buffer : bool
            Drop or not the last buffer. It is required by maxfilter (SSS)
            that only accepts raw files with buffers of the same size.
        proj : bool
            If True the data is saved with the projections applied (active).
            Note: If apply_proj() was used to apply the projections,
            the projectons will be active even if proj is False.
        format : str
            Format to use to save raw data. Valid options are 'double',
            'single', 'int', and 'short' for 64- or 32-bit float, or 32- or
            16-bit integers, respectively. It is STRONGLY recommended to use
            'single', as this is backward-compatible, and is standard for
            maintaining precision. Note that using 'short' or 'int' may result
            in loss of precision, complex data cannot be saved as 'short',
            and neither complex data types nor real data stored as 'double'
            can be loaded with the MNE command-line tools. See raw.orig_format
            to determine the format the original data were stored in.
        overwrite : bool
            If True, the destination file (if it exists) will be overwritten.
            If False (default), an error will be raised if the file exists.
        verbose : bool, str, int, or None
            If not None, override default verbose level (see mne.verbose).
            Defaults to self.verbose.

        Notes
        -----
        If Raw is a concatenation of several raw files, *be warned* that only
        the measurement information from the first raw file is stored. This
        likely means that certain operations with external tools may not
        work properly on a saved concatenated file (e.g., probably some
        or all forms of SSS). It is recommended not to concatenate and
        then save raw files for this reason.
        """
        fname = op.abspath(fname)
        if not self._preloaded and fname in self.info['filenames']:
            raise ValueError('You cannot save data to the same file.'
                             ' Please use a different filename.')

        if self._preloaded:
            if np.iscomplexobj(self._data):
                warnings.warn('Saving raw file with complex data. Loading '
                              'with command-line MNE tools will not work.')

        type_dict = dict(short=FIFF.FIFFT_DAU_PACK16,
                         int=FIFF.FIFFT_INT,
                         single=FIFF.FIFFT_FLOAT,
                         double=FIFF.FIFFT_DOUBLE)
        if not format in type_dict.keys():
            raise ValueError('format must be "short", "int", "single", '
                             'or "double"')
        reset_dict = dict(short=False, int=False, single=True, double=True)

        data_test = self[0, 0][0]
        if format == 'short' and np.iscomplexobj(data_test):
            raise ValueError('Complex data must be saved as "single" or '
                             '"double", not "short"')

        # check for file existence
        _check_fname(fname, overwrite)

        if proj:
            info = copy.deepcopy(self.info)
            projector, info = setup_proj(info)
            activate_proj(info['projs'], copy=False)
        else:
            info = self.info
            projector = None

        # set the correct compensation grade and make inverse compensator
        inv_comp = None
        if self.comp is not None:
            inv_comp = linalg.inv(self.comp)
            set_current_comp(info, self._orig_comp_grade)

        outfid, cals = start_writing_raw(fname, info, picks, type_dict[format],
                                         reset_range=reset_dict[format])
        #
        #   Set up the reading parameters
        #

        #   Convert to samples
        start = int(floor(tmin * self.info['sfreq']))
        first_samp = self.first_samp + start

        if tmax is None:
            stop = self.last_samp + 1 - self.first_samp
        else:
            stop = int(floor(tmax * self.info['sfreq']))

        if buffer_size_sec is None:
            if 'buffer_size_sec' in self.info:
                buffer_size_sec = self.info['buffer_size_sec']
            else:
                buffer_size_sec = 10.0
        buffer_size = int(ceil(buffer_size_sec * self.info['sfreq']))
        #
        #   Read and write all the data
        #
        if first_samp != 0:
            write_int(outfid, FIFF.FIFF_FIRST_SAMPLE, first_samp)
        for first in range(start, stop, buffer_size):
            last = first + buffer_size
            if last >= stop:
                last = stop + 1

            if picks is None:
                data, times = self[:, first:last]
            else:
                data, times = self[picks, first:last]

            if projector is not None:
                data = np.dot(projector, data)

            if ((drop_small_buffer and (first > start)
                 and (len(times) < buffer_size))):
                logger.info('Skipping data chunk due to small buffer ... '
                            '[done]')
                break
            logger.info('Writing ...')
            write_raw_buffer(outfid, data, cals, format, inv_comp)
            logger.info('[done]')

        finish_writing_raw(outfid)

    def plot(raw, events=None, duration=10.0, start=0.0, n_channels=20,
             bgcolor='w', color=None, bad_color=(0.8, 0.8, 0.8),
             event_color='cyan', scalings=None, remove_dc=True, order='type',
             show_options=False, title=None, show=True, block=False):
        """Plot raw data

        Parameters
        ----------
        raw : instance of Raw
            The raw data to plot.
        events : array | None
            Events to show with vertical bars.
        duration : float
            Time window (sec) to plot in a given time.
        start : float
            Initial time to show (can be changed dynamically once plotted).
        n_channels : int
            Number of channels to plot at once.
        bgcolor : color object
            Color of the background.
        color : dict | color object | None
            Color for the data traces. If None, defaults to:
            `dict(mag='darkblue', grad='b', eeg='k', eog='k', ecg='r', emg='k',
                 ref_meg='steelblue', misc='k', stim='k', resp='k', chpi='k')`
        bad_color : color object
            Color to make bad channels.
        event_color : color object
            Color to use for events.
        scalings : dict | None
            Scale factors for the traces. If None, defaults to:
            `dict(mag=1e-12, grad=4e-11, eeg=20e-6,
                  eog=150e-6, ecg=5e-4, emg=1e-3,
                  ref_meg=1e-12, misc=1e-3, stim=1, resp=1, chpi=1e-4)`
        remove_dc : bool
            If True remove DC component when plotting data.
        order : 'type' | 'original' | array
            Order in which to plot data. 'type' groups by channel type,
            'original' plots in the order of ch_names, array gives the
            indices to use in plotting.
        show_options : bool
            If True, a dialog for options related to projection is shown.
        title : str | None
            The title of the window. If None, and either the filename of the
            raw object or '<unknown>' will be displayed as title.
        show : bool
            Show figures if True
        block : bool
            Whether to halt program execution until the figure is closed.
            Useful for setting bad channels on the fly (click on line).

        Returns
        -------
        fig : Instance of matplotlib.figure.Figure
            Raw traces.

        Notes
        -----
        The arrow keys (up/down/left/right) can typically be used to navigate
        between channels and time ranges, but this depends on the backend
        matplotlib is configured to use (e.g., mpl.use('TkAgg') should work).
        To mark or un-mark a channel as bad, click on the rather flat segments
        of a channel's time series. The changes will be reflected immediately
        in the raw object's ``raw.info['bads']`` entry.
        """
        return plot_raw(raw, events, duration, start, n_channels, bgcolor,
                        color, bad_color, event_color, scalings, remove_dc,
                        order, show_options, title, show, block)

    @verbose
    def plot_psds(self, tmin=0.0, tmax=60.0, fmin=0, fmax=np.inf,
                  proj=False, n_fft=2048, picks=None, ax=None, color='black',
                  area_mode='std', area_alpha=0.33, n_jobs=1, verbose=None):
        """Plot the power spectral density across channels

        Parameters
        ----------
        tmin : float
            Start time for calculations.
        tmax : float
            End time for calculations.
        fmin : float
            Start frequency to consider.
        fmax : float
            End frequency to consider.
        proj : bool
            Apply projection.
        n_fft : int
            Number of points to use in Welch FFT calculations.
        picks : list | None
            List of channels to use. Cannot be None if `ax` is supplied. If
            both `picks` and `ax` are None, separate subplots will be created
            for each standard channel type (`mag`, `grad`, and `eeg`).
        ax : instance of matplotlib Axes | None
            Axes to plot into. If None, axes will be created.
        color : str | tuple
            A matplotlib-compatible color to use.
        area_mode : str | None
            How to plot area. If 'std', the mean +/- 1 STD (across channels)
            will be plotted. If 'range', the min and max (across channels)
            will be plotted. Bad channels will be excluded from these
            calculations. If None, no area will be plotted.
        area_alpha : float
            Alpha for the area.
        n_jobs : int
            Number of jobs to run in parallel.
        verbose : bool, str, int, or None
            If not None, override default verbose level (see mne.verbose).
        """
        return plot_raw_psds(self, tmin, tmax, fmin, fmax, proj, n_fft, picks,
                             ax, color, area_mode, area_alpha, n_jobs)

    def time_as_index(self, times, use_first_samp=False):
        """Convert time to indices

        Parameters
        ----------
        times : list-like | float | int
            List of numbers or a number representing points in time.
        use_first_samp : boolean
            If True, time is treated as relative to the session onset, else
            as relative to the recording onset.

        Returns
        -------
        index : ndarray
            Indices corresponding to the times supplied.
        """
        return _time_as_index(times, self.info['sfreq'], self.first_samp,
                              use_first_samp)

    def index_as_time(self, index, use_first_samp=False):
        """Convert time to indices

        Parameters
        ----------
        index : list-like | int
            List of ints or int representing points in time.
        use_first_samp : boolean
            If True, the time returned is relative to the session onset, else
            relative to the recording onset.

        Returns
        -------
        times : ndarray
            Times corresponding to the index supplied.
        """
        return _index_as_time(index, self.info['sfreq'], self.first_samp,
                              use_first_samp)

    def estimate_rank(self, tstart=0.0, tstop=30.0, tol=1e-4,
                      return_singular=False):
        """Estimate rank of the raw data

        This function is meant to provide a reasonable estimate of the rank.
        The true rank of the data depends on many factors, so use at your
        own risk.

        Parameters
        ----------
        tstart : float
            Start time to use for rank estimation. Defaul is 0.0.
        tstop : float | None
            End time to use for rank estimation. Default is 30.0.
            If None, the end time of the raw file is used.
        tol : float
            Tolerance for singular values to consider non-zero in
            calculating the rank. The singular values are calculated
            in this method such that independent data are expected to
            have singular value around one.
        return_singular : bool
            If True, also return the singular values that were used
            to determine the rank.

        Returns
        -------
        rank : int
            Estimated rank of the data.
        s : array
            If return_singular is True, the singular values that were
            thresholded to determine the rank are also returned.

        Notes
        -----
        If data are not pre-loaded, the appropriate data will be loaded
        by this function (can be memory intensive).

        Projectors are not taken into account unless they have been applied
        to the data using apply_proj(), since it is not always possible
        to tell whether or not projectors have been applied previously.

        Bad channels will be excluded from calculations.
        """
        start = max(0, self.time_as_index(tstart)[0])
        if tstop is None:
            stop = self.n_times - 1
        else:
            stop = min(self.n_times - 1, self.time_as_index(tstop)[0])
        tslice = slice(start, stop + 1)
        picks = pick_types(self.info, meg=True, eeg=True, ref_meg=False,
                           exclude='bads')
        # ensure we don't get a view of data
        if len(picks) == 1:
            return 1.0, 1.0
        # this should already be a copy, so we can overwrite it
        data = self[picks, tslice][0]
        return estimate_rank(data, tol, return_singular, copy=False)

    @property
    def ch_names(self):
        return self.info['ch_names']

    @property
    def n_times(self):
        return self.last_samp - self.first_samp + 1

    def __len__(self):
        return self.n_times

    def load_bad_channels(self, bad_file=None, force=False):
        """
        Mark channels as bad from a text file, in the style
        (mostly) of the C function mne_mark_bad_channels

        Parameters
        ----------
        bad_file : string
            File name of the text file containing bad channels
            If bad_file = None, bad channels are cleared, but this
            is more easily done directly as raw.info['bads'] = [].

        force : boolean
            Whether or not to force bad channel marking (of those
            that exist) if channels are not found, instead of
            raising an error.
        """

        if bad_file is not None:
            # Check to make sure bad channels are there
            names = frozenset(self.info['ch_names'])
            bad_names = [l for l in open(bad_file).read().splitlines() if l]
            names_there = [ci for ci in bad_names if ci in names]
            count_diff = len(bad_names) - len(names_there)

            if count_diff > 0:
                if not force:
                    raise ValueError('Bad channels from:\n%s\n not found '
                                     'in:\n%s' % (bad_file,
                                                  self.info['filenames'][0]))
                else:
                    warnings.warn('%d bad channels from:\n%s\nnot found '
                                  'in:\n%s' % (count_diff, bad_file,
                                               self.info['filenames'][0]))
            self.info['bads'] = names_there
        else:
            self.info['bads'] = []

    def append(self, raws, preload=None):
        """Concatenate raw instances as if they were continuous

        Parameters
        ----------
        raws : list, or Raw instance
            list of Raw instances to concatenate to the current instance
            (in order), or a single raw instance to concatenate.

        preload : bool, str, or None (default None)
            Preload data into memory for data manipulation and faster indexing.
            If True, the data will be preloaded into memory (fast, requires
            large amount of memory). If preload is a string, preload is the
            file name of a memory-mapped file which is used to store the data
            on the hard drive (slower, requires less memory). If preload is
            None, preload=True or False is inferred using the preload status
            of the raw files passed in.
        """
        if not isinstance(raws, list):
            raws = [raws]

        # make sure the raws are compatible
        all_raws = [self]
        all_raws += raws
        _check_raw_compatibility(all_raws)

        # deal with preloading data first (while files are separate)
        all_preloaded = self._preloaded and all(r._preloaded for r in raws)
        if preload is None:
            if all_preloaded:
                preload = True
            else:
                preload = False

        if preload is False:
            if self._preloaded:
                self._data = None
                self._times = None
            self._preloaded = False
        else:
            # do the concatenation ourselves since preload might be a string
            nchan = self.info['nchan']
            c_ns = np.cumsum([rr.n_times for rr in ([self] + raws)])
            nsamp = c_ns[-1]

            if not self._preloaded:
                this_data = self._read_segment()[0]
            else:
                this_data = self._data

            # allocate the buffer
            if isinstance(preload, string_types):
                _data = np.memmap(preload, mode='w+', dtype=this_data.dtype,
                                  shape=(nchan, nsamp))
            else:
                _data = np.empty((nchan, nsamp), dtype=this_data.dtype)

            _data[:, 0:c_ns[0]] = this_data

            for ri in range(len(raws)):
                if not raws[ri]._preloaded:
                    # read the data directly into the buffer
                    data_buffer = _data[:, c_ns[ri]:c_ns[ri + 1]]
                    raws[ri]._read_segment(data_buffer=data_buffer)
                else:
                    _data[:, c_ns[ri]:c_ns[ri + 1]] = raws[ri]._data
            self._data = _data
            self._preloaded = True

        # now combine information from each raw file to construct new self
        for r in raws:
            self._first_samps = np.r_[self._first_samps, r._first_samps]
            self._last_samps = np.r_[self._last_samps, r._last_samps]
            self._raw_lengths = np.r_[self._raw_lengths, r._raw_lengths]
            self.rawdirs += r.rawdirs
            self.info['filenames'] += r.info['filenames']
        # reconstruct fids in case some were preloaded and others weren't
        self._initialize_fids()
        self.last_samp = self.first_samp + sum(self._raw_lengths) - 1

        # this has to be done after first and last sample are set appropriately
        if self._preloaded:
            self._times = np.arange(self.n_times) / self.info['sfreq']

    def close(self):
        """Close the files on disk."""
        [f.close() for f in self.fids]
        self.fids = []

    def copy(self):
        """ Return copy of Raw instance
        """
        old_fids = self.fids
        try:
            self.fids = list()
            new = deepcopy(self)
            new._initialize_fids()
        finally:
            self.fids = old_fids
        return new

    def _initialize_fids(self):
        """Initialize self.fids based on self.info['filenames']
        """
        if not self._preloaded:
            self.fids = [open(fname, "rb") for fname in self.info['filenames']]
            for fid in self.fids:
                fid.seek(0, 0)
        else:
            self.fids = []

    def as_data_frame(self, picks=None, start=None, stop=None, scale_time=1e3,
                      scalings=None, use_time_index=True, copy=True):
        """Get the epochs as Pandas DataFrame

        Export raw data in tabular structure with MEG channels.

        Caveat! To save memory, depending on selected data size consider
        setting copy to False.

        Parameters
        ----------
        picks : None | array of int
            If None only MEG and EEG channels are kept
            otherwise the channels indices in picks are kept.
        start : int | None
            Data-extraction start index. If None, data will be exported from
            the first sample.
        stop : int | None
            Data-extraction stop index. If None, data will be exported to the
            last index.
        scale_time : float
            Scaling to be applied to time units.
        scalings : dict | None
            Scaling to be applied to the channels picked. If None, defaults to
            ``scalings=dict(eeg=1e6, grad=1e13, mag=1e15, misc=1.0)`.
        use_time_index : bool
            If False, times will be included as in the data table, else it will
            be used as index object.
        copy : bool
            If true, data will be copied. Else data may be modified in place.

        Returns
        -------
        df : instance of pandas.core.DataFrame
            Raw data exported into tabular data structure.
        """

        pd = _check_pandas_installed()
        if picks is None:
            picks = list(range(self.info['nchan']))

        data, times = self[picks, start:stop]

        if copy:
            data = data.copy()

        types = [channel_type(self.info, idx) for idx in picks]
        n_channel_types = 0
        ch_types_used = []

        scalings = _mutable_defaults(('scalings', scalings))[0]
        for t in scalings.keys():
            if t in types:
                n_channel_types += 1
                ch_types_used.append(t)

        for t in ch_types_used:
            scaling = scalings[t]
            idx = [picks[i] for i in range(len(picks)) if types[i] == t]
            if len(idx) > 0:
                data[idx] *= scaling

        assert times.shape[0] == data.shape[1]
        col_names = [self.ch_names[k] for k in picks]

        df = pd.DataFrame(data.T, columns=col_names)
        df.insert(0, 'time', times * scale_time)

        if use_time_index is True:
            if 'time' in df:
                df['time'] = df['time'].astype(np.int64)
            with warnings.catch_warnings(record=True):
                df.set_index('time', inplace=True)

        return df

    def to_nitime(self, picks=None, start=None, stop=None,
                  use_first_samp=False, copy=True):
        """ Raw data as nitime TimeSeries

        Parameters
        ----------
        picks : array-like | None
            Indices of channels to apply. If None, all channels will be
            exported.
        start : int | None
            Data-extraction start index. If None, data will be exported from
            the first sample.
        stop : int | None
            Data-extraction stop index. If None, data will be exported to the
            last index.
        use_first_samp: bool
            If True, the time returned is relative to the session onset, else
            relative to the recording onset.
        copy : bool
            Whether to copy the raw data or not.

        Returns
        -------
        raw_ts : instance of nitime.TimeSeries
        """
        try:
            from nitime import TimeSeries  # to avoid strong dependency
        except ImportError:
            raise Exception('the nitime package is missing')

        data, _ = self[picks, start:stop]
        if copy:
            data = data.copy()

        start_time = self.index_as_time(start if start else 0, use_first_samp)
        raw_ts = TimeSeries(data, sampling_rate=self.info['sfreq'],
                            t0=start_time)

        raw_ts.ch_names = [self.ch_names[k] for k in picks]

        return raw_ts

    @verbose
    def _read_segment(self, start=0, stop=None, sel=None, data_buffer=None,
                      verbose=None, projector=None):
        """Read a chunk of raw data

        Parameters
        ----------
        start : int, (optional)
            first sample to include (first is 0). If omitted, defaults to the
            first sample in data.
        stop : int, (optional)
            First sample to not include.
            If omitted, data is included to the end.
        sel : array, optional
            Indices of channels to select.
        data_buffer : array or str, optional
            numpy array to fill with data read, must have the correct shape.
            If str, a np.memmap with the correct data type will be used
            to store the data.
        verbose : bool, str, int, or None
            If not None, override default verbose level (see mne.verbose).
        projector : array
            SSP operator to apply to the data.

        Returns
        -------
        data : array, [channels x samples]
           the data matrix (channels x samples).
        times : array, [samples]
            returns the time values corresponding to the samples.
        """
        #  Initial checks
        start = int(start)
        stop = self.n_times if stop is None else min([int(stop), self.n_times])

        if start >= stop:
            raise ValueError('No data in this range')

        logger.info('Reading %d ... %d  =  %9.3f ... %9.3f secs...' %
                    (start, stop - 1, start / float(self.info['sfreq']),
                     (stop - 1) / float(self.info['sfreq'])))

        #  Initialize the data and calibration vector
        nchan = self.info['nchan']

        n_sel_channels = nchan if sel is None else len(sel)
        # convert sel to a slice if possible for efficiency
        if sel is not None and len(sel) > 1 and np.all(np.diff(sel) == 1):
            sel = slice(sel[0], sel[-1] + 1)
        idx = slice(None, None, None) if sel is None else sel
        data_shape = (n_sel_channels, stop - start)
        if isinstance(data_buffer, np.ndarray):
            if data_buffer.shape != data_shape:
                raise ValueError('data_buffer has incorrect shape')
            data = data_buffer
        else:
            data = None  # we will allocate it later, once we know the type

        mult = list()
        for ri in range(len(self._raw_lengths)):
            mult.append(np.diag(self.cals.ravel()))
            if self.comp is not None:
                mult[ri] = np.dot(self.comp, mult[ri])
            if projector is not None:
                mult[ri] = np.dot(projector, mult[ri])
            mult[ri] = mult[ri][idx]

        # deal with having multiple files accessed by the raw object
        cumul_lens = np.concatenate(([0], np.array(self._raw_lengths,
                                                   dtype='int')))
        cumul_lens = np.cumsum(cumul_lens)
        files_used = np.logical_and(np.less(start, cumul_lens[1:]),
                                    np.greater_equal(stop - 1,
                                                     cumul_lens[:-1]))

        first_file_used = False
        s_off = 0
        dest = 0
        if isinstance(idx, slice):
            cals = self.cals.ravel()[idx][:, np.newaxis]
        else:
            cals = self.cals.ravel()[:, np.newaxis]

        for fi in np.nonzero(files_used)[0]:
            start_loc = self._first_samps[fi]
            # first iteration (only) could start in the middle somewhere
            if not first_file_used:
                first_file_used = True
                start_loc += start - cumul_lens[fi]
            stop_loc = np.min([stop - 1 - cumul_lens[fi] +
                               self._first_samps[fi], self._last_samps[fi]])
            if start_loc < self._first_samps[fi]:
                raise ValueError('Bad array indexing, could be a bug')
            if stop_loc > self._last_samps[fi]:
                raise ValueError('Bad array indexing, could be a bug')
            if stop_loc < start_loc:
                raise ValueError('Bad array indexing, could be a bug')
            len_loc = stop_loc - start_loc + 1

            for this in self.rawdirs[fi]:

                #  Do we need this buffer
                if this['last'] >= start_loc:
                    #  The picking logic is a bit complicated
                    if stop_loc > this['last'] and start_loc < this['first']:
                        #    We need the whole buffer
                        first_pick = 0
                        last_pick = this['nsamp']
                        logger.debug('W')

                    elif start_loc >= this['first']:
                        first_pick = start_loc - this['first']
                        if stop_loc <= this['last']:
                            #   Something from the middle
                            last_pick = this['nsamp'] + stop_loc - this['last']
                            logger.debug('M')
                        else:
                            #   From the middle to the end
                            last_pick = this['nsamp']
                            logger.debug('E')
                    else:
                        #    From the beginning to the middle
                        first_pick = 0
                        last_pick = stop_loc - this['first'] + 1
                        logger.debug('B')

                    #   Now we are ready to pick
                    picksamp = last_pick - first_pick
                    if picksamp > 0:
                        # only read data if it exists
                        if this['ent'] is not None:
                            one = read_tag(self.fids[fi], this['ent'].pos,
                                           shape=(this['nsamp'], nchan),
                                           rlims=(first_pick, last_pick)).data
                            if np.isrealobj(one):
                                dtype = np.float
                            else:
                                dtype = np.complex128
                            one.shape = (picksamp, nchan)
                            one = one.T.astype(dtype)
                            # use proj + cal factors in mult
                            if mult is not None:
                                one[idx] = np.dot(mult[fi], one)
                            else:  # apply just the calibration factors
                                # this logic is designed to limit memory copies
                                if isinstance(idx, slice):
                                    # This is a view operation, so it's fast
                                    one[idx] *= cals
                                else:
                                    # Extra operations are actually faster here
                                    # than creating a new array
                                    # (fancy indexing)
                                    one *= cals

                            # if not already done, allocate array with
                            # right type
                            data = _allocate_data(data, data_buffer,
                                                  data_shape, dtype)
                            if isinstance(idx, slice):
                                # faster to slice in data than doing
                                # one = one[idx] sooner
                                data[:, dest:(dest + picksamp)] = one[idx]
                            else:
                                # faster than doing one = one[idx]
                                data_view = data[:, dest:(dest + picksamp)]
                                for ii, ix in enumerate(idx):
                                    data_view[ii] = one[ix]
                        dest += picksamp

                #   Done?
                if this['last'] >= stop_loc:
                    # if not already done, allocate array with float dtype
                    data = _allocate_data(data, data_buffer, data_shape,
                                          np.float)
                    break

            self.fids[fi].seek(0, 0)  # Go back to beginning of the file
            s_off += len_loc
            # double-check our math
            if not s_off == dest:
                raise ValueError('Incorrect file reading')

        logger.info('[done]')
        times = np.arange(start, stop) / self.info['sfreq']

        return data, times

    def __repr__(self):
        s = "n_channels x n_times : %s x %s" % (len(self.info['ch_names']),
                                                self.n_times)
        return "<Raw  |  %s>" % s

def set_eeg_reference(raw, ref_channels, copy=True):
    """Rereference eeg channels to new reference channel(s).

    If multiple reference channels are specified, they will be averaged.

    Parameters
    ----------
    raw : instance of Raw
        Instance of Raw with eeg channels and reference channel(s).

    ref_channels : list of str
        The name(s) of the reference channel(s).

    copy : bool
        Specifies whether instance of Raw will be copied or modified in place.

    Returns
    -------
    raw : instance of Raw
        Instance of Raw with eeg channels rereferenced.

    ref_data : array
        Array of reference data subtracted from eeg channels.
    """
    # Check to see that raw data is preloaded
    if not raw._preloaded:
        raise RuntimeError('Raw data needs to be preloaded. Use '
                           'preload=True (or string) in the constructor.')
    # Make sure that reference channels are loaded as list of string
    if not isinstance(ref_channels, list):
        raise IOError('Reference channel(s) must be a list of string. '
                      'If using a single reference channel, enter as '
                      'a list with one element.')
    # Find the indices to the reference electrodes
    ref_idx = [raw.ch_names.index(c) for c in ref_channels]

    # Get the reference array
    ref_data = raw._data[ref_idx].mean(0)

    # Get the indices to the eeg channels using the pick_types function
    eeg_idx = pick_types(raw.info, exclude="bads", eeg=True, meg=False,
                         ref_meg=False)

    # Copy raw data or modify raw data in place
    if copy:  # copy data
        raw = raw.copy()

    # Rereference the eeg channels
    raw._data[eeg_idx] -= ref_data

    # Return rereferenced data and reference array
    return raw, ref_data

def _allocate_data(data, data_buffer, data_shape, dtype):
    if data is None:
        # if not already done, allocate array with right type
        if isinstance(data_buffer, string_types):
            # use a memmap
            data = np.memmap(data_buffer, mode='w+',
                             dtype=dtype, shape=data_shape)
        else:
            data = np.zeros(data_shape, dtype=dtype)
    return data

def _time_as_index(times, sfreq, first_samp=0, use_first_samp=False):
    """Convert time to indices

    Parameters
    ----------
    times : list-like | float | int
        List of numbers or a number representing points in time.
    use_first_samp : boolean
        If True, time is treated as relative to the session onset, else
        as relative to the recording onset.

    Returns
    -------
    index : ndarray
        Indices corresponding to the times supplied.
    """
    index = np.atleast_1d(times) * sfreq
    index -= (first_samp if use_first_samp else 0)
    return index.astype(int)

def _index_as_time(index, sfreq, first_samp=0, use_first_samp=False):
    """Convert time to indices

    Parameters
    ----------
    index : list-like | int
        List of ints or int representing points in time.
    use_first_samp : boolean
        If True, the time returned is relative to the session onset, else
        relative to the recording onset.

    Returns
    -------
    times : ndarray
        Times corresponding to the index supplied.
    """
    times = np.atleast_1d(index) + (first_samp if use_first_samp else 0)
    return times / sfreq

class _RawShell():
    """Used for creating a temporary raw object"""
    def __init__(self):
        self.first_samp = None
        self.last_samp = None
        self.cals = None
        self.rawdir = None
        self._projector = None

    @property
    def n_times(self):
        return self.last_samp - self.first_samp + 1

###############################################################################
# Writing

from .write import (start_file, end_file, start_block, end_block,
                    write_dau_pack16, write_float, write_double,
                    write_complex64, write_complex128, write_int, write_id)

def start_writing_raw(name, info, sel=None, data_type=FIFF.FIFFT_FLOAT,
                      reset_range=True):
    """Start write raw data in file

    Data will be written in float

    Parameters
    ----------
    name : string
        Name of the file to create.
    info : dict
        Measurement info.
    sel : array of int, optional
        Indices of channels to include. By default all channels are included.
    data_type : int
        The data_type in case it is necessary. Should be 4 (FIFFT_FLOAT),
        5 (FIFFT_DOUBLE), 16 (FIFFT_DAU_PACK16), or 3 (FIFFT_INT) for raw data.
    reset_range : bool
        If True, the info['chs'][k]['range'] parameter will be set to unity.

    Returns
    -------
    fid : file
        The file descriptor.
    cals : list
        calibration factors.
    """
    #
    #  Create the file and save the essentials
    #
    fid = start_file(name)
    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'])
    #
    #    Measurement info
    #
    info = copy.deepcopy(info)
    if sel is not None:
        info['chs'] = [info['chs'][k] for k in sel]
        info['nchan'] = len(sel)

        ch_names = [c['ch_name'] for c in info['chs']]  # name of good channels
        comps = copy.deepcopy(info['comps'])
        for c in comps:
            row_idx = [k for k, n in enumerate(c['data']['row_names'])
                       if n in ch_names]
            row_names = [c['data']['row_names'][i] for i in row_idx]
            rowcals = c['rowcals'][row_idx]
            c['rowcals'] = rowcals
            c['data']['nrow'] = len(row_names)
            c['data']['row_names'] = row_names
            c['data']['data'] = c['data']['data'][row_idx]
        info['comps'] = comps

    cals = []
    for k in range(info['nchan']):
        #
        #   Scan numbers may have been messed up
        #
        info['chs'][k]['scanno'] = k + 1  # scanno starts at 1 in FIF format
        if reset_range is True:
            info['chs'][k]['range'] = 1.0
        cals.append(info['chs'][k]['cal'] * info['chs'][k]['range'])

    write_meas_info(fid, info, data_type=data_type, reset_range=reset_range)

    #
    # Start the raw data
    #
    start_block(fid, FIFF.FIFFB_RAW_DATA)

    return fid, cals

def write_raw_buffer(fid, buf, cals, format, inv_comp):
    """Write raw buffer

    Parameters
    ----------
    fid : file descriptor
        an open raw data file.
    buf : array
        The buffer to write.
    cals : array
        Calibration factors.
    format : str
        'short', 'int', 'single', or 'double' for 16/32 bit int or 32/64 bit
        float for each item. This will be doubled for complex datatypes. Note
        that short and int formats cannot be used for complex data.
    inv_comp : array | None
        The CTF compensation matrix used to revert compensation
        change when reading.
    """
    if buf.shape[0] != len(cals):
        raise ValueError('buffer and calibration sizes do not match')

    if not format in ['short', 'int', 'single', 'double']:
        raise ValueError('format must be "short", "single", or "double"')

    if np.isrealobj(buf):
        if format == 'short':
            write_function = write_dau_pack16
        elif format == 'int':
            write_function = write_int
        elif format == 'single':
            write_function = write_float
        else:
            write_function = write_double
    else:
        if format == 'single':
            write_function = write_complex64
        elif format == 'double':
            write_function = write_complex128
        else:
            raise ValueError('only "single" and "double" supported for '
                             'writing complex data')

    if inv_comp is not None:
        buf = np.dot(inv_comp / np.ravel(cals)[:, None], buf)
    else:
        buf = buf / np.ravel(cals)[:, None]

    write_function(fid, FIFF.FIFF_DATA_BUFFER, buf)

    # make sure we didn't go over the 2GB file size limit
    pos = fid.tell()
    if pos >= 2147483647:  # np.iinfo(np.int32).max
        raise IOError('2GB file size limit reached. Support for larger '
                      'raw files will be added in the future.')

def finish_writing_raw(fid):
    """Finish writing raw FIF file

    Parameters
    ----------
    fid : file descriptor
        an open raw data file.
    """
    end_block(fid, FIFF.FIFFB_RAW_DATA)
    end_block(fid, FIFF.FIFFB_MEAS)
    end_file(fid)

def _envelope(x):
    """ Compute envelope signal """
    return np.abs(hilbert(x))

def _check_raw_compatibility(raw):
    """Check to make sure all instances of Raw
    in the input list raw have compatible parameters"""
    for ri in range(1, len(raw)):
        if not raw[ri].info['nchan'] == raw[0].info['nchan']:
            raise ValueError('raw[%d][\'info\'][\'nchan\'] must match' % ri)
        if not raw[ri].info['bads'] == raw[0].info['bads']:
            raise ValueError('raw[%d][\'info\'][\'bads\'] must match' % ri)
        if not raw[ri].info['sfreq'] == raw[0].info['sfreq']:
            raise ValueError('raw[%d][\'info\'][\'sfreq\'] must match' % ri)
        if not set(raw[ri].info['ch_names']) == set(raw[0].info['ch_names']):
            raise ValueError('raw[%d][\'info\'][\'ch_names\'] must match' % ri)
        if not all(raw[ri].cals == raw[0].cals):
            raise ValueError('raw[%d].cals must match' % ri)
        if len(raw[0].info['projs']) != len(raw[ri].info['projs']):
            raise ValueError('SSP projectors in raw files must be the same')
        if not all(proj_equal(p1, p2) for p1, p2 in
                   zip(raw[0].info['projs'], raw[ri].info['projs'])):
            raise ValueError('SSP projectors in raw files must be the same')
    if not all([r.orig_format == raw[0].orig_format for r in raw]):
        warnings.warn('raw files do not all have the same data format, '
                      'could result in precision mismatch. Setting '
                      'raw.orig_format="unknown"')
        raw[0].orig_format = 'unknown'

def concatenate_raws(raws, preload=None):
    """Concatenate raw instances as if they were continuous. Note that raws[0]
    is modified in-place to achieve the concatenation.

    Parameters
    ----------
    raws : list
        list of Raw instances to concatenate (in order).

    preload : bool, or None
        If None, preload status is inferred using the preload status of the
        raw files passed in. True or False sets the resulting raw file to
        have or not have data preloaded.

    Returns
    -------
    raw : instance of Raw
        The result of the concatenation (first Raw instance passed in).
    """
    raws[0].append(raws[1:], preload)
    return raws[0]

def get_chpi_positions(raw, t_step=None):
    """Extract head positions

    Note that the raw instance must have CHPI channels recorded.

    Parameters
    ----------
    raw : instance of Raw | str
        Raw instance to extract the head positions from. Can also be a
        path to a Maxfilter log file (str).
    t_step : float | None
        Sampling interval to use when converting data. If None, it will
        be automatically determined. By default, a sampling interval of
        1 second is used if processing a raw data. If processing a
        Maxfilter log file, this must be None because the log file
        itself will determine the sampling interval.

    Returns
    -------
    translation : array
        A 2-dimensional array of head position vectors (n_time x 3).
    rotation : array
        A 3-dimensional array of rotation matrices (n_time x 3 x 3).
    t : array
        The time points associated with each position (n_time).

    Notes
    -----
    The digitized HPI head frame y is related to the frame position X as:

        Y = np.dot(rotation, X) + translation

    Note that if a Maxfilter log file is being processed, the start time
    may not use the same reference point as the rest of mne-python (i.e.,
    it could be referenced relative to raw.first_samp or something else).
    """
    if isinstance(raw, Raw):
        # for simplicity, we'll sample at 1 sec intervals like maxfilter
        if t_step is None:
            t_step = 1.0
        if not np.isscalar(t_step):
            raise TypeError('t_step must be a scalar or None')
        picks = pick_types(raw.info, meg=False, ref_meg=False,
                           chpi=True, exclude=[])
        if len(picks) == 0:
            raise RuntimeError('raw file has no CHPI channels')
        time_idx = raw.time_as_index(np.arange(0, raw.n_times
                                               / raw.info['sfreq'], t_step))
        data = [raw[picks, ti] for ti in time_idx]
        t = np.array([d[1] for d in data])
        data = np.array([d[0][:, 0] for d in data])
        data = np.c_[t, data]
    else:
        if not isinstance(raw, string_types):
            raise TypeError('raw must be an instance of fiff.Raw or string')
        if not op.isfile(raw):
            raise IOError('File "%s" does not exist' % raw)
        if t_step is not None:
            raise ValueError('t_step must be None if processing a log')
        data = np.loadtxt(raw, skiprows=1)  # first line is header, skip it
    t = data[:, 0]
    translation = data[:, 4:7].copy()
    rotation = _quart_to_rot(data[:, 1:4])
    return translation, rotation, t

def _quart_to_rot(q):
    """Helper to convert quarternions to rotations"""
    q0 = np.sqrt(1 - np.sum(q[:, 0:3] ** 2, 1))
    q1 = q[:, 0]
    q2 = q[:, 1]
    q3 = q[:, 2]
    rotation = np.array((np.c_[(q0 ** 2 + q1 ** 2 - q2 ** 2 - q3 ** 2,
                               2 * (q1 * q2 - q0 * q3),
                               2 * (q1 * q3 + q0 * q2))],
                         np.c_[(2 * (q1 * q2 + q0 * q3),
                                q0 ** 2 + q2 ** 2 - q1 ** 2 - q3 ** 2,
                                2 * (q2 * q3 - q0 * q1))],
                         np.c_[(2 * (q1 * q3 - q0 * q2),
                                2 * (q2 * q3 + q0 * q1),
                                q0 ** 2 + q3 ** 2 - q1 ** 2 - q2 ** 2)]
                         ))
    rotation = np.swapaxes(rotation, 0, 1).copy()
    return rotation


from __future__ import print_function
import sys
import os
import argparse
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torchvision.transforms as transforms
from torch.autograd import Variable
from data import WIDERFace_ROOT , WIDERFace_CLASSES as labelmap
from PIL import Image
from data import WIDERFaceDetection, WIDERFaceAnnotationTransform, WIDERFace_CLASSES, WIDERFace_ROOT, BaseTransform , TestBaseTransform
from data import *
import torch.utils.data as data
from light_face_ssd import build_ssd
import pdb
import numpy as np
import cv2
import math
import matplotlib.pyplot as plt
import time
from layers.box_utils import nms
import torchvision
#from torchscope import scope

plt.switch_backend('agg')

parser = argparse.ArgumentParser(description='Single Shot MultiBox Detection')
parser.add_argument('--trained_model', default='weights/light_DSFD.pth',
                    type=str, help='Trained state_dict file path to open')
#parser.add_argument('--save_folder', default='eval_tools/WIDERFace_SSD_F_RES152_FPN_MIO_CPM_PA0.70.3_DAS_Anchor1.5_IOU0.4_DeepClsRelu_6W_16bs_pyramidtest_testset/', type=str,
#                    help='Dir to save results')
parser.add_argument('--visual_threshold', default=0.01, type=float,
                    help='Final confidence threshold')
parser.add_argument('--cuda', default=False, type=bool,
                    help='Use cuda to train model')
parser.add_argument('--widerface_root', default=WIDERFace_ROOT, help='Location of VOC root directory')
args = parser.parse_args()

if args.cuda and torch.cuda.is_available():
    torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
    torch.set_default_tensor_type('torch.FloatTensor')
#if not os.path.exists(args.save_folder):
#    os.mkdir(args.save_folder)

def detect_face(image, shrink):
    x = image
    if shrink != 1:
        x = cv2.resize(image, None, None, fx=shrink, fy=shrink, interpolation=cv2.INTER_LINEAR)
    print('shrink:{}'.format(shrink))
    width = x.shape[1]
    height = x.shape[0]
    print (width , height)
    x = x.astype(np.float32)
    x -= np.array([104, 117, 123],dtype=np.float32)

    x = torch.from_numpy(x).permute(2, 0, 1)
    x = x.unsqueeze(0)
    if args.cuda:
        x = Variable(x.cuda(), volatile=True)
    else:
        #with torch.no_grad():
        #    x = Variable(x)
        x = Variable(x, volatile=True)

    #net.priorbox = PriorBoxLayer(width,height)
    #model = torchvision.models.resnet18(pretrained=True)
    time0 = time.time()
    #model(x)
    #time01 = time.time()
    #print ( time01 - time0)
    y = net(x)
    time1 = time.time()
    print ( "infer time: %.4f " % (time1-time0) )
    detections = y.data
    scale = torch.Tensor([width, height, width, height])

    boxes=[]
    scores = []
    for i in range(detections.size(1)):
        j = 0
        while detections[0,i,j,0] >= 0.01:
            score = detections[0,i,j,0]
            pt = (detections[0, i, j, 1:]*scale).cpu().numpy()
            boxes.append([pt[0],pt[1],pt[2],pt[3]])
            scores.append(score)
            j += 1
            if j >= detections.size(2):
                break

    det_conf = np.array(scores)
    boxes = np.array(boxes)

    if boxes.shape[0] == 0:
        return np.array([[0,0,0,0,0.001]])

    det_xmin = boxes[:,0] / shrink
    det_ymin = boxes[:,1] / shrink
    det_xmax = boxes[:,2] / shrink
    det_ymax = boxes[:,3] / shrink
    det = np.column_stack((det_xmin, det_ymin, det_xmax, det_ymax, det_conf))

    keep_index = np.where(det[:, 4] >= 0)[0]
    det = det[keep_index, :]
    return det

# load net
cfg = widerface_640
num_classes = len(WIDERFace_CLASSES) + 1 # +1 background
net = build_ssd('test', cfg['min_dim'], num_classes) # initialize SSD
#net = nn.DataParallel(net)
net.load_state_dict(torch.load(args.trained_model))
if args.cuda:
    net.cuda()
net.eval()
print('Finished loading model!')

#scope( net, input_size=( 3, 720, 1280) , batch_size=1 , device='cpu')

path = './data/yuebing.jpg'
image = cv2.imread(path, cv2.IMREAD_COLOR)
image = cv2.resize(image, (720, 1280), interpolation=cv2.INTER_LINEAR)
#(720,1280)
shrink=1
time1 = time.time()
for i in range(100):
  det = detect_face(image, shrink)  # origin test
time2 = time.time()
print( (time2-time1)/100 )


import sys
sys.path.append("X:\\")
import labscript_utils.h5_lock
from labscript_utils.ls_zprocess import ProcessTree
import zprocess

import pyvisa
import numpy as np
import time
from device_server import DeviceServer
import h5py

process_tree = ProcessTree.instance()
process_tree.zlock_client.set_process_name("rigol-scope")

class ScopeInterface:
    def __init__(self, identifying_string):
        self.resource_manager = pyvisa.ResourceManager()
        all_devices = self.resource_manager.list_resources()
        acceptable_devices = [
            i for i in all_devices if identifying_string in i]
        if acceptable_devices == []:
            raise Exception(
                f"Could not find device with identifying string {identifying_string}"
                + "Found devices:\n" + "\n".join(all_devices)
            )
        self.scope = self.resource_manager.open_resource(acceptable_devices[0])
        self.n_channels = self.scope.query_ascii_values(":SYST:RAM?")[0]
        print(f'Connected to scope with {self.n_channels} channels')

    def get_voltage(self, channel):
        """
        Inputs:
            Channel - integer

        Outputs:
            time_relative_trigger - numpy array of time (in seconds) relative to
            the trigger pulse
            data - numpy array of voltage at each time
        """
        if channel > self.n_channels or channel < 1:
            raise Exception(
                f"Channel must be between 1 and {self.n_channels}.")
        self.scope.write(":STOP")
        self.scope.write(f':WAV:SOUR CHAN{channel}')
        self.scope.write(':WAV:MODE norm')
#        self.scope.write(":TRIG:EDGE:SWE NORM")
        self.scope.write(':WAV:FORM byte')
        self.scope.write(":WAV:DATA?")
        data = self.scope.read_raw()
        _, _, wvf_points, _, xinc, xorigin, xref, yinc, yorigin, yref = self.scope.query_ascii_values(
            ":WAV:PRE?")
        data = np.frombuffer(data, dtype="uint8", offset=12)
        data = (data - yref - yorigin) * yinc
        data = data[:-1]
        #time_relative_trigger = np.arange(0, -wvf_points * xinc, -xinc) - xorigin
        time_relative_trigger = np.linspace(
            0, -(wvf_points - 1) * xinc, int(wvf_points)) - xorigin
        return time_relative_trigger, data

    def run(self):
        self.scope.write(":RUN")

    def get_all_voltages(self, channels):
        """
        Get the voltage trace on the screen of every channel

        Outputs:
            time - numpy array of times that data was taken
            datas - list of voltages, corresponding to channel 1..n
        """
        values = [self.get_voltage(i) for i in channels + 1]
        times, datas = zip(*values)
        return times[0][::-1], datas[::-1]

    def set_timestep(self, timestep):
        """
        Set the scope timebase to timestep secs/division

        Input:
            timestep: float, secs/division
        """
        current_timestep = float(
            self.scope.query_ascii_values(":TIM:SCAL?")[0])
        print(f"Setting timestep to {timestep}")
        print(f"Previous value was {current_timestep}")
        if timestep != current_timestep:
            self.scope.write(f":TIM:SCAL {timestep:f}")
            time.sleep(0.1)

    def set_timeoffset(self, offset):
        """
        Set the scope horizontal offset relative to the trigger
        Input:
            offset: float, secs [Needs to be double checked]
        """
        current_offset = float(self.scope.query_ascii_values(":TIM:OFFS?")[0])
        if offset != current_offset:
            self.scope.write(f":TIM:OFFS {offset:f}")
            time.sleep(0.1)
        print("Set time offset")

class ScopeServer(DeviceServer):

    def __init__(self, scope_name, port, device_name):
        print(port)
        super().__init__(port)
        print("Starting Scope Server")
        self.name = scope_name
        self.interface = ScopeInterface(scope_name)
        self.channels = (1)
        self.device_name = device_name

    def transition_to_buffered(self, h5_filepath):
        """
        Set timebase to correct values
        """
        with h5py.File(h5_filepath, 'r') as f:
            dset = f'/devices/{self.device_name}'
            parameters = f[dset].attrs
            timestep = parameters['timestep']
            offset = parameters['offset']
            self.channels = parameters['channels']
            self.names = parameters['names']
        self.interface.set_timestep(timestep)
        self.interface.set_timeoffset(offset)
        self.interface.run()
        return True

    def transition_to_static(self, h5_filepath):
        """
        Read scope values, write to h5 file.
        """
        times, voltages = self.interface.get_all_voltages(self.channels)
        # while True:
        #     try:
        with h5py.File(h5_filepath, 'a') as f:
            group = f['data']
            try:
                trace_group = f.create_group('ScopeTraces')
            except ValueError:
                trace_group = f['ScopeTraces']
                print("Group Already existed")
            for name, voltage in zip(self.names, voltages):
                trace_group.create_dataset('times' + name, data=times, dtype='float')
                trace_group.create_dataset(name, data=voltage, dtype='float')
            # except OSError:
            #     print("h5 file error")
            #     continue
            # break
    def abort(self):
        """To be overridden by subclasses. Return cameras and any other state
        to one in which transition_to_buffered() can be called again. abort()
        will be called if there was an exception in either
        transition_to_buffered() or transtition_to_static(), and so should
        ideally be written to return things to a sensible state even if those
        methods did not complete. Like any cleanup function, abort() should
        proceed to further cleanups even if earlier cleanups fail. As such it
        should make liberal use of try: except: blocks, so that an exception
        in performing one cleanup operation does not stop it from proceeding
        to subsequent cleanup operations"""
        print("abort")

if __name__ == '__main__':
    port = 2625

    kserver = ScopeServer("USB0::0x1AB1::0x04CE::DS1ZA205020654::INSTR", port, "RigolScope")
    kserver.shutdown_on_interrupt()