# Author: Alexandre Gramfort <alexandre.gramfort@inria.fr>
#         Daniel Strohmeier <daniel.strohmeier@tu-ilmenau.de>
#
# License: Simplified BSD

import os.path as op

import numpy as np
from numpy.testing import (assert_array_almost_equal, assert_allclose,
                           assert_array_less, assert_array_equal)
import pytest

import mne
from mne.datasets import testing
from mne.label import read_label
from mne import (read_cov, read_forward_solution, read_evokeds,
                 convert_forward_solution)
from mne.inverse_sparse import mixed_norm, tf_mixed_norm
from mne.inverse_sparse.mxne_inverse import make_stc_from_dipoles, _split_gof
from mne.inverse_sparse.mxne_inverse import _compute_mxne_sure
from mne.inverse_sparse.mxne_optim import norm_l2inf
from mne.minimum_norm import apply_inverse, make_inverse_operator
from mne.minimum_norm.tests.test_inverse import \
    assert_var_exp_log, assert_stc_res
from mne.utils import assert_stcs_equal, catch_logging, _record_warnings
from mne.dipole import Dipole
from mne.source_estimate import VolSourceEstimate
from mne.simulation import simulate_sparse_stc, simulate_evoked


data_path = testing.data_path(download=False)
# NOTE: These use the ave and cov from sample dataset (no _trunc)
fname_data = op.join(data_path, 'MEG', 'sample', 'sample_audvis-ave.fif')
fname_cov = op.join(data_path, 'MEG', 'sample', 'sample_audvis-cov.fif')
fname_raw = op.join(data_path, 'MEG', 'sample', 'sample_audvis_trunc_raw.fif')
fname_fwd = op.join(data_path, 'MEG', 'sample',
                    'sample_audvis_trunc-meg-eeg-oct-6-fwd.fif')
label = 'Aud-rh'
fname_label = op.join(data_path, 'MEG', 'sample', 'labels', '%s.label' % label)


@pytest.fixture(scope='module', params=[testing._pytest_param])
def forward():
    """Get a forward solution."""
    # module scope it for speed (but don't overwrite in use!)
    return read_forward_solution(fname_fwd)


@testing.requires_testing_data
@pytest.mark.timeout(150)  # ~30 sec on Travis Linux
@pytest.mark.slowtest
def test_mxne_inverse_standard(forward):
    """Test (TF-)MxNE inverse computation."""
    # Read noise covariance matrix
    cov = read_cov(fname_cov)

    # Handling average file
    loose = 0.0
    depth = 0.9

    evoked = read_evokeds(fname_data, condition=0, baseline=(None, 0))
    evoked.crop(tmin=-0.05, tmax=0.2)

    evoked_l21 = evoked.copy()
    evoked_l21.crop(tmin=0.081, tmax=0.1)
    label = read_label(fname_label)
    assert label.hemi == 'rh'

    forward = convert_forward_solution(forward, surf_ori=True)

    # Reduce source space to make test computation faster
    inverse_operator = make_inverse_operator(evoked_l21.info, forward, cov,
                                             loose=loose, depth=depth,
                                             fixed=True, use_cps=True)
    stc_dspm = apply_inverse(evoked_l21, inverse_operator, lambda2=1. / 9.,
                             method='dSPM')
    stc_dspm.data[np.abs(stc_dspm.data) < 12] = 0.0
    stc_dspm.data[np.abs(stc_dspm.data) >= 12] = 1.
    weights_min = 0.5

    # MxNE tests
    alpha = 70  # spatial regularization parameter

    with _record_warnings():  # CD
        stc_cd = mixed_norm(evoked_l21, forward, cov, alpha, loose=loose,
                            depth=depth, maxit=300, tol=1e-8,
                            active_set_size=10, weights=stc_dspm,
                            weights_min=weights_min, solver='cd')
    stc_bcd = mixed_norm(evoked_l21, forward, cov, alpha, loose=loose,
                         depth=depth, maxit=300, tol=1e-8, active_set_size=10,
                         weights=stc_dspm, weights_min=weights_min,
                         solver='bcd')
    assert_array_almost_equal(stc_cd.times, evoked_l21.times, 5)
    assert_array_almost_equal(stc_bcd.times, evoked_l21.times, 5)
    assert_allclose(stc_cd.data, stc_bcd.data, rtol=1e-3, atol=0.0)
    assert stc_cd.vertices[1][0] in label.vertices
    assert stc_bcd.vertices[1][0] in label.vertices

    # vector
    with _record_warnings():  # no convergence
        stc = mixed_norm(evoked_l21, forward, cov, alpha, loose=1, maxit=2)
    with _record_warnings():  # no convergence
        stc_vec = mixed_norm(evoked_l21, forward, cov, alpha, loose=1, maxit=2,
                             pick_ori='vector')
    assert_stcs_equal(stc_vec.magnitude(), stc)
    with _record_warnings(), \
            pytest.raises(ValueError, match='pick_ori='):
        mixed_norm(evoked_l21, forward, cov, alpha, loose=0, maxit=2,
                   pick_ori='vector')

    with _record_warnings(), catch_logging() as log:  # CD
        dips = mixed_norm(evoked_l21, forward, cov, alpha, loose=loose,
                          depth=depth, maxit=300, tol=1e-8, active_set_size=10,
                          weights=stc_dspm, weights_min=weights_min,
                          solver='cd', return_as_dipoles=True, verbose=True)
    stc_dip = make_stc_from_dipoles(dips, forward['src'])
    assert isinstance(dips[0], Dipole)
    assert stc_dip.subject == "sample"
    assert_stcs_equal(stc_cd, stc_dip)
    assert_var_exp_log(log.getvalue(), 51, 53)  # 51.8

    # Single time point things should match
    with _record_warnings(), catch_logging() as log:
        dips = mixed_norm(evoked_l21.copy().crop(0.081, 0.081),
                          forward, cov, alpha, loose=loose,
                          depth=depth, maxit=300, tol=1e-8, active_set_size=10,
                          weights=stc_dspm, weights_min=weights_min,
                          solver='cd', return_as_dipoles=True, verbose=True)
    assert_var_exp_log(log.getvalue(), 37.8, 38.0)  # 37.9
    gof = sum(dip.gof[0] for dip in dips)  # these are now partial exp vars
    assert_allclose(gof, 37.9, atol=0.1)

    with _record_warnings(), catch_logging() as log:
        stc, res = mixed_norm(evoked_l21, forward, cov, alpha, loose=loose,
                              depth=depth, maxit=300, tol=1e-8,
                              weights=stc_dspm,  # gh-6382
                              active_set_size=10, return_residual=True,
                              solver='cd', verbose=True)
    assert_array_almost_equal(stc.times, evoked_l21.times, 5)
    assert stc.vertices[1][0] in label.vertices
    assert_var_exp_log(log.getvalue(), 51, 53)  # 51.8
    assert stc.data.min() < -1e-9  # signed
    assert_stc_res(evoked_l21, stc, forward, res)

    # irMxNE tests
    with _record_warnings(), catch_logging() as log:  # CD
        stc, residual = mixed_norm(
            evoked_l21, forward, cov, alpha, n_mxne_iter=5, loose=0.0001,
            depth=depth, maxit=300, tol=1e-8, active_set_size=10,
            solver='cd', return_residual=True, pick_ori='vector', verbose=True)
    assert_array_almost_equal(stc.times, evoked_l21.times, 5)
    assert stc.vertices[1][0] in label.vertices
    assert stc.vertices == [[63152], [79017]]
    assert_var_exp_log(log.getvalue(), 51, 53)  # 51.8
    assert_stc_res(evoked_l21, stc, forward, residual)

    # Do with TF-MxNE for test memory savings
    alpha = 60.  # overall regularization parameter
    l1_ratio = 0.01  # temporal regularization proportion

    stc, _ = tf_mixed_norm(evoked, forward, cov,
                           loose=loose, depth=depth, maxit=100, tol=1e-4,
                           tstep=4, wsize=16, window=0.1, weights=stc_dspm,
                           weights_min=weights_min, return_residual=True,
                           alpha=alpha, l1_ratio=l1_ratio)
    assert_array_almost_equal(stc.times, evoked.times, 5)
    assert stc.vertices[1][0] in label.vertices

    # vector
    stc_nrm = tf_mixed_norm(
        evoked, forward, cov, loose=1, depth=depth, maxit=2, tol=1e-4,
        tstep=4, wsize=16, window=0.1, weights=stc_dspm,
        weights_min=weights_min, alpha=alpha, l1_ratio=l1_ratio)
    stc_vec, residual = tf_mixed_norm(
        evoked, forward, cov, loose=1, depth=depth, maxit=2, tol=1e-4,
        tstep=4, wsize=16, window=0.1, weights=stc_dspm,
        weights_min=weights_min, alpha=alpha, l1_ratio=l1_ratio,
        pick_ori='vector', return_residual=True)
    assert_stcs_equal(stc_vec.magnitude(), stc_nrm)

    pytest.raises(ValueError, tf_mixed_norm, evoked, forward, cov,
                  alpha=101, l1_ratio=0.03)
    pytest.raises(ValueError, tf_mixed_norm, evoked, forward, cov,
                  alpha=50., l1_ratio=1.01)


@pytest.mark.slowtest
@testing.requires_testing_data
def test_mxne_vol_sphere():
    """Test (TF-)MxNE with a sphere forward and volumic source space."""
    evoked = read_evokeds(fname_data, condition=0, baseline=(None, 0))
    evoked.crop(tmin=-0.05, tmax=0.2)
    cov = read_cov(fname_cov)

    evoked_l21 = evoked.copy()
    evoked_l21.crop(tmin=0.081, tmax=0.1)

    info = evoked.info
    sphere = mne.make_sphere_model(r0=(0., 0., 0.), head_radius=0.080)
    src = mne.setup_volume_source_space(subject=None, pos=15., mri=None,
                                        sphere=(0.0, 0.0, 0.0, 0.08),
                                        bem=None, mindist=5.0,
                                        exclude=2.0, sphere_units='m')
    fwd = mne.make_forward_solution(info, trans=None, src=src,
                                    bem=sphere, eeg=False, meg=True)

    alpha = 80.

    # Computing inverse with restricted orientations should also work, since
    # we have a discrete source space.
    stc = mixed_norm(evoked_l21, fwd, cov, alpha, loose=0.2,
                     return_residual=False, maxit=3, tol=1e-8,
                     active_set_size=10)
    assert_array_almost_equal(stc.times, evoked_l21.times, 5)

    # irMxNE tests
    with catch_logging() as log:
        stc = mixed_norm(evoked_l21, fwd, cov, alpha,
                         n_mxne_iter=1, maxit=30, tol=1e-8,
                         active_set_size=10, verbose=True)
    assert isinstance(stc, VolSourceEstimate)
    assert_array_almost_equal(stc.times, evoked_l21.times, 5)
    assert_var_exp_log(log.getvalue(), 9, 11)  # 10.2

    # Compare orientation obtained using fit_dipole and gamma_map
    # for a simulated evoked containing a single dipole
    stc = mne.VolSourceEstimate(50e-9 * np.random.RandomState(42).randn(1, 4),
                                vertices=[stc.vertices[0][:1]],
                                tmin=stc.tmin,
                                tstep=stc.tstep)
    evoked_dip = mne.simulation.simulate_evoked(fwd, stc, info, cov, nave=1e9,
                                                use_cps=True)

    dip_mxne = mixed_norm(evoked_dip, fwd, cov, alpha=80,
                          n_mxne_iter=1, maxit=30, tol=1e-8,
                          active_set_size=10, return_as_dipoles=True)

    amp_max = [np.max(d.amplitude) for d in dip_mxne]
    dip_mxne = dip_mxne[np.argmax(amp_max)]
    assert dip_mxne.pos[0] in src[0]['rr'][stc.vertices[0]]

    dip_fit = mne.fit_dipole(evoked_dip, cov, sphere)[0]
    assert np.abs(np.dot(dip_fit.ori[0], dip_mxne.ori[0])) > 0.99
    dist = 1000 * np.linalg.norm(dip_fit.pos[0] - dip_mxne.pos[0])
    assert dist < 4.  # within 4 mm

    # Do with TF-MxNE for test memory savings
    alpha = 60.  # overall regularization parameter
    l1_ratio = 0.01  # temporal regularization proportion

    stc, _ = tf_mixed_norm(evoked, fwd, cov, maxit=3, tol=1e-4,
                           tstep=16, wsize=32, window=0.1, alpha=alpha,
                           l1_ratio=l1_ratio, return_residual=True)
    assert isinstance(stc, VolSourceEstimate)
    assert_array_almost_equal(stc.times, evoked.times, 5)


@pytest.mark.parametrize('mod', (
    None, 'mult', 'augment', 'sign', 'zero', 'less'))
def test_split_gof_basic(mod):
    """Test splitting the goodness of fit."""
    # first a trivial case
    gain = np.array([[0., 1., 1.], [1., 1., 0.]]).T
    M = np.ones((3, 1))
    X = np.ones((2, 1))
    M_est = gain @ X
    assert_allclose(M_est, np.array([[1., 2., 1.]]).T)  # a reasonable estimate
    if mod == 'mult':
        gain *= [1., -0.5]
        X[1] *= -2
    elif mod == 'augment':
        gain = np.concatenate((gain, np.zeros((3, 1))), axis=1)
        X = np.concatenate((X, [[1.]]))
    elif mod == 'sign':
        gain[1] *= -1
        M[1] *= -1
        M_est[1] *= -1
    elif mod in ('zero', 'less'):
        gain = np.array([[1, 1., 1.], [1., 1., 1.]]).T
        if mod == 'zero':
            X[:, 0] = [1., 0.]
        else:
            X[:, 0] = [1., 0.5]
        M_est = gain @ X
    else:
        assert mod is None
    res = M - M_est
    gof = 100 * (1. - (res * res).sum() / (M * M).sum())
    gof_split = _split_gof(M, X, gain)
    assert_allclose(gof_split.sum(), gof)
    want = gof_split[[0, 0]]
    if mod == 'augment':
        want = np.concatenate((want, [[0]]))
    if mod in ('mult', 'less'):
        assert_array_less(gof_split[1], gof_split[0])
    elif mod == 'zero':
        assert_allclose(gof_split[0], gof_split.sum(0))
        assert_allclose(gof_split[1], 0., atol=1e-6)
    else:
        assert_allclose(gof_split, want, atol=1e-12)


@testing.requires_testing_data
@pytest.mark.parametrize('idx, weights', [
    # empirically determined approximately orthogonal columns: 0, 15157, 19448
    ([0], [1]),
    ([0, 15157], [1, 1]),
    ([0, 15157], [1, 3]),
    ([0, 15157], [5, -1]),
    ([0, 15157, 19448], [1, 1, 1]),
    ([0, 15157, 19448], [1e-2, 1, 5]),
])
def test_split_gof_meg(forward, idx, weights):
    """Test GOF splitting on MEG data."""
    gain = forward['sol']['data'][:, idx]
    # close to orthogonal
    norms = np.linalg.norm(gain, axis=0)
    triu = np.triu_indices(len(idx), 1)
    prods = np.abs(np.dot(gain.T, gain) / np.outer(norms, norms))[triu]
    assert_array_less(prods, 5e-3)  # approximately orthogonal
    # first, split across time (one dipole per time point)
    M = gain * weights
    gof_split = _split_gof(M, np.diag(weights), gain)
    assert_allclose(gof_split.sum(0), 100., atol=1e-5)  # all sum to 100
    assert_allclose(gof_split, 100 * np.eye(len(weights)), atol=1)  # loc
    # next, summed to a single time point (all dipoles active at one time pt)
    weights = np.array(weights)[:, np.newaxis]
    x = gain @ weights
    assert x.shape == (gain.shape[0], 1)
    gof_split = _split_gof(x, weights, gain)
    want = (norms * weights.T).T ** 2
    want = 100 * want / want.sum()
    assert_allclose(gof_split, want, atol=1e-3, rtol=1e-2)
    assert_allclose(gof_split.sum(), 100, rtol=1e-5)


@pytest.mark.parametrize('n_sensors, n_dipoles, n_times', [
    (10, 15, 7),
    (20, 60, 20),
])
@pytest.mark.parametrize('nnz', [2, 4])
@pytest.mark.parametrize('corr', [0.75])
@pytest.mark.parametrize('n_orient', [1, 3])
def test_mxne_inverse_sure_synthetic(n_sensors, n_dipoles, n_times, nnz, corr,
                                     n_orient, snr=4):
    """Tests SURE criterion for automatic alpha selection on synthetic data."""
    rng = np.random.RandomState(0)
    sigma = np.sqrt(1 - corr ** 2)
    U = rng.randn(n_sensors)
    # generate gain matrix
    G = np.empty([n_sensors, n_dipoles], order='F')
    G[:, :n_orient] = np.expand_dims(U, axis=-1)
    n_dip_per_pos = n_dipoles // n_orient
    for j in range(1, n_dip_per_pos):
        U *= corr
        U += sigma * rng.randn(n_sensors)
        G[:, j * n_orient:(j + 1) * n_orient] = np.expand_dims(U, axis=-1)
    # generate coefficient matrix
    support = rng.choice(n_dip_per_pos, nnz, replace=False)
    X = np.zeros((n_dipoles, n_times))
    for k in support:
        X[k * n_orient:(k + 1) * n_orient, :] = rng.normal(
            size=(n_orient, n_times))
    # generate measurement matrix
    M = G @ X
    noise = rng.randn(n_sensors, n_times)
    sigma = 1 / np.linalg.norm(noise) * np.linalg.norm(M) / snr
    M += sigma * noise
    # inverse modeling with sure
    alpha_max = norm_l2inf(np.dot(G.T, M), n_orient, copy=False)
    alpha_grid = np.geomspace(alpha_max, alpha_max / 10, num=15)
    _, active_set, _ = _compute_mxne_sure(M, G, alpha_grid, sigma=sigma,
                                          n_mxne_iter=5, maxit=3000, tol=1e-4,
                                          n_orient=n_orient,
                                          active_set_size=10, debias=True,
                                          solver="auto", dgap_freq=10,
                                          random_state=0, verbose=False)
    assert np.count_nonzero(active_set, axis=-1) == n_orient * nnz


@pytest.mark.slowtest  # slow on Azure
@testing.requires_testing_data
def test_mxne_inverse_sure():
    """Tests SURE criterion for automatic alpha selection on MEG data."""
    def data_fun(times):
        data = np.zeros(times.shape)
        data[times >= 0] = 50e-9
        return data
    n_dipoles = 2
    raw = mne.io.read_raw_fif(fname_raw)
    info = mne.io.read_info(fname_data)
    with info._unlock():
        info['projs'] = []
    noise_cov = mne.make_ad_hoc_cov(info)
    label_names = ['Aud-lh', 'Aud-rh']
    labels = [
        mne.read_label(data_path / 'MEG' / 'sample' / 'labels' / f'{ln}.label')
        for ln in label_names]
    fname_fwd = op.join(data_path, 'MEG', 'sample',
                        'sample_audvis_trunc-meg-eeg-oct-4-fwd.fif')
    forward = mne.read_forward_solution(fname_fwd)
    forward = mne.pick_types_forward(forward, meg="grad", eeg=False,
                                     exclude=raw.info['bads'])
    times = np.arange(100, dtype=np.float64) / raw.info['sfreq'] - 0.1
    stc = simulate_sparse_stc(forward['src'], n_dipoles=n_dipoles, times=times,
                              random_state=1, labels=labels, data_fun=data_fun)
    nave = 30
    evoked = simulate_evoked(forward, stc, info, noise_cov, nave=nave,
                             use_cps=False, iir_filter=None)
    evoked = evoked.crop(tmin=0, tmax=10e-3)
    stc_ = mixed_norm(evoked, forward, noise_cov, loose=0.9, n_mxne_iter=5,
                      depth=0.9)
    assert_array_equal(stc_.vertices, stc.vertices)


@pytest.mark.slowtest  # slow on Azure
@testing.requires_testing_data
def test_mxne_inverse_empty():
    """Tests solver with too high alpha."""
    evoked = read_evokeds(fname_data, condition=0, baseline=(None, 0))
    evoked.pick("grad", exclude="bads")
    fname_fwd = op.join(data_path, 'MEG', 'sample',
                        'sample_audvis_trunc-meg-eeg-oct-4-fwd.fif')
    forward = mne.read_forward_solution(fname_fwd)
    forward = mne.pick_types_forward(forward, meg="grad", eeg=False,
                                     exclude=evoked.info['bads'])
    cov = read_cov(fname_cov)
    with pytest.warns(RuntimeWarning, match='too big'):
        stc, residual = mixed_norm(
            evoked, forward, cov, n_mxne_iter=3, alpha=99,
            return_residual=True)
        assert stc.data.size == 0
        assert stc.vertices[0].size == 0
        assert stc.vertices[1].size == 0
        assert_allclose(evoked.data, residual.data)