# Authors: The MNE-Python contributors.
# License: BSD-3-Clause
# Copyright the MNE-Python contributors.

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

from mne.inverse_sparse.mxne_optim import (
    _Phi,
    _PhiT,
    dgap_l21l1,
    iterative_mixed_norm_solver,
    iterative_tf_mixed_norm_solver,
    mixed_norm_solver,
    norm_epsilon,
    norm_epsilon_inf,
    tf_mixed_norm_solver,
)
from mne.time_frequency._stft import stft_norm2
from mne.utils import _record_warnings, catch_logging


def _generate_tf_data():
    n, p, t = 30, 40, 64
    rng = np.random.RandomState(0)
    G = rng.randn(n, p)
    G /= np.std(G, axis=0)[None, :]
    X = np.zeros((p, t))
    active_set = [0, 4]
    times = np.linspace(0, 2 * np.pi, t)
    X[0] = np.sin(times)
    X[4] = -2 * np.sin(4 * times)
    X[4, times <= np.pi / 2] = 0
    X[4, times >= np.pi] = 0
    M = np.dot(G, X)
    M += 1 * rng.randn(*M.shape)
    return M, G, active_set


def test_l21_mxne():
    """Test convergence of MxNE solver."""
    n, p, t, alpha = 30, 40, 20, 1.0
    rng = np.random.RandomState(0)
    G = rng.randn(n, p)
    G /= np.std(G, axis=0)[None, :]
    X = np.zeros((p, t))
    X[0] = 3
    X[4] = -2
    M = np.dot(G, X)

    args = (M, G, alpha, 1000, 1e-8)
    with _record_warnings():  # CD
        X_hat_cd, active_set, _, gap_cd = mixed_norm_solver(
            *args, active_set_size=None, debias=True, solver="cd", return_gap=True
        )
    assert_array_less(gap_cd, 1e-8)
    assert_array_equal(np.where(active_set)[0], [0, 4])
    with _record_warnings():  # CD
        X_hat_bcd, active_set, E, gap_bcd = mixed_norm_solver(
            M,
            G,
            alpha,
            maxit=1000,
            tol=1e-8,
            active_set_size=None,
            debias=True,
            solver="bcd",
            return_gap=True,
        )
    assert_array_less(gap_bcd, 9.6e-9)
    assert_array_equal(np.where(active_set)[0], [0, 4])
    assert_allclose(X_hat_bcd, X_hat_cd, rtol=1e-2)

    with _record_warnings():  # CD
        X_hat_cd, active_set, _ = mixed_norm_solver(
            *args, active_set_size=2, debias=True, solver="cd"
        )
    assert_array_equal(np.where(active_set)[0], [0, 4])
    with _record_warnings():  # CD
        X_hat_bcd, active_set, _ = mixed_norm_solver(
            *args, active_set_size=2, debias=True, solver="bcd"
        )
    assert_array_equal(np.where(active_set)[0], [0, 4])
    assert_allclose(X_hat_bcd, X_hat_cd, rtol=1e-2)

    with _record_warnings():  # CD
        X_hat_bcd, active_set, _ = mixed_norm_solver(
            *args, active_set_size=2, debias=True, n_orient=2, solver="bcd"
        )
    assert_array_equal(np.where(active_set)[0], [0, 1, 4, 5])

    # suppress a coordinate-descent warning here
    with pytest.warns(RuntimeWarning, match="descent"):
        X_hat_cd, active_set, _ = mixed_norm_solver(
            *args, active_set_size=2, debias=True, n_orient=2, solver="cd"
        )
    assert_array_equal(np.where(active_set)[0], [0, 1, 4, 5])
    assert_allclose(X_hat_bcd, X_hat_cd, rtol=1e-2)

    with _record_warnings():  # CD
        X_hat_bcd, active_set, _ = mixed_norm_solver(
            *args, active_set_size=2, debias=True, n_orient=5, solver="bcd"
        )
    assert_array_equal(np.where(active_set)[0], [0, 1, 2, 3, 4])
    with pytest.warns(RuntimeWarning, match="descent"):
        X_hat_cd, active_set, _ = mixed_norm_solver(
            *args, active_set_size=2, debias=True, n_orient=5, solver="cd"
        )

    assert_array_equal(np.where(active_set)[0], [0, 1, 2, 3, 4])
    assert_allclose(X_hat_bcd, X_hat_cd)


@pytest.mark.slowtest
def test_non_convergence():
    """Test non-convergence of MxNE solver to catch unexpected bugs."""
    n, p, t, alpha = 30, 40, 20, 1.0
    rng = np.random.RandomState(0)
    G = rng.randn(n, p)
    G /= np.std(G, axis=0)[None, :]
    X = np.zeros((p, t))
    X[0] = 3
    X[4] = -2
    M = np.dot(G, X)

    # Impossible to converge with only 1 iteration and tol 1e-12
    # In case of non-convegence, we test that no error is returned.
    args = (M, G, alpha, 1, 1e-12)
    with catch_logging() as log:
        mixed_norm_solver(
            *args, active_set_size=None, debias=True, solver="bcd", verbose=True
        )
    log = log.getvalue()
    assert "Convergence reached" not in log


def test_tf_mxne():
    """Test convergence of TF-MxNE solver."""
    alpha_space = 10.0
    alpha_time = 5.0

    M, G, active_set = _generate_tf_data()

    with _record_warnings():  # CD
        X_hat_tf, active_set_hat_tf, E, gap_tfmxne = tf_mixed_norm_solver(
            M,
            G,
            alpha_space,
            alpha_time,
            maxit=200,
            tol=1e-8,
            verbose=True,
            n_orient=1,
            tstep=4,
            wsize=32,
            return_gap=True,
        )
    assert_array_less(gap_tfmxne, 1e-8)
    assert_array_equal(np.where(active_set_hat_tf)[0], active_set)


def test_norm_epsilon():
    """Test computation of espilon norm on TF coefficients."""
    tstep = np.array([2])
    wsize = np.array([4])
    n_times = 10
    n_steps = np.ceil(n_times / tstep.astype(float)).astype(int)
    n_freqs = wsize // 2 + 1
    n_coefs = n_steps * n_freqs
    phi = _Phi(wsize, tstep, n_coefs, n_times)

    Y = np.zeros((n_steps * n_freqs).item())
    l1_ratio = 0.03
    assert_allclose(norm_epsilon(Y, l1_ratio, phi), 0.0)

    Y[0] = 2.0
    assert_allclose(norm_epsilon(Y, l1_ratio, phi), np.max(Y))

    l1_ratio = 1.0
    assert_allclose(norm_epsilon(Y, l1_ratio, phi), np.max(Y))
    # dummy value without random:
    Y = np.arange((n_steps * n_freqs).item())
    l1_ratio = 0.0
    assert_allclose(
        norm_epsilon(Y, l1_ratio, phi) ** 2,
        stft_norm2(Y.reshape(-1, n_freqs[0], n_steps[0])),
    )

    l1_ratio = 0.03
    # test that vanilla epsilon norm = weights equal to 1
    w_time = np.ones(n_coefs[0])
    Y = np.abs(np.random.randn(n_coefs[0]))
    assert_allclose(
        norm_epsilon(Y, l1_ratio, phi), norm_epsilon(Y, l1_ratio, phi, w_time=w_time)
    )

    # scaling w_time and w_space by the same amount should divide
    # epsilon norm by the same amount
    Y = np.arange(n_coefs.item()) + 1
    mult = 2.0
    assert_allclose(
        norm_epsilon(Y, l1_ratio, phi, w_space=1, w_time=np.ones(n_coefs.item()))
        / mult,
        norm_epsilon(
            Y, l1_ratio, phi, w_space=mult, w_time=mult * np.ones(n_coefs.item())
        ),
    )


@pytest.mark.slowtest  # slow-ish on Travis OSX
@pytest.mark.timeout(60)  # ~30 s on Travis OSX and Linux OpenBLAS
def test_dgapl21l1():
    """Test duality gap for L21 + L1 regularization."""
    n_orient = 2
    M, G, active_set = _generate_tf_data()
    n_times = M.shape[1]
    n_sources = G.shape[1]
    tstep, wsize = np.array([4, 2]), np.array([64, 16])
    n_steps = np.ceil(n_times / tstep.astype(float)).astype(int)
    n_freqs = wsize // 2 + 1
    n_coefs = n_steps * n_freqs
    phi = _Phi(wsize, tstep, n_coefs, n_times)
    phiT = _PhiT(tstep, n_freqs, n_steps, n_times)

    for l1_ratio in [0.05, 0.1]:
        alpha_max = norm_epsilon_inf(G, M, phi, l1_ratio, n_orient)
        alpha_space = (1.0 - l1_ratio) * alpha_max
        alpha_time = l1_ratio * alpha_max

        Z = np.zeros([n_sources, phi.n_coefs.sum()])
        # for alpha = alpha_max, Z = 0 is the solution so the dgap is 0
        gap = dgap_l21l1(
            M,
            G,
            Z,
            np.ones(n_sources, dtype=bool),
            alpha_space,
            alpha_time,
            phi,
            phiT,
            n_orient,
            -np.inf,
        )[0]

        assert_allclose(0.0, gap)
        # check that solution for alpha smaller than alpha_max is non 0:
        X_hat_tf, active_set_hat_tf, E, gap = tf_mixed_norm_solver(
            M,
            G,
            alpha_space / 1.01,
            alpha_time / 1.01,
            maxit=200,
            tol=1e-8,
            verbose=True,
            debias=False,
            n_orient=n_orient,
            tstep=tstep,
            wsize=wsize,
            return_gap=True,
        )
        # allow possible small numerical errors (negative gap)
        assert_array_less(-1e-10, gap)
        assert_array_less(gap, 1e-8)
        assert_array_less(1, len(active_set_hat_tf))

        X_hat_tf, active_set_hat_tf, E, gap = tf_mixed_norm_solver(
            M,
            G,
            alpha_space / 5.0,
            alpha_time / 5.0,
            maxit=200,
            tol=1e-8,
            verbose=True,
            debias=False,
            n_orient=n_orient,
            tstep=tstep,
            wsize=wsize,
            return_gap=True,
        )
        assert_array_less(-1e-10, gap)
        assert_array_less(gap, 1e-8)
        assert_array_less(1, len(active_set_hat_tf))


def test_tf_mxne_vs_mxne():
    """Test equivalence of TF-MxNE (with alpha_time=0) and MxNE."""
    alpha_space = 60.0
    alpha_time = 0.0

    M, G, active_set = _generate_tf_data()

    X_hat_tf, active_set_hat_tf, E = tf_mixed_norm_solver(
        M,
        G,
        alpha_space,
        alpha_time,
        maxit=200,
        tol=1e-8,
        verbose=True,
        debias=False,
        n_orient=1,
        tstep=4,
        wsize=32,
    )

    # Also run L21 and check that we get the same
    X_hat_l21, _, _ = mixed_norm_solver(
        M,
        G,
        alpha_space,
        maxit=200,
        tol=1e-8,
        verbose=False,
        n_orient=1,
        active_set_size=None,
        debias=False,
    )

    assert_allclose(X_hat_tf, X_hat_l21, rtol=1e-1)


@pytest.mark.slowtest  # slow-ish on Travis OSX
def test_iterative_reweighted_mxne():
    """Test convergence of irMxNE solver."""
    n, p, t, alpha = 30, 40, 20, 1
    rng = np.random.RandomState(0)
    G = rng.randn(n, p)
    G /= np.std(G, axis=0)[None, :]
    X = np.zeros((p, t))
    X[0] = 3
    X[4] = -2
    M = np.dot(G, X)

    with _record_warnings():  # CD
        X_hat_l21, _, _ = mixed_norm_solver(
            M,
            G,
            alpha,
            maxit=1000,
            tol=1e-8,
            verbose=False,
            n_orient=1,
            active_set_size=None,
            debias=False,
            solver="bcd",
        )
    with _record_warnings():  # CD
        X_hat_bcd, active_set, _ = iterative_mixed_norm_solver(
            M,
            G,
            alpha,
            1,
            maxit=1000,
            tol=1e-8,
            active_set_size=None,
            debias=False,
            solver="bcd",
        )
    assert_allclose(X_hat_bcd, X_hat_l21, rtol=1e-3)

    with _record_warnings():  # CD
        X_hat_bcd, active_set, _ = iterative_mixed_norm_solver(
            M,
            G,
            alpha,
            5,
            maxit=1000,
            tol=1e-8,
            active_set_size=2,
            debias=True,
            solver="bcd",
        )
    assert_array_equal(np.where(active_set)[0], [0, 4])
    with _record_warnings():  # CD
        X_hat_cd, active_set, _ = iterative_mixed_norm_solver(
            M,
            G,
            alpha,
            5,
            maxit=1000,
            tol=1e-8,
            active_set_size=None,
            debias=True,
            solver="cd",
        )
    assert_array_equal(np.where(active_set)[0], [0, 4])
    assert_array_almost_equal(X_hat_bcd, X_hat_cd, 5)

    with _record_warnings():  # CD
        X_hat_bcd, active_set, _ = iterative_mixed_norm_solver(
            M,
            G,
            alpha,
            5,
            maxit=1000,
            tol=1e-8,
            active_set_size=2,
            debias=True,
            n_orient=2,
            solver="bcd",
        )
    assert_array_equal(np.where(active_set)[0], [0, 1, 4, 5])
    # suppress a coordinate-descent warning here
    with pytest.warns(RuntimeWarning, match="descent"):
        X_hat_cd, active_set, _ = iterative_mixed_norm_solver(
            M,
            G,
            alpha,
            5,
            maxit=1000,
            tol=1e-8,
            active_set_size=2,
            debias=True,
            n_orient=2,
            solver="cd",
        )
    assert_array_equal(np.where(active_set)[0], [0, 1, 4, 5])
    assert_allclose(X_hat_bcd, X_hat_cd)

    X_hat_bcd, active_set, _ = iterative_mixed_norm_solver(
        M, G, alpha, 5, maxit=1000, tol=1e-8, active_set_size=2, debias=True, n_orient=5
    )
    assert_array_equal(np.where(active_set)[0], [0, 1, 2, 3, 4])
    with pytest.warns(RuntimeWarning, match="descent"):
        X_hat_cd, active_set, _ = iterative_mixed_norm_solver(
            M,
            G,
            alpha,
            5,
            maxit=1000,
            tol=1e-8,
            active_set_size=2,
            debias=True,
            n_orient=5,
            solver="cd",
        )
    assert_array_equal(np.where(active_set)[0], [0, 1, 2, 3, 4])
    assert_allclose(X_hat_bcd, X_hat_cd)


@pytest.mark.slowtest
def test_iterative_reweighted_tfmxne():
    """Test convergence of irTF-MxNE solver."""
    M, G, true_active_set = _generate_tf_data()
    alpha_space = 38.0
    alpha_time = 0.5
    tstep, wsize = [4, 2], [64, 16]

    X_hat_tf, _, _ = tf_mixed_norm_solver(
        M,
        G,
        alpha_space,
        alpha_time,
        maxit=1000,
        tol=1e-4,
        wsize=wsize,
        tstep=tstep,
        verbose=False,
        n_orient=1,
        debias=False,
    )
    X_hat_bcd, active_set, _ = iterative_tf_mixed_norm_solver(
        M,
        G,
        alpha_space,
        alpha_time,
        1,
        wsize=wsize,
        tstep=tstep,
        maxit=1000,
        tol=1e-4,
        debias=False,
        verbose=False,
    )
    assert_allclose(X_hat_tf, X_hat_bcd, rtol=1e-3)
    assert_array_equal(np.where(active_set)[0], true_active_set)

    alpha_space = 50.0
    X_hat_bcd, active_set, _ = iterative_tf_mixed_norm_solver(
        M,
        G,
        alpha_space,
        alpha_time,
        3,
        wsize=wsize,
        tstep=tstep,
        n_orient=5,
        maxit=1000,
        tol=1e-4,
        debias=False,
        verbose=False,
    )
    assert_array_equal(np.where(active_set)[0], [0, 1, 2, 3, 4])

    alpha_space = 40.0
    X_hat_bcd, active_set, _ = iterative_tf_mixed_norm_solver(
        M,
        G,
        alpha_space,
        alpha_time,
        2,
        wsize=wsize,
        tstep=tstep,
        n_orient=2,
        maxit=1000,
        tol=1e-4,
        debias=False,
        verbose=False,
    )
    assert_array_equal(np.where(active_set)[0], [0, 1, 4, 5])