"""Functions to make 3D plots with M/EEG data
"""
from __future__ import print_function

# Authors: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
#          Denis Engemann <denis.engemann@gmail.com>
#          Martin Luessi <mluessi@nmr.mgh.harvard.edu>
#          Eric Larson <larson.eric.d@gmail.com>
#          Mainak Jas <mainak@neuro.hut.fi>
#          Mark Wronkiewicz <wronk.mark@gmail.com>
#
# License: Simplified BSD

from ..externals.six import string_types, advance_iterator

import os.path as op
import inspect
import warnings
from itertools import cycle
import base64

import numpy as np
from scipy import linalg

from ..io.pick import pick_types
from ..io.constants import FIFF
from ..surface import (get_head_surf, get_meg_helmet_surf, read_surface,
                       transform_surface_to)
from ..transforms import (read_trans, _find_trans, apply_trans,
                          combine_transforms, _get_mri_head_t, _ensure_trans,
                          invert_transform)
from ..utils import get_subjects_dir, logger, _check_subject, verbose
from ..defaults import _handle_default
from .utils import mne_analyze_colormap, _prepare_trellis, COLORS
from ..externals.six import BytesIO


def plot_evoked_field(evoked, surf_maps, time=None, time_label='t = %0.0f ms',
                      n_jobs=1):
    """Plot MEG/EEG fields on head surface and helmet in 3D

    Parameters
    ----------
    evoked : instance of mne.Evoked
        The evoked object.
    surf_maps : list
        The surface mapping information obtained with make_field_map.
    time : float | None
        The time point at which the field map shall be displayed. If None,
        the average peak latency (across sensor types) is used.
    time_label : str
        How to print info about the time instant visualized.
    n_jobs : int
        Number of jobs to run in parallel.

    Returns
    -------
    fig : instance of mlab.Figure
        The mayavi figure.
    """
    types = [t for t in ['eeg', 'grad', 'mag'] if t in evoked]

    time_idx = None
    if time is None:
        time = np.mean([evoked.get_peak(ch_type=t)[1] for t in types])

    if not evoked.times[0] <= time <= evoked.times[-1]:
        raise ValueError('`time` (%0.3f) must be inside `evoked.times`' % time)
    time_idx = np.argmin(np.abs(evoked.times - time))

    types = [sm['kind'] for sm in surf_maps]

    # Plot them
    from mayavi import mlab
    alphas = [1.0, 0.5]
    colors = [(0.6, 0.6, 0.6), (1.0, 1.0, 1.0)]
    colormap = mne_analyze_colormap(format='mayavi')
    colormap_lines = np.concatenate([np.tile([0., 0., 255., 255.], (127, 1)),
                                     np.tile([0., 0., 0., 255.], (2, 1)),
                                     np.tile([255., 0., 0., 255.], (127, 1))])

    fig = mlab.figure(bgcolor=(0.0, 0.0, 0.0), size=(600, 600))

    for ii, this_map in enumerate(surf_maps):
        surf = this_map['surf']
        map_data = this_map['data']
        map_type = this_map['kind']
        map_ch_names = this_map['ch_names']

        if map_type == 'eeg':
            pick = pick_types(evoked.info, meg=False, eeg=True)
        else:
            pick = pick_types(evoked.info, meg=True, eeg=False, ref_meg=False)

        ch_names = [evoked.ch_names[k] for k in pick]

        set_ch_names = set(ch_names)
        set_map_ch_names = set(map_ch_names)
        if set_ch_names != set_map_ch_names:
            message = ['Channels in map and data do not match.']
            diff = set_map_ch_names - set_ch_names
            if len(diff):
                message += ['%s not in data file. ' % list(diff)]
            diff = set_ch_names - set_map_ch_names
            if len(diff):
                message += ['%s not in map file.' % list(diff)]
            raise RuntimeError(' '.join(message))

        data = np.dot(map_data, evoked.data[pick, time_idx])

        x, y, z = surf['rr'].T
        nn = surf['nn']
        # make absolutely sure these are normalized for Mayavi
        nn = nn / np.sum(nn * nn, axis=1)[:, np.newaxis]

        # Make a solid surface
        vlim = np.max(np.abs(data))
        alpha = alphas[ii]
        with warnings.catch_warnings(record=True):  # traits
            mesh = mlab.pipeline.triangular_mesh_source(x, y, z, surf['tris'])
        mesh.data.point_data.normals = nn
        mesh.data.cell_data.normals = None
        mlab.pipeline.surface(mesh, color=colors[ii], opacity=alpha)

        # Now show our field pattern
        with warnings.catch_warnings(record=True):  # traits
            mesh = mlab.pipeline.triangular_mesh_source(x, y, z, surf['tris'],
                                                        scalars=data)
        mesh.data.point_data.normals = nn
        mesh.data.cell_data.normals = None
        with warnings.catch_warnings(record=True):  # traits
            fsurf = mlab.pipeline.surface(mesh, vmin=-vlim, vmax=vlim)
        fsurf.module_manager.scalar_lut_manager.lut.table = colormap

        # And the field lines on top
        with warnings.catch_warnings(record=True):  # traits
            mesh = mlab.pipeline.triangular_mesh_source(x, y, z, surf['tris'],
                                                        scalars=data)
        mesh.data.point_data.normals = nn
        mesh.data.cell_data.normals = None
        with warnings.catch_warnings(record=True):  # traits
            cont = mlab.pipeline.contour_surface(mesh, contours=21,
                                                 line_width=1.0,
                                                 vmin=-vlim, vmax=vlim,
                                                 opacity=alpha)
        cont.module_manager.scalar_lut_manager.lut.table = colormap_lines

    if '%' in time_label:
        time_label %= (1e3 * evoked.times[time_idx])
    mlab.text(0.01, 0.01, time_label, width=0.4)
    mlab.view(10, 60)
    return fig


def _plot_mri_contours(mri_fname, surf_fnames, orientation='coronal',
                       slices=None, show=True, img_output=False):
    """Plot BEM contours on anatomical slices.

    Parameters
    ----------
    mri_fname : str
        The name of the file containing anatomical data.
    surf_fnames : list of str
        The filenames for the BEM surfaces in the format
        ['inner_skull.surf', 'outer_skull.surf', 'outer_skin.surf'].
    orientation : str
        'coronal' or 'transverse' or 'sagittal'
    slices : list of int
        Slice indices.
    show : bool
        Call pyplot.show() at the end.
    img_output : None | tuple
        If tuple (width and height), images will be produced instead of a
        single figure with many axes. This mode is designed to reduce the
        (substantial) overhead associated with making tens to hundreds
        of matplotlib axes, instead opting to re-use a single Axes instance.

    Returns
    -------
    fig : Instance of matplotlib.figure.Figure | list
        The figure. Will instead be a list of png images if
        img_output is a tuple.
    """
    import matplotlib.pyplot as plt
    import nibabel as nib

    if orientation not in ['coronal', 'axial', 'sagittal']:
        raise ValueError("Orientation must be 'coronal', 'axial' or "
                         "'sagittal'. Got %s." % orientation)

    # Load the T1 data
    nim = nib.load(mri_fname)
    data = nim.get_data()
    affine = nim.get_affine()

    n_sag, n_axi, n_cor = data.shape
    orientation_name2axis = dict(sagittal=0, axial=1, coronal=2)
    orientation_axis = orientation_name2axis[orientation]

    if slices is None:
        n_slices = data.shape[orientation_axis]
        slices = np.linspace(0, n_slices, 12, endpoint=False).astype(np.int)

    # create of list of surfaces
    surfs = list()

    trans = linalg.inv(affine)
    # XXX : next line is a hack don't ask why
    trans[:3, -1] = [n_sag // 2, n_axi // 2, n_cor // 2]

    for surf_fname in surf_fnames:
        surf = dict()
        surf['rr'], surf['tris'] = read_surface(surf_fname)
        # move back surface to MRI coordinate system
        surf['rr'] = nib.affines.apply_affine(trans, surf['rr'])
        surfs.append(surf)

    if img_output is None:
        fig, axs = _prepare_trellis(len(slices), 4)
    else:
        fig, ax = plt.subplots(1, 1, figsize=(7.0, 7.0))
        axs = [ax] * len(slices)

        fig_size = fig.get_size_inches()
        w, h = img_output[0], img_output[1]
        w2 = fig_size[0]
        fig.set_size_inches([(w2 / float(w)) * w, (w2 / float(w)) * h])
        plt.close(fig)

    inds = dict(coronal=[0, 1, 2], axial=[2, 0, 1],
                sagittal=[2, 1, 0])[orientation]
    outs = []
    for ax, sl in zip(axs, slices):
        # adjust the orientations for good view
        if orientation == 'coronal':
            dat = data[:, :, sl].transpose()
        elif orientation == 'axial':
            dat = data[:, sl, :]
        elif orientation == 'sagittal':
            dat = data[sl, :, :]

        # First plot the anatomical data
        if img_output is not None:
            ax.clear()
        ax.imshow(dat, cmap=plt.cm.gray)
        ax.axis('off')

        # and then plot the contours on top
        for surf in surfs:
            ax.tricontour(surf['rr'][:, inds[0]], surf['rr'][:, inds[1]],
                          surf['tris'], surf['rr'][:, inds[2]],
                          levels=[sl], colors='yellow', linewidths=2.0)
        if img_output is not None:
            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_xlim(0, img_output[1])
            ax.set_ylim(img_output[0], 0)
            output = BytesIO()
            fig.savefig(output, bbox_inches='tight',
                        pad_inches=0, format='png')
            outs.append(base64.b64encode(output.getvalue()).decode('ascii'))
    if show:
        plt.subplots_adjust(left=0., bottom=0., right=1., top=1., wspace=0.,
                            hspace=0.)
        plt.show()

    return fig if img_output is None else outs


@verbose
def plot_trans(info, trans='auto', subject=None, subjects_dir=None,
               ch_type=None, source=('bem', 'head'), coord_frame='head',
               meg_sensors=False, dig=False, verbose=None):
    """Plot MEG/EEG head surface and helmet in 3D.

    Parameters
    ----------
    info : dict
        The measurement info.
    trans : str | 'auto' | dict
        The full path to the head<->MRI transform ``*-trans.fif`` file
        produced during coregistration.
    subject : str | None
        The subject name corresponding to FreeSurfer environment
        variable SUBJECT.
    subjects_dir : str
        The path to the freesurfer subjects reconstructions.
        It corresponds to Freesurfer environment variable SUBJECTS_DIR.
    ch_type : None | 'eeg' | 'meg'
        If None, both the MEG helmet and EEG electrodes will be shown.
        If 'meg', only the MEG helmet will be shown. If 'eeg', only the
        EEG electrodes will be shown.
    source : str
        Type to load. Common choices would be `'bem'` or `'head'`. We first
        try loading `'$SUBJECTS_DIR/$SUBJECT/bem/$SUBJECT-$SOURCE.fif'`, and
        then look for `'$SUBJECT*$SOURCE.fif'` in the same directory. Defaults
        to 'bem'. Note. For single layer bems it is recommended to use 'head'.
    coord_frame : str
        Coordinate frame to use, 'head', 'meg', or 'mri'.
    meg_sensors : bool
        If True, plot MEG sensors as points in addition to showing the helmet.
    dig : bool
        If True, plot the digitization points.
    verbose : bool, str, int, or None
        If not None, override default verbose level (see mne.verbose).

    Returns
    -------
    fig : instance of mlab.Figure
        The mayavi figure.
    """
    if coord_frame not in ['head', 'meg', 'mri']:
        raise ValueError('coord_frame must be "head" or "meg"')
    if ch_type not in [None, 'eeg', 'meg']:
        raise ValueError('Argument ch_type must be None | eeg | meg. Got %s.'
                         % ch_type)

    if isinstance(trans, string_types):
        if trans == 'auto':
            # let's try to do this in MRI coordinates so they're easy to plot
            trans = _find_trans(subject, subjects_dir)
        trans = read_trans(trans)
    elif not isinstance(trans, dict):
        raise TypeError('trans must be str or dict')
    head_mri_t = _ensure_trans(trans, 'head', 'mri')
    del trans

    # both the head and helmet will be in MRI coordinates after this
    surfs = [get_head_surf(subject, source=source, subjects_dir=subjects_dir)]
    if ch_type is None or ch_type == 'meg':
        surfs.append(get_meg_helmet_surf(info, head_mri_t))
    if coord_frame == 'meg':
        surf_trans = combine_transforms(info['dev_head_t'], head_mri_t,
                                        'meg', 'mri')
    elif coord_frame == 'head':
        surf_trans = head_mri_t
    else:  # coord_frame == 'mri'
        surf_trans = None
    surfs = [transform_surface_to(surf, coord_frame, surf_trans)
             for surf in surfs]
    del surf_trans

    # determine points
    meg_loc = list()
    ext_loc = list()
    car_loc = list()
    if ch_type is None or ch_type == 'eeg':
        eeg_loc = np.array([info['chs'][k]['loc'][:3]
                           for k in pick_types(info, meg=False, eeg=True)])
        if len(eeg_loc) > 0:
            # Transform EEG electrodes from head coordinates if necessary
            if coord_frame == 'meg':
                eeg_loc = apply_trans(invert_transform(info['dev_head_t']),
                                      eeg_loc)
            elif coord_frame == 'mri':
                eeg_loc = apply_trans(invert_transform(head_mri_t), eeg_loc)
        else:
            # only warn if EEG explicitly requested, or EEG channels exist but
            # no locations are provided
            if (ch_type is not None or
                    len(pick_types(info, meg=False, eeg=True)) > 0):
                warnings.warn('EEG electrode locations not found. '
                              'Cannot plot EEG electrodes.')
    if meg_sensors:
        meg_loc = np.array([info['chs'][k]['loc'][:3]
                           for k in pick_types(info)])
        if len(meg_loc) > 0:
            # Transform MEG coordinates from meg if necessary
            if coord_frame == 'head':
                meg_loc = apply_trans(info['dev_head_t'], meg_loc)
            elif coord_frame == 'mri':
                t = combine_transforms(info['dev_head_t'], head_mri_t,
                                       'meg', 'mri')
                meg_loc = apply_trans(t, meg_loc)
        else:
            warnings.warn('MEG electrodes not found. '
                          'Cannot plot MEG locations.')
    if dig:
        ext_loc = np.array([d['r'] for d in info['dig']
                           if d['kind'] == FIFF.FIFFV_POINT_EXTRA])
        car_loc = np.array([d['r'] for d in info['dig']
                            if d['kind'] == FIFF.FIFFV_POINT_CARDINAL])
        if coord_frame == 'meg':
            t = invert_transform(info['dev_head_t'])
            ext_loc = apply_trans(t, ext_loc)
            car_loc = apply_trans(t, car_loc)
        elif coord_frame == 'mri':
            ext_loc = apply_trans(head_mri_t, ext_loc)
            car_loc = apply_trans(head_mri_t, car_loc)
        if len(car_loc) == len(ext_loc) == 0:
            warnings.warn('Digitization points not found. '
                          'Cannot plot digitization.')

    # do the plotting, surfaces then points
    from mayavi import mlab
    fig = mlab.figure(bgcolor=(0.0, 0.0, 0.0), size=(600, 600))

    alphas = [1.0, 0.5]  # head, helmet
    colors = [(0.6, 0.6, 0.6), (0.0, 0.0, 0.6)]
    for surf, alpha, color in zip(surfs, alphas, colors):
        x, y, z = surf['rr'].T
        nn = surf['nn']
        # make absolutely sure these are normalized for Mayavi
        nn = nn / np.sum(nn * nn, axis=1)[:, np.newaxis]

        # Make a solid surface
        with warnings.catch_warnings(record=True):  # traits
            mesh = mlab.pipeline.triangular_mesh_source(x, y, z, surf['tris'])
        mesh.data.point_data.normals = nn
        mesh.data.cell_data.normals = None
        mlab.pipeline.surface(mesh, color=color, opacity=alpha)

    datas = (eeg_loc, meg_loc, car_loc, ext_loc)
    colors = ((1., 0., 0.), (0., 0.25, 0.5), (1., 1., 0.), (1., 0.5, 0.))
    alphas = (1.0, 0.25, 0.5, 0.25)
    scales = (0.005, 0.0025, 0.015, 0.0075)
    for data, color, alpha, scale in zip(datas, colors, alphas, scales):
        if len(data) > 0:
            with warnings.catch_warnings(record=True):  # traits
                mlab.points3d(data[:, 0], data[:, 1], data[:, 2],
                              color=color, scale_factor=scale, opacity=alpha)
    mlab.view(90, 90)
    return fig


def _limits_to_control_points(clim, stc_data, colormap):
    """Private helper function to convert limits (values or percentiles)
    to control points.

    Note: If using 'mne', generate cmap control points for a directly
    mirrored cmap for simplicity (i.e., no normalization is computed to account
    for a 2-tailed mne cmap).

    Parameters
    ----------
    clim : str | dict
        Desired limits use to set cmap control points.

    Returns
    -------
    ctrl_pts : list (length 3)
        Array of floats corresponding to values to use as cmap control points.
    colormap : str
        The colormap.
    """

    # Based on type of limits specified, get cmap control points
    if colormap == 'auto':
        if clim == 'auto':
            colormap = 'mne' if (stc_data < 0).any() else 'hot'
        else:
            if 'lims' in clim:
                colormap = 'hot'
            else:  # 'pos_lims' in clim
                colormap = 'mne'
    if clim == 'auto':
        # Set upper and lower bound based on percent, and get average between
        ctrl_pts = np.percentile(np.abs(stc_data), [96, 97.5, 99.95])
    elif isinstance(clim, dict):
        # Get appropriate key for clim if it's a dict
        limit_key = ['lims', 'pos_lims'][colormap in ('mne', 'mne_analyze')]
        if colormap != 'auto' and limit_key not in clim.keys():
            raise KeyError('"pos_lims" must be used with "mne" colormap')
        clim['kind'] = clim.get('kind', 'percent')
        if clim['kind'] == 'percent':
            ctrl_pts = np.percentile(np.abs(stc_data),
                                     list(np.abs(clim[limit_key])))
        elif clim['kind'] == 'value':
            ctrl_pts = np.array(clim[limit_key])
            if (np.diff(ctrl_pts) < 0).any():
                raise ValueError('value colormap limits must be strictly '
                                 'nondecreasing')
        else:
            raise ValueError('If clim is a dict, clim[kind] must be '
                             ' "value" or "percent"')
    else:
        raise ValueError('"clim" must be "auto" or dict')
    if len(ctrl_pts) != 3:
        raise ValueError('"lims" or "pos_lims" is length %i. It must be length'
                         ' 3' % len(ctrl_pts))
    ctrl_pts = np.array(ctrl_pts, float)
    if len(set(ctrl_pts)) != 3:
        if len(set(ctrl_pts)) == 1:  # three points match
            if ctrl_pts[0] == 0:  # all are zero
                warnings.warn('All data were zero')
                ctrl_pts = np.arange(3, dtype=float)
            else:
                ctrl_pts *= [0., 0.5, 1]  # all nonzero pts == max
        else:  # two points match
            # if points one and two are identical, add a tiny bit to the
            # control point two; if points two and three are identical,
            # subtract a tiny bit from point two.
            bump = 1e-5 if ctrl_pts[0] == ctrl_pts[1] else -1e-5
            ctrl_pts[1] = ctrl_pts[0] + bump * (ctrl_pts[2] - ctrl_pts[0])

    return ctrl_pts, colormap


def plot_source_estimates(stc, subject=None, surface='inflated', hemi='lh',
                          colormap='auto', time_label='time=%0.2f ms',
                          smoothing_steps=10, transparent=None, alpha=1.0,
                          time_viewer=False, config_opts=None,
                          subjects_dir=None, figure=None, views='lat',
                          colorbar=True, clim='auto'):
    """Plot SourceEstimates with PySurfer

    Note: PySurfer currently needs the SUBJECTS_DIR environment variable,
    which will automatically be set by this function. Plotting multiple
    SourceEstimates with different values for subjects_dir will cause
    PySurfer to use the wrong FreeSurfer surfaces when using methods of
    the returned Brain object. It is therefore recommended to set the
    SUBJECTS_DIR environment variable or always use the same value for
    subjects_dir (within the same Python session).

    Parameters
    ----------
    stc : SourceEstimates
        The source estimates to plot.
    subject : str | None
        The subject name corresponding to FreeSurfer environment
        variable SUBJECT. If None stc.subject will be used. If that
        is None, the environment will be used.
    surface : str
        The type of surface (inflated, white etc.).
    hemi : str, 'lh' | 'rh' | 'split' | 'both'
        The hemisphere to display.
    colormap : str | np.ndarray of float, shape(n_colors, 3 | 4)
        Name of colormap to use or a custom look up table. If array, must
        be (n x 3) or (n x 4) array for with RGB or RGBA values between
        0 and 255. If 'auto', either 'hot' or 'mne' will be chosen
        based on whether 'lims' or 'pos_lims' are specified in `clim`.
    time_label : str
        How to print info about the time instant visualized.
    smoothing_steps : int
        The amount of smoothing
    transparent : bool | None
        If True, use a linear transparency between fmin and fmid.
        None will choose automatically based on colormap type.
    alpha : float
        Alpha value to apply globally to the overlay.
    time_viewer : bool
        Display time viewer GUI.
    config_opts : dict
        Keyword arguments for Brain initialization.
        See pysurfer.viz.Brain.
    subjects_dir : str
        The path to the freesurfer subjects reconstructions.