# Author: Alexandre Gramfort <alexandre.gramfort@inria.fr>
#         Daniel Strohmeier <daniel.strohmeier@gmail.com>
#         Mathurin Massias <mathurin.massias@gmail.com>
# License: Simplified BSD

import functools
from math import sqrt

import numpy as np

from .mxne_debiasing import compute_bias
from ..utils import (logger, verbose, sum_squared, warn, _get_blas_funcs,
                     _validate_type, _check_option)
from ..time_frequency._stft import stft_norm1, stft_norm2, stft, istft


@functools.lru_cache(None)
def _get_dgemm():
    return _get_blas_funcs(np.float64, 'gemm')


def groups_norm2(A, n_orient):
    """Compute squared L2 norms of groups inplace."""
    n_positions = A.shape[0] // n_orient
    return np.sum(np.power(A, 2, A).reshape(n_positions, -1), axis=1)


def norm_l2inf(A, n_orient, copy=True):
    """L2-inf norm."""
    if A.size == 0:
        return 0.0
    if copy:
        A = A.copy()
    return sqrt(np.max(groups_norm2(A, n_orient)))


def norm_l21(A, n_orient, copy=True):
    """L21 norm."""
    if A.size == 0:
        return 0.0
    if copy:
        A = A.copy()
    return np.sum(np.sqrt(groups_norm2(A, n_orient)))


def _primal_l21(M, G, X, active_set, alpha, n_orient):
    """Primal objective for the mixed-norm inverse problem.

    See :footcite:`GramfortEtAl2012`.

    Parameters
    ----------
    M : array, shape (n_sensors, n_times)
        The data.
    G : array, shape (n_sensors, n_active)
        The gain matrix a.k.a. lead field.
    X : array, shape (n_active, n_times)
        Sources.
    active_set : array of bool, shape (n_sources,)
        Mask of active sources.
    alpha : float
        The regularization parameter.
    n_orient : int
        Number of dipoles per locations (typically 1 or 3).

    Returns
    -------
    p_obj : float
        Primal objective.
    R : array, shape (n_sensors, n_times)
        Current residual (M - G * X).
    nR2 : float
        Data-fitting term.
    GX : array, shape (n_sensors, n_times)
        Forward prediction.
    """
    GX = np.dot(G[:, active_set], X)
    R = M - GX
    penalty = norm_l21(X, n_orient, copy=True)
    nR2 = sum_squared(R)
    p_obj = 0.5 * nR2 + alpha * penalty
    return p_obj, R, nR2, GX


def dgap_l21(M, G, X, active_set, alpha, n_orient):
    """Duality gap for the mixed norm inverse problem.

    See :footcite:`GramfortEtAl2012`.

    Parameters
    ----------
    M : array, shape (n_sensors, n_times)
        The data.
    G : array, shape (n_sensors, n_active)
        The gain matrix a.k.a. lead field.
    X : array, shape (n_active, n_times)
        Sources.
    active_set : array of bool, shape (n_sources, )
        Mask of active sources.
    alpha : float
        The regularization parameter.
    n_orient : int
        Number of dipoles per locations (typically 1 or 3).

    Returns
    -------
    gap : float
        Dual gap.
    p_obj : float
        Primal objective.
    d_obj : float
        Dual objective. gap = p_obj - d_obj.
    R : array, shape (n_sensors, n_times)
        Current residual (M - G * X).

    References
    ----------
    .. footbibilography::
    """
    p_obj, R, nR2, GX = _primal_l21(M, G, X, active_set, alpha, n_orient)
    dual_norm = norm_l2inf(np.dot(G.T, R), n_orient, copy=False)
    scaling = alpha / dual_norm
    scaling = min(scaling, 1.0)
    d_obj = (scaling - 0.5 * (scaling ** 2)) * nR2 + scaling * np.sum(R * GX)

    gap = p_obj - d_obj
    return gap, p_obj, d_obj, R


@verbose
def _mixed_norm_solver_cd(M, G, alpha, lipschitz_constant, maxit=10000,
                          tol=1e-8, verbose=None, init=None, n_orient=1,
                          dgap_freq=10):
    """Solve L21 inverse problem with coordinate descent."""
    from sklearn.linear_model import MultiTaskLasso

    assert M.ndim == G.ndim and M.shape[0] == G.shape[0]

    clf = MultiTaskLasso(alpha=alpha / len(M), tol=tol / sum_squared(M),
                         fit_intercept=False, max_iter=maxit, warm_start=True)
    if init is not None:
        clf.coef_ = init.T
    else:
        clf.coef_ = np.zeros((G.shape[1], M.shape[1])).T
    clf.fit(G, M)

    X = clf.coef_.T
    active_set = np.any(X, axis=1)
    X = X[active_set]
    gap, p_obj, d_obj, _ = dgap_l21(M, G, X, active_set, alpha, n_orient)
    return X, active_set, p_obj


@verbose
def _mixed_norm_solver_bcd(M, G, alpha, lipschitz_constant, maxit=200,
                           tol=1e-8, verbose=None, init=None, n_orient=1,
                           dgap_freq=10, use_accel=True, K=5):
    """Solve L21 inverse problem with block coordinate descent."""
    _, n_times = M.shape
    _, n_sources = G.shape
    n_positions = n_sources // n_orient

    if init is None:
        X = np.zeros((n_sources, n_times))
        R = M.copy()
    else:
        X = init
        R = M - np.dot(G, X)

    E = []  # track primal objective function
    highest_d_obj = - np.inf
    active_set = np.zeros(n_sources, dtype=bool)  # start with full AS

    alpha_lc = alpha / lipschitz_constant

    if use_accel:
        last_K_X = np.empty((K + 1, n_sources, n_times))
        U = np.zeros((K, n_sources * n_times))

    # First make G fortran for faster access to blocks of columns
    G = np.asfortranarray(G)
    # Ensure these are correct for dgemm
    assert R.dtype == np.float64
    assert G.dtype == np.float64
    one_ovr_lc = 1. / lipschitz_constant

    # assert that all the multiplied matrices are fortran contiguous
    assert X.T.flags.f_contiguous
    assert R.T.flags.f_contiguous
    assert G.flags.f_contiguous
    # storing list of contiguous arrays
    list_G_j_c = []
    for j in range(n_positions):
        idx = slice(j * n_orient, (j + 1) * n_orient)
        list_G_j_c.append(np.ascontiguousarray(G[:, idx]))

    for i in range(maxit):
        _bcd(G, X, R, active_set, one_ovr_lc, n_orient, alpha_lc, list_G_j_c)

        if (i + 1) % dgap_freq == 0:
            _, p_obj, d_obj, _ = dgap_l21(M, G, X[active_set], active_set,
                                          alpha, n_orient)
            highest_d_obj = max(d_obj, highest_d_obj)
            gap = p_obj - highest_d_obj
            E.append(p_obj)
            logger.debug("Iteration %d :: p_obj %f :: dgap %f :: n_active %d" %
                         (i + 1, p_obj, gap, np.sum(active_set) / n_orient))

            if gap < tol:
                logger.debug('Convergence reached ! (gap: %s < %s)'
                             % (gap, tol))
                break

        # using Anderson acceleration of the primal variable for faster
        # convergence
        if use_accel:
            last_K_X[i % (K + 1)] = X

            if i % (K + 1) == K:
                for k in range(K):
                    U[k] = last_K_X[k + 1].ravel() - last_K_X[k].ravel()
                C = U @ U.T
                # at least on ARM64 we can't rely on np.linalg.solve to
                # reliably raise LinAlgError here, so use SVD instead
                # equivalent to:
                # z = np.linalg.solve(C, np.ones(K))
                u, s, _ = np.linalg.svd(C, hermitian=True)
                if s[-1] <= 1e-6 * s[0] or not np.isfinite(s).all():
                    logger.debug("Iteration %d: LinAlg Error" % (i + 1))
                    continue
                z = ((u * 1 / s) @ u.T).sum(0)
                c = z / z.sum()
                X_acc = np.sum(
                    last_K_X[:-1] * c[:, None, None], axis=0
                )
                _grp_norm2_acc = groups_norm2(X_acc, n_orient)
                active_set_acc = _grp_norm2_acc != 0
                if n_orient > 1:
                    active_set_acc = np.kron(
                        active_set_acc, np.ones(n_orient, dtype=bool)
                    )
                p_obj = _primal_l21(M, G, X[active_set], active_set, alpha,
                                    n_orient)[0]
                p_obj_acc = _primal_l21(M, G, X_acc[active_set_acc],
                                        active_set_acc, alpha, n_orient)[0]
                if p_obj_acc < p_obj:
                    X = X_acc
                    active_set = active_set_acc
                    R = M - G[:, active_set] @ X[active_set]

    X = X[active_set]

    return X, active_set, E


def _bcd(G, X, R, active_set, one_ovr_lc, n_orient, alpha_lc, list_G_j_c):
    """Implement one full pass of BCD.

    BCD stands for Block Coordinate Descent.
    This function make use of scipy.linalg.get_blas_funcs to speed reasons.

    Parameters
    ----------
    G : array, shape (n_sensors, n_active)
        The gain matrix a.k.a. lead field.
    X : array, shape (n_sources, n_times)
        Sources, modified in place.
    R : array, shape (n_sensors, n_times)
        The residuals: R = M - G @ X, modified in place.
    active_set : array of bool, shape (n_sources, )
        Mask of active sources, modified in place.
    one_ovr_lc : array, shape (n_positions, )
        One over the lipschitz constants.
    n_orient : int
        Number of dipoles per positions (typically 1 or 3).
    n_positions : int
        Number of source positions.
    alpha_lc: array, shape (n_positions, )
        alpha * (Lipschitz constants).
    """
    X_j_new = np.zeros_like(X[:n_orient, :], order='C')
    dgemm = _get_dgemm()

    for j, G_j_c in enumerate(list_G_j_c):
        idx = slice(j * n_orient, (j + 1) * n_orient)
        G_j = G[:, idx]
        X_j = X[idx]
        dgemm(alpha=one_ovr_lc[j], beta=0., a=R.T, b=G_j, c=X_j_new.T,
              overwrite_c=True)
        # X_j_new = G_j.T @ R
        # Mathurin's trick to avoid checking all the entries
        was_non_zero = X_j[0, 0] != 0
        # was_non_zero = np.any(X_j)
        if was_non_zero:
            dgemm(alpha=1., beta=1., a=X_j.T, b=G_j_c.T, c=R.T,
                  overwrite_c=True)
            # R += np.dot(G_j, X_j)
            X_j_new += X_j
        block_norm = sqrt(sum_squared(X_j_new))
        if block_norm <= alpha_lc[j]:
            X_j.fill(0.)
            active_set[idx] = False
        else:
            shrink = max(1.0 - alpha_lc[j] / block_norm, 0.0)
            X_j_new *= shrink
            dgemm(alpha=-1., beta=1., a=X_j_new.T, b=G_j_c.T, c=R.T,
                  overwrite_c=True)
            # R -= np.dot(G_j, X_j_new)
            X_j[:] = X_j_new
            active_set[idx] = True


@verbose
def mixed_norm_solver(M, G, alpha, maxit=3000, tol=1e-8, verbose=None,
                      active_set_size=50, debias=True, n_orient=1,
                      solver='auto', return_gap=False, dgap_freq=10,
                      active_set_init=None, X_init=None):
    """Solve L1/L2 mixed-norm inverse problem with active set strategy.

    See references :footcite:`GramfortEtAl2012,StrohmeierEtAl2016,
    BertrandEtAl2020`.

    Parameters
    ----------
    M : array, shape (n_sensors, n_times)
        The data.
    G : array, shape (n_sensors, n_dipoles)
        The gain matrix a.k.a. lead field.
    alpha : float
        The regularization parameter. It should be between 0 and 100.
        A value of 100 will lead to an empty active set (no active source).
    maxit : int
        The number of iterations.
    tol : float
        Tolerance on dual gap for convergence checking.
    %(verbose)s
    active_set_size : int
        Size of active set increase at each iteration.
    debias : bool
        Debias source estimates.
    n_orient : int
        The number of orientation (1 : fixed or 3 : free or loose).
    solver : 'cd' | 'bcd' | 'auto'
        The algorithm to use for the optimization. Block Coordinate Descent
        (BCD) uses Anderson acceleration for faster convergence.
    return_gap : bool
        Return final duality gap.
    dgap_freq : int
        The duality gap is computed every dgap_freq iterations of the solver on
        the active set.
    active_set_init : array, shape (n_dipoles,) or None
        The initial active set (boolean array) used at the first iteration.
        If None, the usual active set strategy is applied.
    X_init : array, shape (n_dipoles, n_times) or None
        The initial weight matrix used for warm starting the solver. If None,
        the weights are initialized at zero.

    Returns
    -------
    X : array, shape (n_active, n_times)
        The source estimates.
    active_set : array, shape (new_active_set_size,)
        The mask of active sources. Note that new_active_set_size is the size
        of the active set after convergence of the solver.
    E : list
        The value of the objective function over the iterations.
    gap : float
        Final duality gap. Returned only if return_gap is True.

    References
    ----------
    .. footbibliography::
    """
    n_dipoles = G.shape[1]
    n_positions = n_dipoles // n_orient
    _, n_times = M.shape
    alpha_max = norm_l2inf(np.dot(G.T, M), n_orient, copy=False)
    logger.info("-- ALPHA MAX : %s" % alpha_max)
    alpha = float(alpha)
    X = np.zeros((n_dipoles, n_times), dtype=G.dtype)

    has_sklearn = True
    try:
        from sklearn.linear_model import MultiTaskLasso  # noqa: F401
    except ImportError:
        has_sklearn = False

    _validate_type(solver, str, 'solver')
    _check_option('solver', solver, ('cd', 'bcd', 'auto'))
    if solver == 'auto':
        if has_sklearn and (n_orient == 1):
            solver = 'cd'
        else:
            solver = 'bcd'

    if solver == 'cd':
        if n_orient == 1 and not has_sklearn:
            warn('Scikit-learn >= 0.12 cannot be found. Using block coordinate'
                 ' descent instead of coordinate descent.')
            solver = 'bcd'
        if n_orient > 1:
            warn('Coordinate descent is only available for fixed orientation. '
                 'Using block coordinate descent instead of coordinate '
                 'descent')
            solver = 'bcd'

    if solver == 'cd':
        logger.info("Using coordinate descent")
        l21_solver = _mixed_norm_solver_cd
        lc = None
    else:
        assert solver == 'bcd'
        logger.info("Using block coordinate descent")
        l21_solver = _mixed_norm_solver_bcd
        G = np.asfortranarray(G)
        if n_orient == 1:
            lc = np.sum(G * G, axis=0)
        else:
            lc = np.empty(n_positions)
            for j in range(n_positions):
                G_tmp = G[:, (j * n_orient):((j + 1) * n_orient)]
                lc[j] = np.linalg.norm(np.dot(G_tmp.T, G_tmp), ord=2)

    if active_set_size is not None:
        E = list()
        highest_d_obj = - np.inf
        if X_init is not None and X_init.shape != (n_dipoles, n_times):
            raise ValueError('Wrong dim for initialized coefficients.')
        active_set = (active_set_init if active_set_init is not None else
                      np.zeros(n_dipoles, dtype=bool))
        idx_large_corr = np.argsort(groups_norm2(np.dot(G.T, M), n_orient))
        new_active_idx = idx_large_corr[-active_set_size:]
        if n_orient > 1:
            new_active_idx = (n_orient * new_active_idx[:, None] +
                              np.arange(n_orient)[None, :]).ravel()
        active_set[new_active_idx] = True
        as_size = np.sum(active_set)
        gap = np.inf
        for k in range(maxit):
            if solver == 'bcd':
                lc_tmp = lc[active_set[::n_orient]]
            elif solver == 'cd':
                lc_tmp = None
            else:
                lc_tmp = 1.01 * np.linalg.norm(G[:, active_set], ord=2) ** 2
            X, as_, _ = l21_solver(M, G[:, active_set], alpha, lc_tmp,
                                   maxit=maxit, tol=tol, init=X_init,
                                   n_orient=n_orient, dgap_freq=dgap_freq)
            active_set[active_set] = as_.copy()
            idx_old_active_set = np.where(active_set)[0]

            _, p_obj, d_obj, R = dgap_l21(M, G, X, active_set, alpha,
                                          n_orient)
            highest_d_obj = max(d_obj, highest_d_obj)
            gap = p_obj - highest_d_obj
            E.append(p_obj)
            logger.info("Iteration %d :: p_obj %f :: dgap %f :: "
                        "n_active_start %d :: n_active_end %d" % (
                            k + 1, p_obj, gap, as_size // n_orient,
                            np.sum(active_set) // n_orient))
            if gap < tol:
                logger.info('Convergence reached ! (gap: %s < %s)'
                            % (gap, tol))
                break

            # add sources if not last iteration
            if k < (maxit - 1):
                idx_large_corr = np.argsort(groups_norm2(np.dot(G.T, R),
                                                         n_orient))
                new_active_idx = idx_large_corr[-active_set_size:]
                if n_orient > 1:
                    new_active_idx = (n_orient * new_active_idx[:, None] +
                                      np.arange(n_orient)[None, :])
                    new_active_idx = new_active_idx.ravel()
                active_set[new_active_idx] = True
                idx_active_set = np.where(active_set)[0]
                as_size = np.sum(active_set)
                X_init = np.zeros((as_size, n_times), dtype=X.dtype)
                idx = np.searchsorted(idx_active_set, idx_old_active_set)
                X_init[idx] = X
        else:
            warn('Did NOT converge ! (gap: %s > %s)' % (gap, tol))
    else:
        X, active_set, E = l21_solver(M, G, alpha, lc, maxit=maxit,
                                      tol=tol, n_orient=n_orient, init=None)
        if return_gap:
            gap = dgap_l21(M, G, X, active_set, alpha, n_orient)[0]

    if np.any(active_set) and debias:
        bias = compute_bias(M, G[:, active_set], X, n_orient=n_orient)
        X *= bias[:, np.newaxis]

    logger.info('Final active set size: %s' % (np.sum(active_set) // n_orient))

    if return_gap:
        return X, active_set, E, gap
    else:
        return X, active_set, E


@verbose
def iterative_mixed_norm_solver(M, G, alpha, n_mxne_iter, maxit=3000,
                                tol=1e-8, verbose=None, active_set_size=50,
                                debias=True, n_orient=1, dgap_freq=10,
                                solver='auto', weight_init=None):
    """Solve L0.5/L2 mixed-norm inverse problem with active set strategy.

    See reference :footcite:`StrohmeierEtAl2016`.

    Parameters
    ----------
    M : array, shape (n_sensors, n_times)
        The data.
    G : array, shape (n_sensors, n_dipoles)
        The gain matrix a.k.a. lead field.
    alpha : float
        The regularization parameter. It should be between 0 and 100.
        A value of 100 will lead to an empty active set (no active source).
    n_mxne_iter : int
        The number of MxNE iterations. If > 1, iterative reweighting
        is applied.
    maxit : int
        The number of iterations.
    tol : float
        Tolerance on dual gap for convergence checking.
    %(verbose)s
    active_set_size : int
        Size of active set increase at each iteration.
    debias : bool
        Debias source estimates.
    n_orient : int
        The number of orientation (1 : fixed or 3 : free or loose).
    dgap_freq : int or np.inf
        The duality gap is evaluated every dgap_freq iterations.
    solver : 'cd' | 'bcd' | 'auto'
        The algorithm to use for the optimization.
    weight_init : array, shape (n_dipoles,) or None
        The initial weight used for reweighting the gain matrix. If None, the
        weights are initialized with ones.

    Returns
    -------
    X : array, shape (n_active, n_times)
        The source estimates.
    active_set : array
        The mask of active sources.
    E : list
        The value of the objective function over the iterations.

    References
    ----------
    .. footbibliography::
    """
    def g(w):
        return np.sqrt(np.sqrt(groups_norm2(w.copy(), n_orient)))

    def gprime(w):
        return 2. * np.repeat(g(w), n_orient).ravel()

    E = list()

    if weight_init is not None and weight_init.shape != (G.shape[1],):
        raise ValueError('Wrong dimension for weight initialization. Got %s. '
                         'Expected %s.' % (weight_init.shape, (G.shape[1],)))

    weights = weight_init if weight_init is not None else np.ones(G.shape[1])
    active_set = (weights != 0)
    weights = weights[active_set]
    X = np.zeros((G.shape[1], M.shape[1]))

    for k in range(n_mxne_iter):
        X0 = X.copy()
        active_set_0 = active_set.copy()
        G_tmp = G[:, active_set] * weights[np.newaxis, :]

        if active_set_size is not None:
            if np.sum(active_set) > (active_set_size * n_orient):
                X, _active_set, _ = mixed_norm_solver(
                    M, G_tmp, alpha, debias=False, n_orient=n_orient,
                    maxit=maxit, tol=tol, active_set_size=active_set_size,
                    dgap_freq=dgap_freq, solver=solver, verbose=verbose)
            else:
                X, _active_set, _ = mixed_norm_solver(
                    M, G_tmp, alpha, debias=False, n_orient=n_orient,
                    maxit=maxit, tol=tol, active_set_size=None,
                    dgap_freq=dgap_freq, solver=solver, verbose=verbose)
        else:
            X, _active_set, _ = mixed_norm_solver(
                M, G_tmp, alpha, debias=False, n_orient=n_orient,
                maxit=maxit, tol=tol, active_set_size=None,
                dgap_freq=dgap_freq, solver=solver, verbose=verbose)

        logger.info('active set size %d' % (_active_set.sum() / n_orient))

        if _active_set.sum() > 0:
            active_set[active_set] = _active_set
            # Reapply weights to have correct unit
            X *= weights[_active_set][:, np.newaxis]
            weights = gprime(X)
            p_obj = 0.5 * np.linalg.norm(M - np.dot(G[:, active_set], X),
                                         'fro') ** 2. + alpha * np.sum(g(X))
            E.append(p_obj)

            # Check convergence
            if ((k >= 1) and np.all(active_set == active_set_0) and
                    np.all(np.abs(X - X0) < tol)):
                print('Convergence reached after %d reweightings!' % k)
                break
        else:
            active_set = np.zeros_like(active_set)
            p_obj = 0.5 * np.linalg.norm(M) ** 2.
            E.append(p_obj)
            break

    if np.any(active_set) and debias:
        bias = compute_bias(M, G[:, active_set], X, n_orient=n_orient)
        X *= bias[:, np.newaxis]

    return X, active_set, E


###############################################################################
# TF-MxNE

@verbose
def tf_lipschitz_constant(M, G, phi, phiT, tol=1e-3, verbose=None):
    """Compute lipschitz constant for FISTA.

    It uses a power iteration method.
    """
    n_times = M.shape[1]
    n_points = G.shape[1]
    iv = np.ones((n_points, n_times), dtype=np.float64)
    v = phi(iv)
    L = 1e100
    for it in range(100):
        L_old = L
        logger.info('Lipschitz estimation: iteration = %d' % it)
        iv = np.real(phiT(v))
        Gv = np.dot(G, iv)
        GtGv = np.dot(G.T, Gv)
        w = phi(GtGv)
        L = np.max(np.abs(w))  # l_inf norm
        v = w / L
        if abs((L - L_old) / L_old) < tol:
            break
    return L


def safe_max_abs(A, ia):
    """Compute np.max(np.abs(A[ia])) possible with empty A."""
    if np.sum(ia):  # ia is not empty
        return np.max(np.abs(A[ia]))
    else:
        return 0.


def safe_max_abs_diff(A, ia, B, ib):
    """Compute np.max(np.abs(A)) possible with empty A."""
    A = A[ia] if np.sum(ia) else 0.0
    B = B[ib] if np.sum(ia) else 0.0
    return np.max(np.abs(A - B))


class _Phi(object):
    """Have phi stft as callable w/o using a lambda that does not pickle."""

    def __init__(self, wsize, tstep, n_coefs, n_times):  # noqa: D102
        self.wsize = np.atleast_1d(wsize)
        self.tstep = np.atleast_1d(tstep)
        self.n_coefs = np.atleast_1d(n_coefs)
        self.n_dicts = len(tstep)
        self.n_freqs = wsize // 2 + 1
        self.n_steps = self.n_coefs // self.n_freqs
        self.n_times = n_times
        # ravel freq+time here
        self.ops = list()
        for ws, ts in zip(self.wsize, self.tstep):
            self.ops.append(
                stft(np.eye(n_times), ws, ts,
                     verbose=False).reshape(n_times, -1))

    def __call__(self, x):  # noqa: D105
        if self.n_dicts == 1:
            return x @ self.ops[0]
        else:
            return np.hstack(
                [x @ op for op in self.ops]) / np.sqrt(self.n_dicts)

    def norm(self, z, ord=2):
        """Squared L2 norm if ord == 2 and L1 norm if order == 1."""
        if ord not in (1, 2):
            raise ValueError('Only supported norm order are 1 and 2. '
                             'Got ord = %s' % ord)
        stft_norm = stft_norm1 if ord == 1 else stft_norm2
        norm = 0.
        if len(self.n_coefs) > 1:
            z_ = np.array_split(np.atleast_2d(z), np.cumsum(self.n_coefs)[:-1],
                                axis=1)
        else:
            z_ = [np.atleast_2d(z)]
        for i in range(len(z_)):
            norm += stft_norm(
                z_[i].reshape(-1, self.n_freqs[i], self.n_steps[i]))
        return norm


class _PhiT(object):
    """Have phi.T istft as callable w/o using a lambda that does not pickle."""

    def __init__(self, tstep, n_freqs, n_steps, n_times):  # noqa: D102
        self.tstep = tstep
        self.n_freqs = n_freqs
        self.n_steps = n_steps
        self.n_times = n_times
        self.n_dicts = len(tstep) if isinstance(tstep, np.ndarray) else 1
        self.n_coefs = list()
        self.op_re = list()
        self.op_im = list()
        for nf, ns, ts in zip(self.n_freqs, self.n_steps, self.tstep):
            nc = nf * ns
            self.n_coefs.append(nc)
            eye = np.eye(nc).reshape(nf, ns, nf, ns)
            self.op_re.append(istft(
                eye, ts, n_times).reshape(nc, n_times))
            self.op_im.append(istft(
                eye * 1j, ts, n_times).reshape(nc, n_times))

    def __call__(self, z):  # noqa: D105
        if self.n_dicts == 1:
            return z.real @ self.op_re[0] + z.imag @ self.op_im[0]
        else:
            x_out = np.zeros((z.shape[0], self.n_times))
            z_ = np.array_split(z, np.cumsum(self.n_coefs)[:-1], axis=1)
            for this_z, op_re, op_im in zip(z_, self.op_re, self.op_im):
                x_out += this_z.real @ op_re + this_z.imag @ op_im
            return x_out / np.sqrt(self.n_dicts)


def norm_l21_tf(Z, phi, n_orient, w_space=None):
    """L21 norm for TF."""
    if Z.shape[0]:
        l21_norm = np.sqrt(
            phi.norm(Z, ord=2).reshape(-1, n_orient).sum(axis=1))
        if w_space is not None:
            l21_norm *= w_space
        l21_norm = l21_norm.sum()
    else:
        l21_norm = 0.
    return l21_norm


def norm_l1_tf(Z, phi, n_orient, w_time):
    """L1 norm for TF."""
    if Z.shape[0]:
        n_positions = Z.shape[0] // n_orient
        Z_ = np.sqrt(np.sum(
            (np.abs(Z) ** 2.).reshape((n_orient, -1), order='F'), axis=0))
        Z_ = Z_.reshape((n_positions, -1), order='F')
        if w_time is not None:
            Z_ *= w_time
        l1_norm = phi.norm(Z_, ord=1).sum()
    else:
        l1_norm = 0.
    return l1_norm


def norm_epsilon(Y, l1_ratio, phi, w_space=1., w_time=None):
    """Weighted epsilon norm.

    The weighted epsilon norm is the dual norm of::

    w_{space} * (1. - l1_ratio) * ||Y||_2 + l1_ratio * ||Y||_{1, w_{time}}.

    where `||Y||_{1, w_{time}} = (np.abs(Y) * w_time).sum()`

    Warning: it takes into account the fact that Y only contains coefficients
    corresponding to the positive frequencies (see `stft_norm2()`): some
    entries will be counted twice. It is also assumed that all entries of both
    Y and w_time are non-negative. See
    :footcite:`NdiayeEtAl2016,BurdakovMerkulov2001`.

    Parameters
    ----------
    Y : array, shape (n_coefs,)
        The input data.
    l1_ratio : float between 0 and 1
        Tradeoff between L2 and L1 regularization. When it is 0, no temporal
        regularization is applied.
    phi : instance of _Phi
        The TF operator.
    w_space : float
        Scalar weight of the L2 norm. By default, it is taken equal to 1.
    w_time : array, shape (n_coefs, ) | None
        Weights of each TF coefficient in the L1 norm. If None, weights equal
        to 1 are used.


    Returns
    -------
    nu : float
        The value of the dual norm evaluated at Y.

    References
    ----------
    .. footbibliography::
    """
    # since the solution is invariant to flipped signs in Y, all entries
    # of Y are assumed positive

    # Add negative freqs: count all freqs twice except first and last:
    freqs_count = np.full(len(Y), 2)
    for i, fc in enumerate(np.array_split(freqs_count,
                                          np.cumsum(phi.n_coefs)[:-1])):
        fc[:phi.n_steps[i]] = 1
        fc[-phi.n_steps[i]:] = 1

    # exclude 0 weights:
    if w_time is not None:
        nonzero_weights = (w_time != 0.0)
        Y = Y[nonzero_weights]
        freqs_count = freqs_count[nonzero_weights]
        w_time = w_time[nonzero_weights]

    norm_inf_Y = np.max(Y / w_time) if w_time is not None else np.max(Y)
    if l1_ratio == 1.:
        # dual norm of L1 weighted is Linf with inverse weights
        return norm_inf_Y
    elif l1_ratio == 0.:
        # dual norm of L2 is L2
        return np.sqrt(phi.norm(Y[None, :], ord=2).sum())

    if norm_inf_Y == 0.:
        return 0.

    # ignore some values of Y by lower bound on dual norm:
    if w_time is None:
        idx = Y > l1_ratio * norm_inf_Y
    else:
        idx = Y > l1_ratio * np.max(Y / (w_space * (1. - l1_ratio) +
                                    l1_ratio * w_time))

    if idx.sum() == 1:
        return norm_inf_Y

    # sort both Y / w_time and freqs_count at the same time
    if w_time is not None:
        idx_sort = np.argsort(Y[idx] / w_time[idx])[::-1]
        w_time = w_time[idx][idx_sort]
    else:
        idx_sort = np.argsort(Y[idx])[::-1]

    Y = Y[idx][idx_sort]
    freqs_count = freqs_count[idx][idx_sort]

    Y = np.repeat(Y, freqs_count)
    if w_time is not None:
        w_time = np.repeat(w_time, freqs_count)

    K = Y.shape[0]
    if w_time is None:
        p_sum_Y2 = np.cumsum(Y ** 2)
        p_sum_w2 = np.arange(1, K + 1)
        p_sum_Yw = np.cumsum(Y)
        upper = p_sum_Y2 / Y ** 2 - 2. * p_sum_Yw / Y + p_sum_w2
    else:
        p_sum_Y2 = np.cumsum(Y ** 2)
        p_sum_w2 = np.cumsum(w_time ** 2)
        p_sum_Yw = np.cumsum(Y * w_time)
        upper = (p_sum_Y2 / (Y / w_time) ** 2 -
                 2. * p_sum_Yw / (Y / w_time) + p_sum_w2)
    upper_greater = np.where(upper > w_space ** 2 * (1. - l1_ratio) ** 2 /
                             l1_ratio ** 2)[0]

    i0 = upper_greater[0] - 1 if upper_greater.size else K - 1

    p_sum_Y2 = p_sum_Y2[i0]
    p_sum_w2 = p_sum_w2[i0]
    p_sum_Yw = p_sum_Yw[i0]

    denom = l1_ratio ** 2 * p_sum_w2 - w_space ** 2 * (1. - l1_ratio) ** 2
    if np.abs(denom) < 1e-10:
        return p_sum_Y2 / (2. * l1_ratio * p_sum_Yw)
    else:
        delta = (l1_ratio * p_sum_Yw) ** 2 - p_sum_Y2 * denom
        return (l1_ratio * p_sum_Yw - np.sqrt(delta)) / denom


def norm_epsilon_inf(G, R, phi, l1_ratio, n_orient, w_space=None, w_time=None):
    """Weighted epsilon-inf norm of phi(np.dot(G.T, R)).

    Parameters
    ----------
    G : array, shape (n_sensors, n_sources)
        Gain matrix a.k.a. lead field.
    R : array, shape (n_sensors, n_times)
        Residual.
    phi : instance of _Phi
        The TF operator.
    l1_ratio : float between 0 and 1
        Parameter controlling the tradeoff between L21 and L1 regularization.
        0 corresponds to an absence of temporal regularization, ie MxNE.
    n_orient : int
        Number of dipoles per location (typically 1 or 3).
    w_space : array, shape (n_positions,) or None.
        Weights for the L2 term of the epsilon norm. If None, weights are
        all equal to 1.
    w_time : array, shape (n_positions, n_coefs) or None
        Weights for the L1 term of the epsilon norm. If None, weights are
        all equal to 1.

    Returns
    -------
    nu : float
        The maximum value of the epsilon norms over groups of n_orient dipoles
        (consecutive rows of phi(np.dot(G.T, R))).
    """
    n_positions = G.shape[1] // n_orient
    GTRPhi = np.abs(phi(np.dot(G.T, R)))
    # norm over orientations:
    GTRPhi = GTRPhi.reshape((n_orient, -1), order='F')
    GTRPhi = np.linalg.norm(GTRPhi, axis=0)
    GTRPhi = GTRPhi.reshape((n_positions, -1), order='F')
    nu = 0.
    for idx in range(n_positions):
        GTRPhi_ = GTRPhi[idx]
        w_t = w_time[idx] if w_time is not None else None
        w_s = w_space[idx] if w_space is not None else 1.
        norm_eps = norm_epsilon(GTRPhi_, l1_ratio, phi, w_space=w_s,
                                w_time=w_t)
        if norm_eps > nu:
            nu = norm_eps

    return nu


def dgap_l21l1(M, G, Z, active_set, alpha_space, alpha_time, phi, phiT,
               n_orient, highest_d_obj, w_space=None, w_time=None):
    """Duality gap for the time-frequency mixed norm inverse problem.

    See :footcite:`GramfortEtAl2012,NdiayeEtAl2016`

    Parameters
    ----------
    M : array, shape (n_sensors, n_times)
        The data.
    G : array, shape (n_sensors, n_sources)
        Gain matrix a.k.a. lead field.
    Z : array, shape (n_active, n_coefs)
        Sources in TF domain.
    active_set : array of bool, shape (n_sources, )
        Mask of active sources.
    alpha_space : float
        The spatial regularization parameter.
    alpha_time : float
        The temporal regularization parameter. The higher it is the smoother
        will be the estimated time series.
    phi : instance of _Phi
        The TF operator.
    phiT : instance of _PhiT
        The transpose of the TF operator.
    n_orient : int
        Number of dipoles per locations (typically 1 or 3).
    highest_d_obj : float
        The highest value of the dual objective so far.
    w_space : array, shape (n_positions, )
        Array of spatial weights.
    w_time : array, shape (n_positions, n_coefs)
        Array of TF weights.

    Returns
    -------
    gap : float
        Dual gap
    p_obj : float
        Primal objective
    d_obj : float
        Dual objective. gap = p_obj - d_obj
    R : array, shape (n_sensors, n_times)
        Current residual (M - G * X)

    References
    ----------
    .. footbibliography::
    """
    X = phiT(Z)
    GX = np.dot(G[:, active_set], X)
    R = M - GX

    # some functions need w_time only on active_set, other need it completely
    if w_time is not None:
        w_time_as = w_time[active_set[::n_orient]]
    else:
        w_time_as = None
    if w_space is not None:
        w_space_as = w_space[active_set[::n_orient]]
    else:
        w_space_as = None

    penaltyl1 = norm_l1_tf(Z, phi, n_orient, w_time_as)
    penaltyl21 = norm_l21_tf(Z, phi, n_orient, w_space_as)
    nR2 = sum_squared(R)
    p_obj = 0.5 * nR2 + alpha_space * penaltyl21 + alpha_time * penaltyl1

    l1_ratio = alpha_time / (alpha_space + alpha_time)
    dual_norm = norm_epsilon_inf(G, R, phi, l1_ratio, n_orient,
                                 w_space=w_space, w_time=w_time)
    scaling = min(1., (alpha_space + alpha_time) / dual_norm)

    d_obj = (scaling - 0.5 * (scaling ** 2)) * nR2 + scaling * np.sum(R * GX)
    d_obj = max(d_obj, highest_d_obj)

    gap = p_obj - d_obj
    return gap, p_obj, d_obj, R


def _tf_mixed_norm_solver_bcd_(M, G, Z, active_set, candidates, alpha_space,
                               alpha_time, lipschitz_constant, phi, phiT,
                               w_space=None, w_time=None, n_orient=1,
                               maxit=200, tol=1e-8, dgap_freq=10, perc=None,
                               timeit=True, verbose=None):
    n_sources = G.shape[1]
    n_positions = n_sources // n_orient

    # First make G fortran for faster access to blocks of columns
    Gd = np.asfortranarray(G)
    G = np.ascontiguousarray(
        Gd.T.reshape(n_positions, n_orient, -1).transpose(0, 2, 1))

    R = M.copy()  # residual
    active = np.where(active_set[::n_orient])[0]
    for idx in active:
        R -= np.dot(G[idx], phiT(Z[idx]))

    E = []  # track primal objective function

    if w_time is None:
        alpha_time_lc = alpha_time / lipschitz_constant
    else:
        alpha_time_lc = alpha_time * w_time / lipschitz_constant[:, None]
    if w_space is None:
        alpha_space_lc = alpha_space / lipschitz_constant
    else:
        alpha_space_lc = alpha_space * w_space / lipschitz_constant

    converged = False
    d_obj = - np.inf

    for i in range(maxit):
        for jj in candidates:
            ids = jj * n_orient
            ide = ids + n_orient

            G_j = G[jj]
            Z_j = Z[jj]
            active_set_j = active_set[ids:ide]

            was_active = np.any(active_set_j)

            # gradient step
            GTR = np.dot(G_j.T, R) / lipschitz_constant[jj]
            X_j_new = GTR.copy()

            if was_active:
                X_j = phiT(Z_j)
                R += np.dot(G_j, X_j)
                X_j_new += X_j

            rows_norm = np.linalg.norm(X_j_new, 'fro')
            if rows_norm <= alpha_space_lc[jj]:
                if was_active:
                    Z[jj] = 0.0
                    active_set_j[:] = False
            else:
                GTR_phi = phi(GTR)
                if was_active:
                    Z_j_new = Z_j + GTR_phi
                else:
                    Z_j_new = GTR_phi
                col_norm = np.linalg.norm(Z_j_new, axis=0)

                if np.all(col_norm <= alpha_time_lc[jj]):
                    Z[jj] = 0.0
                    active_set_j[:] = False
                else:
                    # l1
                    shrink = np.maximum(1.0 - alpha_time_lc[jj] / np.maximum(
                                        col_norm, alpha_time_lc[jj]), 0.0)
                    if w_time is not None:
                        shrink[w_time[jj] == 0.0] = 0.0
                    Z_j_new *= shrink[np.newaxis, :]

                    # l21
                    shape_init = Z_j_new.shape
                    row_norm = np.sqrt(phi.norm(Z_j_new, ord=2).sum())
                    if row_norm <= alpha_space_lc[jj]:
                        Z[jj] = 0.0
                        active_set_j[:] = False
                    else:
                        shrink = np.maximum(
                            1.0 - alpha_space_lc[jj] /
                            np.maximum(row_norm, alpha_space_lc[jj]), 0.0)
                        Z_j_new *= shrink
                        Z[jj] = Z_j_new.reshape(-1, *shape_init[1:]).copy()
                        active_set_j[:] = True
                        Z_j_phi_T = phiT(Z[jj])
                        R -= np.dot(G_j, Z_j_phi_T)

        if (i + 1) % dgap_freq == 0:
            Zd = np.vstack([Z[pos] for pos in range(n_positions)
                            if np.any(Z[pos])])
            gap, p_obj, d_obj, _ = dgap_l21l1(
                M, Gd, Zd, active_set, alpha_space, alpha_time, phi, phiT,
                n_orient, d_obj, w_space=w_space, w_time=w_time)
            converged = (gap < tol)
            E.append(p_obj)
            logger.info("\n    Iteration %d :: n_active %d" % (
                        i + 1, np.sum(active_set) / n_orient))
            logger.info("    dgap %.2e :: p_obj %f :: d_obj %f" % (
                        gap, p_obj, d_obj))

        if converged:
            break

        if perc is not None:
            if np.sum(active_set) / float(n_orient) <= perc * n_positions:
                break

    return Z, active_set, E, converged


@verbose
def _tf_mixed_norm_solver_bcd_active_set(M, G, alpha_space, alpha_time,
                                         lipschitz_constant, phi, phiT,
                                         Z_init=None, w_space=None,
                                         w_time=None, n_orient=1, maxit=200,
                                         tol=1e-8, dgap_freq=10,
                                         verbose=None):

    n_sensors, n_times = M.shape
    n_sources = G.shape[1]
    n_positions = n_sources // n_orient

    Z = dict.fromkeys(np.arange(n_positions), 0.0)
    active_set = np.zeros(n_sources, dtype=bool)
    active = []
    if Z_init is not None:
        if Z_init.shape != (n_sources, phi.n_coefs.sum()):
            raise Exception('Z_init must be None or an array with shape '
                            '(n_sources, n_coefs).')
        for ii in range(n_positions):
            if np.any(Z_init[ii * n_orient:(ii + 1) * n_orient]):
                active_set[ii * n_orient:(ii + 1) * n_orient] = True
                active.append(ii)
        if len(active):
            Z.update(dict(zip(active,
                              np.vsplit(Z_init[active_set], len(active)))))

    E = []
    candidates = range(n_positions)
    d_obj = -np.inf

    while True:
        # single BCD pass on all positions:
        Z_init = dict.fromkeys(np.arange(n_positions), 0.0)
        Z_init.update(dict(zip(active, Z.values())))
        Z, active_set, E_tmp, _ = _tf_mixed_norm_solver_bcd_(
            M, G, Z_init, active_set, candidates, alpha_space, alpha_time,
            lipschitz_constant, phi, phiT, w_space=w_space, w_time=w_time,
            n_orient=n_orient, maxit=1, tol=tol, perc=None, verbose=verbose)

        E += E_tmp

        # multiple BCD pass on active positions:
        active = np.where(active_set[::n_orient])[0]
        Z_init = dict(zip(range(len(active)), [Z[idx] for idx in active]))
        candidates_ = range(len(active))
        if w_space is not None:
            w_space_as = w_space[active_set[::n_orient]]
        else:
            w_space_as = None
        if w_time is not None:
            w_time_as = w_time[active_set[::n_orient]]
        else:
            w_time_as = None

        Z, as_, E_tmp, converged = _tf_mixed_norm_solver_bcd_(
            M, G[:, active_set], Z_init,
            np.ones(len(active) * n_orient, dtype=bool),
            candidates_, alpha_space, alpha_time,
            lipschitz_constant[active_set[::n_orient]], phi, phiT,
            w_space=w_space_as, w_time=w_time_as,
            n_orient=n_orient, maxit=maxit, tol=tol,
            dgap_freq=dgap_freq, perc=0.5,
            verbose=verbose)
        active = np.where(active_set[::n_orient])[0]
        active_set[active_set] = as_.copy()
        E += E_tmp

        converged = True
        if converged:
            Zd = np.vstack([Z[pos] for pos in range(len(Z)) if np.any(Z[pos])])
            gap, p_obj, d_obj, _ = dgap_l21l1(
                M, G, Zd, active_set, alpha_space, alpha_time,
                phi, phiT, n_orient, d_obj, w_space, w_time)
            logger.info("\ndgap %.2e :: p_obj %f :: d_obj %f :: n_active %d"
                        % (gap, p_obj, d_obj, np.sum(active_set) / n_orient))
            if gap < tol:
                logger.info("\nConvergence reached!\n")
                break

    if active_set.sum():
        Z = np.vstack([Z[pos] for pos in range(len(Z)) if np.any(Z[pos])])
        X = phiT(Z)
    else:
        Z = np.zeros((0, phi.n_coefs.sum()), dtype=np.complex128)
        X = np.zeros((0, n_times))

    return X, Z, active_set, E, gap


@verbose
def tf_mixed_norm_solver(M, G, alpha_space, alpha_time, wsize=64, tstep=4,
                         n_orient=1, maxit=200, tol=1e-8,
                         active_set_size=None, debias=True, return_gap=False,
                         dgap_freq=10, verbose=None):
    """Solve TF L21+L1 inverse solver with BCD and active set approach.

    See :footcite:`GramfortEtAl2013b,GramfortEtAl2011,BekhtiEtAl2016`.

    Parameters
    ----------
    M : array, shape (n_sensors, n_times)
        The data.
    G : array, shape (n_sensors, n_dipoles)
        The gain matrix a.k.a. lead field.
    alpha_space : float
        The spatial regularization parameter.
    alpha_time : float
        The temporal regularization parameter. The higher it is the smoother
        will be the estimated time series.
    wsize: int or array-like
        Length of the STFT window in samples (must be a multiple of 4).
        If an array is passed, multiple TF dictionaries are used (each having
        its own wsize and tstep) and each entry of wsize must be a multiple
        of 4.
    tstep: int or array-like
        Step between successive windows in samples (must be a multiple of 2,
        a divider of wsize and smaller than wsize/2) (default: wsize/2).
        If an array is passed, multiple TF dictionaries are used (each having
        its own wsize and tstep), and each entry of tstep must be a multiple
        of 2 and divide the corresponding entry of wsize.
    n_orient : int
        The number of orientation (1 : fixed or 3 : free or loose).
    maxit : int
        The number of iterations.
    tol : float
        If absolute difference between estimates at 2 successive iterations
        is lower than tol, the convergence is reached.
    debias : bool
        Debias source estimates.
    return_gap : bool
        Return final duality gap.
    dgap_freq : int or np.inf
        The duality gap is evaluated every dgap_freq iterations.
    %(verbose)s

    Returns
    -------
    X : array, shape (n_active, n_times)
        The source estimates.
    active_set : array
        The mask of active sources.
    E : list
        The value of the objective function every dgap_freq iteration. If
        log_objective is False or dgap_freq is np.inf, it will be empty.
    gap : float
        Final duality gap. Returned only if return_gap is True.

    References
    ----------
    .. footbibliography::
    """
    n_sensors, n_times = M.shape
    n_sensors, n_sources = G.shape
    n_positions = n_sources // n_orient

    tstep = np.atleast_1d(tstep)
    wsize = np.atleast_1d(wsize)
    if len(tstep) != len(wsize):
        raise ValueError('The same number of window sizes and steps must be '
                         'passed. Got tstep = %s and wsize = %s' %
                         (tstep, wsize))

    n_steps = np.ceil(M.shape[1] / 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)

    if n_orient == 1:
        lc = np.sum(G * G, axis=0)
    else:
        lc = np.empty(n_positions)
        for j in range(n_positions):
            G_tmp = G[:, (j * n_orient):((j + 1) * n_orient)]
            lc[j] = np.linalg.norm(np.dot(G_tmp.T, G_tmp), ord=2)

    logger.info("Using block coordinate descent with active set approach")
    X, Z, active_set, E, gap = _tf_mixed_norm_solver_bcd_active_set(
        M, G, alpha_space, alpha_time, lc, phi, phiT,
        Z_init=None, n_orient=n_orient, maxit=maxit, tol=tol,
        dgap_freq=dgap_freq, verbose=None)

    if np.any(active_set) and debias:
        bias = compute_bias(M, G[:, active_set], X, n_orient=n_orient)
        X *= bias[:, np.newaxis]

    if return_gap:
        return X, active_set, E, gap
    else:
        return X, active_set, E


@verbose
def iterative_tf_mixed_norm_solver(M, G, alpha_space, alpha_time,
                                   n_tfmxne_iter, wsize=64, tstep=4,
                                   maxit=3000, tol=1e-8, debias=True,
                                   n_orient=1, dgap_freq=10, verbose=None):
    """Solve TF L0.5/L1 + L0.5 inverse problem with BCD + active set approach.

    Parameters
    ----------
    M: array, shape (n_sensors, n_times)
        The data.
    G: array, shape (n_sensors, n_dipoles)
        The gain matrix a.k.a. lead field.
    alpha_space: float
        The spatial regularization parameter. The higher it is the less there
        will be active sources.
    alpha_time : float
        The temporal regularization parameter. The higher it is the smoother
        will be the estimated time series. 0 means no temporal regularization,
        a.k.a. irMxNE.
    n_tfmxne_iter : int
        Number of TF-MxNE iterations. If > 1, iterative reweighting is applied.
    wsize : int or array-like
        Length of the STFT window in samples (must be a multiple of 4).
        If an array is passed, multiple TF dictionaries are used (each having
        its own wsize and tstep) and each entry of wsize must be a multiple
        of 4.
    tstep : int or array-like
        Step between successive windows in samples (must be a multiple of 2,
        a divider of wsize and smaller than wsize/2) (default: wsize/2).
        If an array is passed, multiple TF dictionaries are used (each having
        its own wsize and tstep), and each entry of tstep must be a multiple
        of 2 and divide the corresponding entry of wsize.
    maxit : int
        The maximum number of iterations for each TF-MxNE problem.
    tol : float
        If absolute difference between estimates at 2 successive iterations
        is lower than tol, the convergence is reached. Also used as criterion
        on duality gap for each TF-MxNE problem.
    debias : bool
        Debias source estimates.
    n_orient : int
        The number of orientation (1 : fixed or 3 : free or loose).
    dgap_freq : int or np.inf
        The duality gap is evaluated every dgap_freq iterations.
    %(verbose)s

    Returns
    -------
    X : array, shape (n_active, n_times)
        The source estimates.
    active_set : array
        The mask of active sources.
    E : list
        The value of the objective function over iterations.
    """
    n_sensors, n_times = M.shape
    n_sources = G.shape[1]
    n_positions = n_sources // n_orient

    tstep = np.atleast_1d(tstep)
    wsize = np.atleast_1d(wsize)
    if len(tstep) != len(wsize):
        raise ValueError('The same number of window sizes and steps must be '
                         'passed. Got tstep = %s and wsize = %s' %
                         (tstep, wsize))

    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)

    if n_orient == 1:
        lc = np.sum(G * G, axis=0)
    else:
        lc = np.empty(n_positions)
        for j in range(n_positions):
            G_tmp = G[:, (j * n_orient):((j + 1) * n_orient)]
            lc[j] = np.linalg.norm(np.dot(G_tmp.T, G_tmp), ord=2)

    # space and time penalties, and inverse of their derivatives:
    def g_space(Z):
        return np.sqrt(np.sqrt(phi.norm(Z, ord=2).reshape(
            -1, n_orient).sum(axis=1)))

    def g_space_prime_inv(Z):
        return 2. * g_space(Z)

    def g_time(Z):
        return np.sqrt(np.sqrt(np.sum((np.abs(Z) ** 2.).reshape(
            (n_orient, -1), order='F'), axis=0)).reshape(
            (-1, Z.shape[1]), order='F'))

    def g_time_prime_inv(Z):
        return 2. * g_time(Z)

    E = list()

    active_set = np.ones(n_sources, dtype=bool)
    Z = np.zeros((n_sources, phi.n_coefs.sum()), dtype=np.complex128)

    for k in range(n_tfmxne_iter):
        active_set_0 = active_set.copy()
        Z0 = Z.copy()

        if k == 0:
            w_space = None
            w_time = None
        else:
            w_space = 1. / g_space_prime_inv(Z)
            w_time = g_time_prime_inv(Z)
            w_time[w_time == 0.0] = -1.
            w_time = 1. / w_time
            w_time[w_time < 0.0] = 0.0

        X, Z, active_set_, E_, _ = _tf_mixed_norm_solver_bcd_active_set(
            M, G[:, active_set], alpha_space, alpha_time,
            lc[active_set[::n_orient]], phi, phiT,
            Z_init=Z, w_space=w_space, w_time=w_time, n_orient=n_orient,
            maxit=maxit, tol=tol, dgap_freq=dgap_freq, verbose=None)

        active_set[active_set] = active_set_

        if active_set.sum() > 0:
            l21_penalty = np.sum(g_space(Z.copy()))
            l1_penalty = phi.norm(g_time(Z.copy()), ord=1).sum()

            p_obj = (0.5 * np.linalg.norm(M - np.dot(G[:, active_set], X),
                     'fro') ** 2. + alpha_space * l21_penalty +
                     alpha_time * l1_penalty)
            E.append(p_obj)

            logger.info('Iteration %d: active set size=%d, E=%f' % (
                        k + 1, active_set.sum() / n_orient, p_obj))

            # Check convergence
            if np.array_equal(active_set, active_set_0):
                max_diff = np.amax(np.abs(Z - Z0))
                if (max_diff < tol):
                    print('Convergence reached after %d reweightings!' % k)
                    break
        else:
            p_obj = 0.5 * np.linalg.norm(M) ** 2.
            E.append(p_obj)
            logger.info('Iteration %d: as_size=%d, E=%f' % (
                        k + 1, active_set.sum() / n_orient, p_obj))
            break

    if debias:
        if active_set.sum() > 0:
            bias = compute_bias(M, G[:, active_set], X, n_orient=n_orient)
            X *= bias[:, np.newaxis]

    return X, active_set, E