# Authors: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
#          Matti Hamalainen <msh@nmr.mgh.harvard.edu>
#          Martin Luessi <mluessi@nmr.mgh.harvard.edu>
#          Denis Engemann <denis.engemann@gmail.com>
#          Teon Brooks <teon.brooks@gmail.com>
#
# License: BSD (3-clause)

import copy
import os
import os.path as op

import numpy as np

from ..constants import FIFF
from ..open import fiff_open, _fiff_get_fid, _get_next_fname
from ..meas_info import read_meas_info
from ..tree import dir_tree_find
from ..tag import read_tag, read_tag_info
from ..proj import make_eeg_average_ref_proj, _needs_eeg_average_ref_proj
from ..base import _BaseRaw, _RawShell, _check_raw_compatibility
from ..utils import _mult_cal_one

from ...annotations import Annotations, _combine_annotations
from ...externals.six import string_types
from ...utils import check_fname, logger, verbose, warn


class Raw(_BaseRaw):
    """Raw data in FIF format

    Parameters
    ----------
    fname : str
        The raw file to load. For files that have automatically been split,
        the split part will be automatically loaded. Filenames should end
        with raw.fif, raw.fif.gz, raw_sss.fif, raw_sss.fif.gz,
        raw_tsss.fif or raw_tsss.fif.gz.
    allow_maxshield : bool | str (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.
        Can also be "yes" to load without eliciting a warning.
    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
        Deprecated. Use :meth:`mne.io.Raw.apply_gradient_compensation`
        instead.
    add_eeg_ref : bool
        If True, add average EEG reference projector (if it's not already
        present).
    fnames : list or str
        Deprecated.
    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.
    preload : bool
        Indicates whether raw data are in memory.
    verbose : bool, str, int, or None
        See above.
    """
    @verbose
    def __init__(self, fname, allow_maxshield=False, preload=False,
                 proj=False, compensation=None, add_eeg_ref=True,
                 fnames=None, verbose=None):
        dep = ('Supplying a list of filenames with "fnames" to the Raw class '
               'has been deprecated and will be removed in 0.13. Use multiple '
               'calls to read_raw_fif with the "fname" argument followed by '
               'and concatenate_raws instead.')
        if fnames is not None:
            warn(dep, DeprecationWarning)
        else:
            fnames = fname
        del fname
        if not isinstance(fnames, list):
            fnames = [fnames]
        else:
            warn(dep, DeprecationWarning)
        fnames = [op.realpath(f) for f in fnames]
        split_fnames = []

        raws = []
        for ii, fname in enumerate(fnames):
            do_check_fname = fname not in split_fnames
            raw, next_fname = self._read_raw_file(fname, allow_maxshield,
                                                  preload, do_check_fname)
            raws.append(raw)
            if next_fname is not None:
                if not op.exists(next_fname):
                    warn('Split raw file detected but next file %s does not '
                         'exist.' % next_fname)
                    continue
                if next_fname in fnames:
                    # the user manually specified the split files
                    logger.info('Note: %s is part of a split raw file. It is '
                                'not necessary to manually specify the parts '
                                'in this case; simply construct Raw using '
                                'the name of the first file.' % next_fname)
                    continue

                # process this file next
                fnames.insert(ii + 1, next_fname)
                split_fnames.append(next_fname)

        _check_raw_compatibility(raws)

        super(Raw, self).__init__(
            copy.deepcopy(raws[0].info), False,
            [r.first_samp for r in raws], [r.last_samp for r in raws],
            [r.filename for r in raws], [r._raw_extras for r in raws],
            raws[0].orig_format, None, verbose=verbose)

        # combine information from each raw file to construct self
        if add_eeg_ref and _needs_eeg_average_ref_proj(self.info):
            eeg_ref = make_eeg_average_ref_proj(self.info, activate=False)
            self.add_proj(eeg_ref)

        # combine annotations
        self.annotations = raws[0].annotations
        if any([r.annotations for r in raws[1:]]):
            first_samps = list()
            last_samps = list()
            for r in raws:
                first_samps = np.r_[first_samps, r.first_samp]
                last_samps = np.r_[last_samps, r.last_samp]
                self.annotations = _combine_annotations((self.annotations,
                                                         r.annotations),
                                                        last_samps,
                                                        first_samps,
                                                        r.info['sfreq'])
        if compensation is not None:
            warn('The "compensation" argument has been deprecated '
                 'in favor of the "raw.apply_gradient_compensation" '
                 'method and will be removed in 0.14',
                 DeprecationWarning)
            self.apply_gradient_compensation(compensation)
        if preload:
            self._preload_data(preload)
        else:
            self.preload = False

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

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

        if do_check_fname:
            check_fname(fname, 'raw', ('raw.fif', 'raw_sss.fif',
                                       'raw_tsss.fif', 'raw.fif.gz',
                                       'raw_sss.fif.gz', 'raw_tsss.fif.gz'))

        #   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
        ff, tree, _ = fiff_open(fname, preload=whole_file)
        with ff as fid:
            #   Read the measurement info

            info, meas = read_meas_info(fid, tree, clean_bads=True)

            annotations = None
            annot_data = dir_tree_find(tree, FIFF.FIFFB_MNE_ANNOTATIONS)
            if len(annot_data) > 0:
                annot_data = annot_data[0]
                for k in range(annot_data['nent']):
                    kind = annot_data['directory'][k].kind
                    pos = annot_data['directory'][k].pos
                    orig_time = None
                    tag = read_tag(fid, pos)
                    if kind == FIFF.FIFF_MNE_BASELINE_MIN:
                        onset = tag.data
                    elif kind == FIFF.FIFF_MNE_BASELINE_MAX:
                        duration = tag.data - onset
                    elif kind == FIFF.FIFF_COMMENT:
                        description = tag.data.split(':')
                        description = [d.replace(';', ':') for d in
                                       description]
                    elif kind == FIFF.FIFF_MEAS_DATE:
                        orig_time = float(tag.data)
                annotations = Annotations(onset, duration, description,
                                          orig_time)

            #   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 (len(raw_node) == 0):
                    raw_node = dir_tree_find(meas, FIFF.FIFFB_SMSH_RAW_DATA)
                    msg = ('This file contains raw Internal Active '
                           'Shielding data. It may be distorted. Elekta '
                           'recommends it be run through MaxFilter to '
                           'produce reliable results. Consider closing '
                           'the file and running MaxFilter on the data.')
                    if (len(raw_node) == 0):
                        raise ValueError('No raw data in %s' % fname)
                    elif allow_maxshield:
                        info['maxshield'] = True
                        if not (isinstance(allow_maxshield, string_types) and
                                allow_maxshield == 'yes'):
                            warn(msg)
                    else:
                        msg += (' Use allow_maxshield=True if you are sure you'
                                ' want to load the data despite this warning.')
                        raise ValueError(msg)

            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
                _check_entry(first, nent)

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

            raw = _RawShell()
            raw.filename = fname
            raw.first_samp = first_samp
            raw.annotations = annotations

            #   Go through the remaining tags in the directory
            raw_extras = 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:
                        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'