# Authors: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
#          Matti Hamalainen <msh@nmr.mgh.harvard.edu>
#
# License: BSD (3-clause)

from copy import deepcopy
from functools import partial
from gzip import GzipFile
import os
import os.path as op

import numpy as np
from scipy import sparse, linalg

from .io.constants import FIFF
from .io.meas_info import create_info
from .io.tree import dir_tree_find
from .io.tag import find_tag, read_tag
from .io.open import fiff_open
from .io.write import (start_block, end_block, write_int,
                       write_float_sparse_rcs, write_string,
                       write_float_matrix, write_int_matrix,
                       write_coord_trans, start_file, end_file, write_id)
from .bem import read_bem_surfaces, ConductorModel
from .surface import (read_surface, _create_surf_spacing, _get_ico_surface,
                      _tessellate_sphere_surf, _get_surf_neighbors,
                      _normalize_vectors, _get_solids, _triangle_neighbors,
                      complete_surface_info, _compute_nearest, fast_cross_3d,
                      mesh_dist)
from .utils import (get_subjects_dir, run_subprocess, has_freesurfer,
                    has_nibabel, check_fname, logger, verbose, _ensure_int,
                    check_version, _get_call_line, warn, _check_fname)
from .parallel import parallel_func, check_n_jobs
from .transforms import (invert_transform, apply_trans, _print_coord_trans,
                         combine_transforms, _get_trans,
                         _coord_frame_name, Transform, _str_to_frame,
                         _ensure_trans, _read_fs_xfm)


def _get_lut():
    """Get the FreeSurfer LUT."""
    data_dir = op.join(op.dirname(__file__), 'data')
    lut_fname = op.join(data_dir, 'FreeSurferColorLUT.txt')
    dtype = [('id', '<i8'), ('name', 'U47'),
             ('R', '<i8'), ('G', '<i8'), ('B', '<i8'), ('A', '<i8')]
    return np.genfromtxt(lut_fname, dtype=dtype)


def _get_lut_id(lut, label, use_lut):
    """Convert a label to a LUT ID number."""
    if not use_lut:
        return 1
    assert isinstance(label, str)
    mask = (lut['name'] == label)
    assert mask.sum() == 1
    return lut['id'][mask]


_src_kind_dict = {
    'vol': 'volume',
    'surf': 'surface',
    'discrete': 'discrete',
}


class SourceSpaces(list):
    """Represent a list of source space.

    Currently implemented as a list of dictionaries containing the source
    space information

    Parameters
    ----------
    source_spaces : list
        A list of dictionaries containing the source space information.
    info : dict
        Dictionary with information about the creation of the source space
        file. Has keys 'working_dir' and 'command_line'.

    Attributes
    ----------
    info : dict
        Dictionary with information about the creation of the source space
        file. Has keys 'working_dir' and 'command_line'.
    """

    def __init__(self, source_spaces, info=None):  # noqa: D102
        super(SourceSpaces, self).__init__(source_spaces)
        if info is None:
            self.info = dict()
        else:
            self.info = dict(info)

    @verbose
    def plot(self, head=False, brain=None, skull=None, subjects_dir=None,
             trans=None, verbose=None):
        """Plot the source space.

        Parameters
        ----------
        head : bool
            If True, show head surface.
        brain : bool | str
            If True, show the brain surfaces. Can also be a str for
            surface type (e.g., 'pial', same as True). Default is None,
            which means 'white' for surface source spaces and False otherwise.
        skull : bool | str | list of str | list of dict | None
            Whether to plot skull surface. If string, common choices would be
            'inner_skull', or 'outer_skull'. Can also be a list to plot
            multiple skull surfaces. If a list of dicts, each dict must
            contain the complete surface info (such as you get from
            :func:`mne.make_bem_model`). True is an alias of 'outer_skull'.
            The subjects bem and bem/flash folders are searched for the 'surf'
            files. Defaults to None, which is False for surface source spaces,
            and True otherwise.
        subjects_dir : string, or None
            Path to SUBJECTS_DIR if it is not set in the environment.
        trans : str | 'auto' | dict | None
            The full path to the head<->MRI transform ``*-trans.fif`` file
            produced during coregistration. If trans 
is None, an identity
            matr
ix is assumed. This is only needed when the source space is in
            head coordinates.
        %(verbose_meth)s

        Returns
        -------
        fig : instance of mayavi.mlab.Figure
            The figure.
        """
        from .viz import plot_alignment

        surfaces = list()
        bem = None

        if brain is None:
            brain = 'white' if any(ss['type'] == 'surf'
                                   for ss in self) else False

        if isinstance(brain, str):
            surfaces.append(brain)
        elif brain:
            surfaces.append('brain')

        if skull is None:
            skull = False if self.kind == 'surface' else True

        if isinstance(skull, str):
            surfaces.append(skull)
        elif skull is True:
            surfaces.append('outer_skull')
        elif skull is not False:  # list
            if isinstance(skull[0], dict):  # bem
                skull_map = {FIFF.FIFFV_BEM_SURF_ID_BRAIN: 'inner_skull',
                             FIFF.FIFFV_BEM_SURF_ID_SKULL: 'outer_skull',
                             FIFF.FIFFV_BEM_SURF_ID_HEAD: 'outer_skin'}
                for this_skull in skull:
                    surfaces.append(skull_map[this_skull['id']])
                bem = skull
            else:  # list of str
                for surf in skull:
                    surfaces.append(surf)

        if head:
            surfaces.append('head')

        if self[0]['coord_frame'] == FIFF.FIFFV_COORD_HEAD:
            coord_frame = 'head'
            if trans is None:
                raise ValueError('Source space is in head coordinates, but no '
                                 'head<->MRI transform was given. Please '
                                 'specify the full path to the appropriate '
                                 '*-trans.fif file as the "trans" parameter.')
        else:
            coord_frame = 'mri'

        info = create_info(0, 1000., 'eeg')

        return plot_alignment(
            info, trans=trans, subject=self[0]['subject_his_id'],
            subjects_dir=subjects_dir, surfaces=surfaces,
            coord_frame=coord_frame, meg=(), eeg=False, dig=False, ecog=False,
            bem=bem, src=self
        )

    def __repr__(self):  # noqa: D105
        ss_repr = []
        for ss in self:
            ss_type = ss['type']
            r = _src_kind_dict[ss_type]
            if ss_type == 'vol':
                if 'seg_name' in ss:
                    r += " (%s)" % (ss['seg_name'],)
                else:
                    r += ", shape=%s" % (ss['shape'],)
            elif ss_type == 'surf':
                r += (" (%s), n_vertices=%i" % (_get_hemi(ss)[0], ss['np']))
            r += (', n_used=%i, coordinate_frame=%s'
                  % (ss['nuse'], _coord_frame_name(int(ss['coord_frame']))))
            ss_repr.append('<%s>' % r)
        return "<SourceSpaces: [%s]>" % ', '.join(ss_repr)

    @property
    def kind(self):
        """The kind of source space (surface, volume, discrete, mixed)."""
        ss_types = list({ss['type'] for ss in self})
        if len(ss_types) != 1:
            return 'mixed'
        return _src_kind_dict[ss_types[0]]

    def __add__(self, other):
        """Combine source spaces."""
        return SourceSpaces(list.__add__(self, other))

    def copy(self):
        """Make a copy of the source spaces.

        Returns
        -------
        src : instance of SourceSpaces
            The copied source spaces.
        """
        src = deepcopy(self)
        return src

  