# Authors: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
#          Eric Larson <larson.eric.d@gmail.com>
#
# License: Simplified BSD

import numpy as np
from scipy import linalg
from copy import deepcopy
import re

from .cov import read_cov, _get_whitener_data
from .io.constants import FIFF
from .io.pick import pick_types
from .io.proj import make_projector
from .bem import _fit_sphere
from .transforms import (_print_coord_trans, _coord_frame_name,
                         apply_trans, invert_transform)

from .forward._make_forward import (_get_mri_head_t, _setup_bem,
                                    _prep_channels)
from .forward._compute_forward import (_compute_forwards_meeg,
                                       _prep_field_computation)

from .externals.six import string_types
from .surface import (_bem_find_surface, transform_surface_to,
                      _normalize_vectors, _get_ico_surface,
                      _bem_explain_surface, _compute_nearest)
from .source_space import (_make_volume_source_space, SourceSpaces,
                           _points_outside_surface)
from .parallel import parallel_func
from .fixes import partial
from .utils import logger, verbose, _time_mask


class Dipole(object):
    """Dipole class

    Used to store positions, orientations, amplitudes, times, goodness of fit
    of dipoles, typically obtained with Neuromag/xfit, mne_dipole_fit
    or certain inverse solvers.

    Parameters
    ----------
    times : array, shape (n_dipoles,)
        The time instants at which each dipole was fitted (sec).
    pos : array, shape (n_dipoles, 3)
        The dipoles positions (m).
    amplitude : array, shape (n_dipoles,)
        The amplitude of the dipoles (nAm).
    ori : array, shape (n_dipoles, 3)
        The dipole orientations (normalized to unit length).
    gof : array, shape (n_dipoles,)
        The goodness of fit.
    name : str | None
        Name of the dipole.
    """
    def __init__(self, times, pos, amplitude, ori, gof, name=None):
        self.times = times
        self.pos = pos
        self.amplitude = amplitude
        self.ori = ori
        self.gof = gof
        self.name = name

    def __repr__(self):
        s = "n_times : %s" % len(self.times)
        s += ", tmin : %s" % np.min(self.times)
        s += ", tmax : %s" % np.max(self.times)
        return "<Dipole  |  %s>" % s

    def save(self, fname):
        """Save dipole in a .dip file

        Parameters
        ----------
        fname : str
            The name of the .dip file.
        """
        fmt = "  %7.1f %7.1f %8.2f %8.2f %8.2f %8.3f %8.3f %8.3f %8.3f %6.1f"
        with open(fname, 'wb') as fid:
            fid.write('# CoordinateSystem "Head"\n'.encode('utf-8'))
            fid.write('#   begin     end   X (mm)   Y (mm)   Z (mm)'
                      '   Q(nAm)  Qx(nAm)  Qy(nAm)  Qz(nAm)    g/%\n'
                      .encode('utf-8'))
            t = self.times[:, np.newaxis] * 1000.
            gof = self.gof[:, np.newaxis]
            amp = 1e9 * self.amplitude[:, np.newaxis]
            out = np.concatenate((t, t, self.pos / 1e-3, amp,
                                  self.ori * amp, gof), axis=-1)
            np.savetxt(fid, out, fmt=fmt)
            if self.name is not None:
                fid.write(('## Name "%s dipoles" Style "Dipoles"'
                           % self.name).encode('utf-8'))

    def crop(self, tmin=None, tmax=None):
        """Crop data to a given time interval

        Parameters
        ----------
        tmin : float | None
            Start time of selection in seconds.
        tmax : float | None
            End time of selection in seconds.
        """
        mask = _time_mask(self.times, tmin, tmax)
        for attr in ('times', 'pos', 'gof', 'amplitude', 'ori'):
            setattr(self, attr, getattr(self, attr)[mask])

    def copy(self):
        """Copy the Dipoles object

        Returns
        -------
        dip : instance of Dipole
            The copied dipole instance.
        """
        return deepcopy(self)

    @verbose
    def plot_locations(self, trans, subject, subjects_dir=None,
                       bgcolor=(1, 1, 1), opacity=0.3,
                       brain_color=(0.7, 0.7, 0.7), mesh_color=(1, 1, 0),
                       fig_name=None, fig_size=(600, 600), mode='cone',
                       scale_factor=0.1e-1, colors=None, verbose=None):
        """Plot dipole locations as arrows

        Parameters
        ----------
        trans : dict
            The mri to head trans.
        subject : str
            The subject name corresponding to FreeSurfer environment
            variable SUBJECT.
        subjects_dir : None | str
            The path to the freesurfer subjects reconstructions.
            It corresponds to Freesurfer environment variable SUBJECTS_DIR.
            The default is None.
        bgcolor : tuple of length 3
            Background color in 3D.
        opacity : float in [0, 1]
            Opacity of brain mesh.
        brain_color : tuple of length 3
            Brain color.
        mesh_color : tuple of length 3
            Mesh color.
        fig_name : tuple of length 2
            Mayavi figure name.
        fig_size : tuple of length 2
            Mayavi figure size.
        mode : str
            Should be ``'cone'`` or ``'sphere'`` to specify how the
            dipoles should be shown.
        scale_factor : float
            The scaling applied to amplitudes for the plot.
        colors: list of colors | None
            Color to plot with each dipole. If None defaults colors are used.
        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.
        """
        from .viz import plot_dipole_locations
        dipoles = []
        for t in self.times:
            dipoles.append(self.copy())
            dipoles[-1].crop(t, t)
        return plot_dipole_locations(
            dipoles, trans, subject, subjects_dir, bgcolor, opacity,
            brain_color, mesh_color, fig_name, fig_size, mode, scale_factor,
            colors)

    def plot_amplitudes(self, color='k', show=True):
        """Plot the dipole amplitudes as a function of time

        Parameters
        ----------
        color: matplotlib Color
            Color to use for the trace.
        show : bool
            Show figure if True.

        Returns
        -------
        fig : matplotlib.figure.Figure
            The figure object containing the plot.
        """
        from .viz import plot_dipole_amplitudes
        return plot_dipole_amplitudes([self], [color], show)

    def __getitem__(self, idx_slice):
        """Handle indexing"""
        if isinstance(idx_slice, int):  # make sure attributes stay 2d
            idx_slice = [idx_slice]

        selected_times = self.times[idx_slice].copy()
        selected_pos = self.pos[idx_slice, :].copy()
        selected_amplitude = self.amplitude[idx_slice].copy()
        selected_ori = self.ori[idx_slice, :].copy()
        selected_gof = self.gof[idx_slice].copy()
        selected_name = self.name

        new_dipole = Dipole(selected_times, selected_pos,
                            selected_amplitude, selected_ori,
                            selected_gof, selected_name)
        return new_dipole

    def __len__(self):
        """Handle len function"""
        return self.pos.shape[0]


# #############################################################################
# IO

@verbose
def read_dipole(fname, verbose=None):
    """Read .dip file from Neuromag/xfit or MNE

    Parameters
    ----------
    fname : str
        The name of the .dip file.
    verbose : bool, str, int, or None
        If not None, override default verbose level (see mne.verbose).

    Returns
    -------
    dipole : instance of Dipole
        The dipole.
    """
    try:
        data = np.loadtxt(fname, comments='%')
    except:
        data = np.loadtxt(fname, comments='#')  # handle 2 types of comments...
    name = None
    with open(fname, 'r') as fid:
        for line in fid.readlines():
            if line.startswith('##') or line.startswith('%%'):
                m = re.search('Name "(.*) dipoles"', line)
                if m:
                    name = m.group(1)
                    break
    if data.ndim == 1:
        data = data[None, :]
    logger.info("%d dipole(s) found" % len(data))
    times = data[:, 0] / 1000.
    pos = 1e-3 * data[:, 2:5]  # put data in meters
    amplitude = data[:, 5]
    norm = amplitude.copy()
    amplitude /= 1e9
    norm[norm == 0] = 1
    ori = data[:, 6:9] / norm[:, np.newaxis]
    gof = data[:, 9]
    return Dipole(times, pos, amplitude, ori, gof, name)


# #############################################################################
# Fitting

def _dipole_forwards(fwd_data, whitener, rr, n_jobs=1):
    """Compute the forward solution and do other nice stuff"""
    B = _compute_forwards_meeg(rr, fwd_data, n_jobs, verbose=False)
    B = np.concatenate(B, axis=1)
    B_orig = B.copy()

    # Apply projection and whiten (cov has projections already)
    B = np.dot(B, whitener.T)

    # column normalization doesn't affect our fitting, so skip for now
    # S = np.sum(B * B, axis=1)  # across channels
    # scales = np.repeat(3. / np.sqrt(np.sum(np.reshape(S, (len(rr), 3)),
    #                                        axis=1)), 3)
    # B *= scales[:, np.newaxis]
    scales = np.ones(3)
    return B, B_orig, scales


def _make_guesses(surf_or_rad, r0, grid, exclude, mindist, n_jobs):
    """Make a guess space inside a sphere or BEM surface"""
    if isinstance(surf_or_rad, dict):
        surf = surf_or_rad
        logger.info('Guess surface (%s) is in %s coordinates'
                    % (_bem_explain_surface(surf['id']),
                       _coord_frame_name(surf['coord_frame'])))
    else:
        radius = surf_or_rad[0]
        logger.info('Making a spherical guess space with radius %7.1f mm...'
                    % (1000 * radius))
        surf = _get_ico_surface(3)
        _normalize_vectors(surf['rr'])
        surf['rr'] *= radius
        surf['rr'] += r0
    logger.info('Filtering (grid = %6.f mm)...' % (1000 * grid))
    src = _make_volume_source_space(surf, grid, exclude, 1000 * mindist,
                                    do_neighbors=False, n_jobs=n_jobs)
    # simplify the result to make things easier later
    src = dict(rr=src['rr'][src['vertno']], nn=src['nn'][src['vertno']],
               nuse=src['nuse'], coord_frame=src['coord_frame'],
               vertno=np.arange(src['nuse']))
    return SourceSpaces([src])


def _fit_eval(rd, B, B2, fwd_svd=None, fwd_data=None, whitener=None,
              constraint=None):
    """Calculate the residual sum of squares"""
    if fwd_svd is None:
        dist = constraint(rd)