cluster_level_back.py

#!/usr/bin/env python

# Authors: Thorsten Kranz <thorstenkranz@gmail.com>
#          Alexandre Gramfort <alexandre.gramfort@inria.fr>
#          Martin Luessi <mluessi@nmr.mgh.harvard.edu>
#          Eric Larson <larson.eric.d@gmail.com>
#          Denis Engemann <denis.engemann@gmail.com>
#          Fernando Perez (bin_perm_rep function)
#
# License: BSD-3-Clause
# Copyright the MNE-Python contributors.

import numpy as np
from scipy import ndimage, sparse
from scipy.sparse.csgraph import connected_components
from scipy.stats import f as fstat
from scipy.stats import t as tstat

from ..fixes import has_numba, jit
from ..parallel import parallel_func
from ..source_estimate import MixedSourceEstimate, SourceEstimate, VolSourceEstimate
from ..source_space import SourceSpaces
from ..utils import (
    ProgressBar,
    _check_option,
    _pl,
    _validate_type,
    check_random_state,
    logger,
    split_list,
    verbose,
    warn,
)
from .parametric import f_oneway, ttest_1samp_no_p


def _get_buddies_fallback(r, s, neighbors, indices=None):
    if indices is None:
        buddies = np.where(r)[0]
    else:
        buddies = indices[r[indices]]
    buddies = buddies[np.isin(s[buddies], neighbors, assume_unique=True)]
    r[buddies] = False
    return buddies.tolist()


def _get_selves_fallback(r, s, ind, inds, t, t_border, max_step):
    start = t_border[max(t[ind] - max_step, 0)]
    stop = t_border[min(t[ind] + max_step + 1, len(t_border) - 1)]
    indices = inds[start:stop]
    selves = indices[r[indices]]
    selves = selves[s[ind] == s[selves]]
    r[selves] = False
    return selves.tolist()


def _where_first_fallback(x):
    # this is equivalent to np.where(r)[0] for these purposes, but it's
    # a little bit faster. Unfortunately there's no way to tell numpy
    # just to find the first instance (to save checking every one):
    next_ind = int(np.argmax(x))
    if next_ind == 0:
        next_ind = -1
    return next_ind


if has_numba:  # pragma: no cover

    @jit()
    def _get_buddies(r, s, neighbors, indices=None):
        buddies = list()
        # At some point we might be able to use the sorted-ness of s or
        # neighbors to further speed this up
        if indices is None:
            n_check = len(r)
        else:
            n_check = len(indices)
        for ii in range(n_check):
            if indices is None:
                this_idx = ii
            else:
                this_idx = indices[ii]
            if r[this_idx]:
                this_s = s[this_idx]
                for ni in range(len(neighbors)):
                    if this_s == neighbors[ni]:
                        buddies.append(this_idx)
                        r[this_idx] = False
                        break
        return buddies


    @jit()
    def _get_selves(r, s, ind, inds, t, t_border, max_step):
        selves = list()
        start = t_border[max(t[ind] - max_step, 0)]
        stop = t_border[min(t[ind] + max_step + 1, len(t_border) - 1)]
        for ii in range(start, stop):
            this_idx = inds[ii]
            if r[this_idx] and s[ind] == s[this_idx]:
                selves.append(this_idx)
                r[this_idx] = False
        return selves


    @jit()
    def _where_first(x):
        for ii in range(len(x)):
            if x[ii]:
                return ii
        return -1

else:  # pragma: no cover
    # fastest ways we've found with NumPy
    _get_buddies = _get_buddies_fallback
    _get_selves = _get_selves_fallback
    _where_first = _where_first_fallback


@jit()
def _masked_sum(x, c):
    return np.sum(x[c])


@jit()
def _masked_sum_power(x, c, t_power):
    return np.sum(np.sign(x[c]) * np.abs(x[c]) ** t_power)


@jit()
def _sum_cluster_data(data, tstep):
    return np.sign(data) * np.logical_not(data == 0) * tstep


def _get_clusters_spatial(s, neighbors):
    """Form spatial clusters using neighbor lists.

    This is equivalent to _get_components with n_times = 1, with a properly
    reconfigured adjacency matrix (formed as "neighbors" list)
    """
    # s is a vector of spatial indices that are significant, like:
    #     s = np.where(x_in)[0]
    # for x_in representing a single time-instant
    r = np.ones(s.shape, bool)
    clusters = list()
    next_ind = 0 if s.size > 0 else -1
    while next_ind >= 0:
        # put first point in a cluster, adjust remaining
        t_inds = [next_ind]
        r[next_ind] = 0
        icount = 1  # count of nodes in the current cluster
        while icount <= len(t_inds):
            ind = t_inds[icount - 1]
            # look across other vertices
            buddies = _get_buddies(r, s, neighbors[s[ind]])
            t_inds.extend(buddies)
            icount += 1
        next_ind = _where_first(r)
        clusters.append(s[t_inds])
    return clusters


def _reassign(check, clusters, base, num):
    """Reassign cluster numbers."""
    # reconfigure check matrix
    check[check == num] = base
    # concatenate new values into clusters array
    clusters[base - 1] = np.concatenate((clusters[base - 1], clusters[num - 1]))
    clusters[num - 1] = np.array([], dtype=int)


def _get_clusters_st_1step(keepers, neighbors):
    """Directly calculate clusters.

    This uses knowledge that time points are
    only adjacent to immediate neighbors for data organized as time x space.

    This algorithm time increases linearly with the number of time points,
    compared to with the square for the standard (graph) algorithm.

    This algorithm creates clusters for each time point using a method more
    efficient than the standard graph method (but otherwise equivalent), then
    combines these clusters across time points in a reasonable way.
    """
    n_src = len(neighbors)
    n_times = len(keepers)
    # start cluster numbering at 1 for diffing convenience
    enum_offset = 1
    check = np.zeros((n_times, n_src), dtype=int)
    clusters = list()
    for ii, k in enumerate(keepers):
        c = _get_clusters_spatial(k, neighbors)
        for ci, cl in enumerate(c):
            check[ii, cl] = ci + enum_offset
        enum_offset += len(c)
        # give them the correct offsets
        c = [cl + ii * n_src for cl in c]
        clusters += c

    # now that each cluster has been assigned a unique number, combine them
    # by going through each time point
    for check1, check2, k in zip(check[:-1], check[1:], keepers[:-1]):
        # go through each one that needs reassignment
        inds = k[check2[k] - check1[k] > 0]
        check1_d = check1[inds]
        n = check2[inds]
        nexts = np.unique(n)
        for num in nexts:
            prevs = check1_d[n == num]
            base = np.min(prevs)
            for pr in np.unique(prevs[prevs != base]):
                _reassign(check1, clusters, base, pr)
            # reassign values
            _reassign(check2, clusters, base, num)
    # clean up clusters
    clusters = [cl for cl in clusters if len(cl) > 0]
    return clusters


def _get_clusters_st_multistep(keepers, neighbors, max_step=1):
    """Directly calculate clusters.

    This uses knowledge that time points are
    only adjacent to immediate neighbors for data organized as time x space.

    This algorithm time increases linearly with the number of time points,
    compared to with the square for the standard (graph) algorithm.
    """
    n_src = len(neighbors)
    n_times = len(keepers)
    t_border = list()
    t_border.append(0)
    for ki, k in enumerate(keepers):
        keepers[ki] = k + ki * n_src
        t_border.append(t_border[ki] + len(k))
    t_border = np.array(t_border)
    keepers = np.concatenate(keepers)
    v = keepers
    t, s = divmod(v, n_src)

    r = np.ones(t.shape, dtype=bool)
    clusters = list()
    inds = np.arange(t_border[0], t_border[n_times])
    next_ind = 0 if s.size > 0 else -1
    while next_ind >= 0:
        # put first point in a cluster, adjust remaining
        t_inds = [next_ind]
        r[next_ind] = False
        icount = 1  # count of nodes in the current cluster
        # look for significant values at the next time point,
        # same sensor, not placed yet, and add those
        while icount <= len(t_inds):
            ind = t_inds[icount - 1]
            selves = _get_selves(r, s, ind, inds, t, t_border, max_step)

            # look at current time point across other vertices
            these_inds = inds[t_border[t[ind]]: t_border[t[ind] + 1]]
            buddies = _get_buddies(r, s, neighbors[s[ind]], these_inds)

            t_inds += buddies + selves
            icount += 1
        next_ind = _where_first(r)
        clusters.append(v[t_inds])

    return clusters


def _get_clusters_st(x_in, neighbors, max_step=1):
    """Choose the most efficient version."""
    n_src = len(neighbors)
    n_times = x_in.size // n_src
    cl_goods = np.where(x_in)[0]
    if len(cl_goods) > 0:
        keepers = [np.array([], dtype=int)] * n_times
        row, col = np.unravel_index(cl_goods, (n_times, n_src))
        lims = [0]
        if isinstance(row, int):
            row = [row]
            col = [col]
        else:
            order = np.argsort(row)
            row = row[order]
            col = col[order]
            lims += (np.where(np.diff(row) > 0)[0] + 1).tolist()
            lims.append(len(row))

        for start, end in zip(lims[:-1], lims[1:]):
            keepers[row[start]] = np.sort(col[start:end])
        if max_step == 1:
            return _get_clusters_st_1step(keepers, neighbors)
        else:
            return _get_clusters_st_multistep(keepers, neighbors, max_step)
    else:
        return []


def _get_components(x_in, adjacency, return_list=True):
    """Get connected components from a mask and a adjacency matrix."""
    if adjacency is False:
        components = np.arange(len(x_in))
    else:
        mask = np.logical_and(x_in[adjacency.row], x_in[adjacency.col])
        data = adjacency.data[mask]
        row = adjacency.row[mask]
        col = adjacency.col[mask]
        shape = adjacency.shape
        idx = np.where(x_in)[0]
        row = np.concatenate((row, idx))
        col = np.concatenate((col, idx))
        data = np.concatenate((data, np.ones(len(idx), dtype=data.dtype)))
        adjacency = sparse.coo_matrix((data, (row, col)), shape=shape)
        _, components = connected_components(adjacency)
    if return_list:
        start = np.min(components)
        stop = np.max(components)
        comp_list = [list() for i in range(start, stop + 1, 1)]
        mask = np.zeros(len(comp_list), dtype=bool)
        for ii, comp in enumerate(components):
            comp_list[comp].append(ii)
            mask[comp] += x_in[ii]
        clusters = [np.array(k) for k, m in zip(comp_list, mask) if m]
        return clusters
    else:
        return components


'''
if tfce = True, then the function will return the TFCE scores of each point BUT cluster!!!!
so in this case, treat this function as a data correction function..  convert t to tfce
'''


def _find_clusters(
        x,
        threshold,
        tail=0,
        adjacency=None,
        max_step=1,
        include=None,
        partitions=None,
        t_power=1,
        show_info=False,
):
    """Find all clusters which are above/below a certain threshold.

    When doing a two-tailed test (tail == 0), only points with the same
    sign will be clustered together.

    Parameters
    ----------
    x : 1D array
        Data
    threshold : float | dict
        Where to threshold the statistic. Should be negative for tail == -1,
        and positive for tail == 0 or 1. Can also be an dict for
        threshold-free cluster enhancement.
    tail : -1 | 0 | 1
        Type of comparison
    adjacency : scipy.sparse.coo_matrix, None, or list
        Defines adjacency between features. The matrix is assumed to
        be symmetric and only the upper triangular half is used.
        If adjacency is a list, it is assumed that each entry stores the
        indices of the spatial neighbors in a spatio-temporal dataset x.
        Default is None, i.e, a regular lattice adjacency.
        False means no adjacency.
    max_step : int
        If adjacency is a list, this defines the maximal number of steps
        between vertices along the second dimension (typically time) to be
        considered adjacent.
    include : 1D bool array or None
        Mask to apply to the data of points to cluster. If None, all points
        are used.
    partitions : array of int or None
        An array (same size as X) of integers indicating which points belong
        to each partition.
    t_power : float
        Power to raise the statistical values (usually t-values) by before
        summing (sign will be retained). Note that t_power == 0 will give a
        count of nodes in each cluster, t_power == 1 will weight each node by
        its statistical score.
    show_info : bool
        If True, display information about thresholds used (for TFCE). Should
        only be done for the standard permutation.

    Returns
    -------
    clusters : list of slices or list of arrays (boolean masks)
        We use slices for 1D signals and mask to multidimensional
        arrays.
    sums : array
        Sum of x values in clusters.
    """
    _check_option("tail", tail, [-1, 0, 1])

    x = np.asanyarray(x)

    if not np.isscalar(threshold):
        if not isinstance(threshold, dict):
            raise TypeError(
                "threshold must be a number, or a dict for "
                "threshold-free cluster enhancement"
            )
        if not all(key in threshold for key in ["start", "step"]):
            raise KeyError(
                "threshold, if dict, must have at least " '"start" and "step"'
            )
        tfce = True
        use_x = x[np.isfinite(x)]
        if use_x.size == 0:
            raise RuntimeError(
                "No finite values found in the observed statistic values"
            )
        if tail == -1:
            if threshold["start"] > 0:
                raise ValueError('threshold["start"] must be <= 0 for ' "tail == -1")
            if threshold["step"] >= 0:
                raise ValueError('threshold["step"] must be < 0 for ' "tail == -1")
            stop = np.min(use_x)
        elif tail == 1:
            stop = np.max(use_x)
        else:  # tail == 0
            stop = max(np.max(use_x), -np.min(use_x))
        del use_x
        thresholds = np.arange(threshold["start"], stop, threshold["step"], float)
        h_power = threshold.get("h_power", 2)
        e_power = threshold.get("e_power", 0.5)
        if show_info is True:
            if len(thresholds) == 0:
                warn(
                    'threshold["start"] (%s) is more extreme than data '
                    "statistics with most extreme value %s" % (threshold["start"], stop)
                )
            else:
                logger.info(
                    "Using %d thresholds from %0.2f to %0.2f for TFCE "
                    "computation (h_power=%0.2f, e_power=%0.2f)"
                    % (len(thresholds), thresholds[0], thresholds[-1], h_power, e_power)
                )
        scores = np.zeros(x.size)
    else:
        thresholds = [threshold]
        tfce = False

    # include all points by default
    if include is None:
        include = np.ones(x.shape, dtype=bool)

    if tail in [0, 1] and not np.all(np.diff(thresholds) > 0):
        raise ValueError("Thresholds must be monotonically increasing")
    if tail == -1 and not np.all(np.diff(thresholds) < 0):
        raise ValueError("Thresholds must be monotonically decreasing")

    # set these here just in case thresholds == []
    clusters = list()
    sums = list()
    for ti, thresh in enumerate(thresholds):
        # these need to be reset on each run
        clusters = list()
        if tail == 0:
            x_ins = [
                np.logical_and(x > thresh, include),
                np.logical_and(x < -thresh, include),
            ]
        elif tail == -1:
            x_ins = [np.logical_and(x < thresh, include)]
        else:  # tail == 1
            x_ins = [np.logical_and(x > thresh, include)]
        # loop over tails
        for x_in in x_ins:
            if np.any(x_in):
                out = _find_clusters_1dir_parts(
                    x, x_in, adjacency, max_step, partitions, t_power, ndimage
                )
                clusters += out[0]
                sums.append(out[1])
        if tfce:
            # the score of each point is the sum of the h^H * e^E for each
            # supporting section "rectangle" h x e.
            if ti == 0:
                h = abs(thresh)
            else:
                h = abs(thresh - thresholds[ti - 1])
            h = h ** h_power
            for c in clusters:
                # triage based on cluster storage type
                if isinstance(c, slice):
                    len_c = c.stop - c.start
                elif isinstance(c, tuple):
                    len_c = len(c)
                elif c.dtype == bool:
                    len_c = np.sum(c)
                else:
                    len_c = len(c)
                scores[c] += h * (len_c ** e_power)
    # turn sums into array
    sums = np.concatenate(sums) if sums else np.array([])
    if tfce:
        # each point gets treated independently
        clusters = np.arange(x.size)
        if adjacency is None or adjacency is False:
            if x.ndim == 1:
                # slices
                clusters = [slice(c, c + 1) for c in clusters]
            else:
                # boolean masks (raveled)
                clusters = [(clusters == ii).ravel() for ii in range(len(clusters))]
        else:
            clusters = [np.array([c]) for c in clusters]
        sums = scores
    return clusters, sums


def _find_clusters_1dir_parts(
        x, x_in, adjacency, max_step, partitions, t_power, ndimage
):
    """Deal with partitions, and pass the work to _find_clusters_1dir."""
    if partitions is None:
        clusters, sums = _find_clusters_1dir(
            x, x_in, adjacency, max_step, t_power, ndimage
        )
    else:
        # cluster each partition separately
        clusters = list()
        sums = list()
        for p in range(np.max(partitions) + 1):
            x_i = np.logical_and(x_in, partitions == p)
            out = _find_clusters_1dir(x, x_i, adjacency, max_step, t_power, ndimage)
            clusters += out[0]
            sums.append(out[1])
        sums = np.concatenate(sums)
    return clusters, sums


def _find_clusters_1dir(x, x_in, adjacency, max_step, t_power, ndimage):
    """Actually call the clustering algorithm."""
    if adjacency is None:
        labels, n_labels = ndimage.label(x_in)

        if x.ndim == 1:
            # slices
            clusters = ndimage.find_objects(labels, n_labels)
            # equivalent to if len(clusters) == 0 but faster
            if not clusters:
                sums = list()
            else:
                index = list(range(1, n_labels + 1))
                if t_power == 1:
                    sums = ndimage.sum(x, labels, index=index)
                else:
                    sums = ndimage.sum(
                        np.sign(x) * np.abs(x) ** t_power, labels, index=index
                    )
        else:
            # boolean masks (raveled)
            clusters = list()
            sums = np.empty(n_labels)
            for label in range(n_labels):
                c = labels == label + 1
                clusters.append(c.ravel())
                if t_power == 1:
                    sums[label] = np.sum(x[c])
                else:
                    sums[label] = np.sum(np.sign(x[c]) * np.abs(x[c]) ** t_power)
    else:
        if x.ndim > 1:
            raise Exception(
                "Data should be 1D when using a adjacency " "to define clusters."
            )
        if isinstance(adjacency, sparse.spmatrix) or adjacency is False:
            clusters = _get_components(x_in, adjacency)
        elif isinstance(adjacency, list):  # use temporal adjacency
            clusters = _get_clusters_st(x_in, adjacency, max_step)
        else:
            raise ValueError("adjacency must be a sparse matrix or list")
        if t_power == 1:
            sums = [_masked_sum(x, c) for c in clusters]
        else:
            sums = [_masked_sum_power(x, c, t_power) for c in clusters]

    return clusters, np.atleast_1d(sums)


def _cluster_indices_to_mask(components, n_tot):
    """Convert to the old format of clusters, which were bool arrays."""
    for ci, c in enumerate(components):
        components[ci] = np.zeros((n_tot), dtype=bool)
        components[ci][c] = True
    return components


def _cluster_mask_to_indices(components, shape):
    """Convert to the old format of clusters, which were bool arrays."""
    for ci, c in enumerate(components):
        if isinstance(c, np.ndarray):  # mask
            components[ci] = np.where(c.reshape(shape))
        elif isinstance(c, slice):
            components[ci] = np.arange(c.start, c.stop)
        else:
            assert isinstance(c, tuple), type(c)
            c = list(c)  # tuple->list
            for ii, cc in enumerate(c):
                if isinstance(cc, slice):
                    c[ii] = np.arange(cc.start, cc.stop)
                else:
                    c[ii] = np.where(cc)[0]
            components[ci] = tuple(c)
    return components


def _pval_from_histogram(T, H0, tail):
    """Get p-values from stats values given an H0 distribution.

    For each stat compute a p-value as percentile of its statistics
    within all statistics in surrogate data
    """
    # from pct to fraction
    if tail == -1:  # up tail
        pval = np.array([np.mean(H0 <= t) for t in T])
    elif tail == 1:  # low tail
        pval = np.array([np.mean(H0 >= t) for t in T])
    else:  # both tails
        pval = np.array([np.mean(abs(H0) >= abs(t)) for t in T])

    return pval


def _setup_adjacency(adjacency, n_tests, n_times):
    if not sparse.issparse(adjacency):
        raise ValueError(
            "If adjacency matrix is given, it must be a " "SciPy sparse matrix."
        )
    if adjacency.shape[0] == n_tests:  # use global algorithm
        adjacency = adjacency.tocoo()
    else:  # use temporal adjacency algorithm
        got_times, mod = divmod(n_tests, adjacency.shape[0])
        if got_times != n_times or mod != 0:
            raise ValueError(
                "adjacency (len %d) must be of the correct size, i.e. be "
                "equal to or evenly divide the number of tests (%d).\n\n"
                "If adjacency was computed for a source space, try using "
                'the fwd["src"] or inv["src"] as some original source space '
                "vertices can be excluded during forward computation"
                % (adjacency.shape[0], n_tests)
            )
        # we claim to only use upper triangular part... not true here
        adjacency = (adjacency + adjacency.transpose()).tocsr()
        adjacency = [
            adjacency.indices[adjacency.indptr[i]: adjacency.indptr[i + 1]]
            for i in range(len(adjacency.indptr) - 1)
        ]
    return adjacency


def _do_permutations(
        X_full,
        slices,
        threshold,
        tail,
        adjacency,
        stat_fun,
        max_step,
        include,
        partitions,
        t_power,
        orders,
        sample_shape,
        buffer_size,
        progress_bar,
):
    n_samp, n_vars = X_full.shape

    if buffer_size is not None and n_vars <= buffer_size:
        buffer_size = None  # don't use buffer for few variables

    # allocate space for output
    max_cluster_sums = np.empty(len(orders), dtype=np.double)

    if buffer_size is not None:
        # allocate buffer, so we don't need to allocate memory during loop
        X_buffer = [
            np.empty((len(X_full[s]), buffer_size), dtype=X_full.dtype) for s in slices
        ]

    for seed_idx, order in enumerate(orders):
        # shuffle sample indices
        assert order is not None
        idx_shuffle_list = [order[s] for s in slices]

        if buffer_size is None:
            # shuffle all data at once
            X_shuffle_list = [X_full[idx, :] for idx in idx_shuffle_list]
            t_obs_surr = stat_fun(*X_shuffle_list)
        else:
            # only shuffle a small data buffer, so we need less memory
            t_obs_surr = np.empty(n_vars, dtype=X_full.dtype)

            for pos in range(0, n_vars, buffer_size):
                # number of variables for this loop
                n_var_loop = min(pos + buffer_size, n_vars) - pos

                # fill buffer
                for i, idx in enumerate(idx_shuffle_list):
                    X_buffer[i][:, :n_var_loop] = X_full[idx, pos: pos + n_var_loop]

                # apply stat_fun and store result
                tmp = stat_fun(*X_buffer)
                t_obs_surr[pos: pos + n_var_loop] = tmp[:n_var_loop]

        # The stat should have the same shape as the samples for no adj.
        if adjacency is None:
            t_obs_surr.shape = sample_shape

        # Find cluster on randomized stats
        out = _find_clusters(
            t_obs_surr,
            threshold=threshold,
            tail=tail,
            max_step=max_step,
            adjacency=adjacency,
            partitions=partitions,
            include=include,
            t_power=t_power,
        )
        perm_clusters_sums = out[1]

        if len(perm_clusters_sums) > 0:
            max_cluster_sums[seed_idx] = np.max(perm_clusters_sums)
        else:
            max_cluster_sums[seed_idx] = 0

        progress_bar.update(seed_idx + 1)

    return max_cluster_sums


def _do_1samp_permutations(
        X,
        slices,
        threshold,
        tail,
        adjacency,
        stat_fun,
        max_step,
        include,
        partitions,
        t_power,
        orders,
        sample_shape,
        buffer_size,
        progress_bar,
):
    n_samp, n_vars = X.shape
    assert slices is None  # should be None for the 1 sample case

    if buffer_size is not None and n_vars <= buffer_size:
        buffer_size = None  # don't use buffer for few variables

    # allocate space for output
    max_cluster_sums = np.empty(len(orders), dtype=np.double)

    if buffer_size is not None:
        # allocate a buffer so we don't need to allocate memory in loop
        X_flip_buffer = np.empty((n_samp, buffer_size), dtype=X.dtype)

    for seed_idx, order in enumerate(orders):
        assert isinstance(order, np.ndarray)
        # new surrogate data with specified sign flip
        assert order.size == n_samp  # should be guaranteed by parent
        signs = 2 * order[:, None].astype(int) - 1
        if not np.all(np.equal(np.abs(signs), 1)):
            raise ValueError("signs from rng must be +/- 1")

        if buffer_size is None:
            # be careful about non-writable memmap (GH#1507)
            if X.flags.writeable:
                X *= signs
                # Recompute statistic on randomized data
                t_obs_surr = stat_fun(X)
                # Set X back to previous state (trade memory eff. for CPU use)
                X *= signs
            else:
                t_obs_surr = stat_fun(X * signs)
        else:
            # only sign-flip a small data buffer, so we need less memory
            t_obs_surr = np.empty(n_vars, dtype=X.dtype)

            for pos in range(0, n_vars, buffer_size):
                # number of variables for this loop
                n_var_loop = min(pos + buffer_size, n_vars) - pos

                X_flip_buffer[:, :n_var_loop] = signs * X[:, pos: pos + n_var_loop]

                # apply stat_fun and store result
                tmp = stat_fun(X_flip_buffer)
                t_obs_surr[pos: pos + n_var_loop] = tmp[:n_var_loop]

        # The stat should have the same shape as the samples for no adj.
        if adjacency is None:
            t_obs_surr.shape = sample_shape

        # Find cluster on randomized stats
        out = _find_clusters(
            t_obs_surr,
            threshold=threshold,
            tail=tail,
            max_step=max_step,
            adjacency=adjacency,
            partitions=partitions,
            include=include,
            t_power=t_power,
        )
        perm_clusters_sums = out[1]
        if len(perm_clusters_sums) > 0:
            # get max with sign info
            idx_max = np.argmax(np.abs(perm_clusters_sums))
            max_cluster_sums[seed_idx] = perm_clusters_sums[idx_max]
        else:
            max_cluster_sums[seed_idx] = 0

        progress_bar.update(seed_idx + 1)

    return max_cluster_sums


'''

rewirite permutaiton test function to match the requirements of 
mix model.


'''
import numpy as np
import pandas as pd


def shuffle_distance_within_subId(X):
    """
    Shuffles the 'distance' column within each 'subId' group for each DataFrame in the ndarray X.

    Parameters:
    - X: ndarray, shape (1, n) containing pandas DataFrames with columns
         ['distance', 'subId', 'logic1', 'logic2', 'RT1', 'RT2'].

    Returns:
    - X: ndarray with shuffled 'distance' within each 'subId' group.
    """
    for i in range(X.shape[1]):
        # Access the DataFrame
        df = X[0, i]

        # Check if the DataFrame contains the expected columns
        if {'distance', 'subId'}.issubset(df.columns):
            # Shuffle 'distance' within each 'subId' group
            df['distance'] = df.groupby('subId')['distance'].transform(lambda x: np.random.permutation(x))

            # Update the DataFrame in the array
            X[0, i] = df
        else:
            print(f"DataFrame at index {i} does not have the required columns.")

    print("Shuffling completed.")
    return X


# Example usage:
# X = np.array([[pd.DataFrame(...), pd.DataFrame(...), ...]])  # Your ndarray of DataFrames
# X = shuffle_distance_within_subId(X)

'''
consider about not chage df in X?

'''

import numpy as np
import pandas as pd
import copy


def shuffle_distance_within_subId(X):
    """
    Shuffles the 'distance' column within each 'subId' group for each DataFrame in the ndarray X.

    Parameters:
    - X: ndarray, shape (1, n) containing pandas DataFrames with columns
         ['distance', 'subId', 'logic1', 'logic2', 'RT1', 'RT2'].

    Returns:
    - X_copy: ndarray with shuffled 'distance' within each 'subId' group.
    """
    # Create a deep copy of X to avoid modifying the original
    X_copy = copy.deepcopy(X)

    for i in range(X_copy.shape[1]):
        # Access the DataFrame
        df = X_copy[0, i]

        # Check if the DataFrame contains the expected columns
        if {'distance', 'subId'}.issubset(df.columns):
            # Shuffle 'distance' within each 'subId' group
            df['distance'] = df.groupby('subId')['distance'].transform(lambda x: np.random.permutation(x))

            # Update the DataFrame in the array
            X_copy[0, i] = df
        else:
            print(f"DataFrame at index {i} does not have the required columns.")

    print("Shuffling completed.")
    return X_copy


# Example usage:
# X = np.array([[pd.DataFrame(...), pd.DataFrame(...), ...]])  # Your ndarray of DataFrames
# X_shuffled = shuffle_distance_within_subId(X)
'''
rewrite the permutation test function to match the requirements of mix model .

improvement1: save and read the t values for every permutation. then you don't need to permutation every time.

'''


def _do_1samp_permutations(
        X,
        slices,
        threshold,
        tail,
        adjacency,
        stat_fun,
        max_step,
        include,
        partitions,
        t_power,
        # orders,
        sample_shape,
        buffer_size,
        # progress_bar,
):
    print("permutation test for mix model")

    '''
   what n_samp and n_vars mean?
   in my conditon, n_samp = 1, n_vars = n_chan * n_time 

   n_samp even not used in latter sentense..


    '''

    n_samp, n_vars = X.shape
    assert slices is None  # should be None for the 1 sample case

    # allocate space for output
    n_permutations = 100
    max_cluster_sums = np.empty(n_permutations, dtype=np.double)

    for perm_idx in range(n_permutations):
        # Shuffle the data using your provided shuffle function
        X_shuffled = shuffle_distance_within_subId(X)

        if buffer_size is None:
            # Recompute statistic on shuffled data
            t_obs_surr = stat_fun(X_shuffled)

        # The stat should have the same shape as the samples for no adj.
        if adjacency is None:
            t_obs_surr.shape = sample_shape

        # Find clusters on randomized stats
        out = _find_clusters(
            t_obs_surr,
            threshold=threshold,  # TFCE is just for cluster finding
            tail=tail,
            max_step=max_step,
            adjacency=adjacency,
            partitions=partitions,
            include=include,
            t_power=t_power,
        )
        perm_clusters_sums = out[1]
        if len(perm_clusters_sums) > 0:
            # get max with sign info
            idx_max = np.argmax(np.abs(perm_clusters_sums))
            max_cluster_sums[perm_idx] = perm_clusters_sums[idx_max]
        else:
            max_cluster_sums[perm_idx] = 0

        # progress_bar.update(perm_idx + 1)

    return max_cluster_sums


'''

improvement1: save and read the t values for every permutation. then you don't need to permutation every time.

'''


def _do_1samp_permutations(
        X,
        slices,
        threshold,
        tail,
        adjacency,
        stat_fun,
        max_step,
        include,
        partitions,
        t_power,
        sample_shape,
        buffer_size,
        n_permutations=1,
        save_t_values=True,  # New parameter to decide whether to save t-values
        t_values_filename='./results_permutation50_28subs'  # Filename to save t-values if needed
):
    """
    Perform a one-sample permutation test with the option to save t-values.

    Parameters:
    - X: The data to be permuted and tested.
    - slices: Should be None for the 1-sample case.
    - threshold, tail, adjacency, stat_fun, max_step, include, partitions, t_power: Various parameters for cluster testing.
    - sample_shape: The shape of the sample to ensure proper reshaping.
    - buffer_size: Buffer size for memory efficiency.
    - n_permutations: The number of permutations to perform.
    - save_t_values: Boolean flag to decide whether to save t-values.
    - t_values_filename: Filename to save the t-values if save_t_values is True.

    Returns:
    - max_cluster_sums: The maximum cluster sums for each permutation.
    - all_t_values: The t-values for each permutation if save_t_values is True.
    """

    print("Permutation test for mixed model")
    print("save t values.")

    n_samp, n_vars = X.shape
    assert slices is None  # should be None for the 1 sample case

    # Allocate space for output
    max_cluster_sums = np.empty(n_permutations, dtype=np.double)
    all_t_values = [] if save_t_values else None

    for perm_idx in range(n_permutations):
        # Shuffle the data using the provided shuffle function
        X_shuffled = shuffle_distance_within_subId(X)

        if buffer_size is None:
            # Recompute statistic on shuffled data
            t_obs_surr = stat_fun(X_shuffled)

            # Save the t-values if required
            if save_t_values:
                all_t_values.append(t_obs_surr)

        # Ensure t_obs_surr has the correct shape
        if adjacency is None:
            t_obs_surr.shape = sample_shape

        # Find clusters on randomized stats
        out = _find_clusters(
            t_obs_surr,
            threshold=threshold,
            tail=tail,
            max_step=max_step,
            adjacency=adjacency,
            partitions=partitions,
            include=include,
            t_power=t_power,
        )
        perm_clusters_sums = out[1]
        if len(perm_clusters_sums) > 0:
            # Get max with sign info
            idx_max = np.argmax(np.abs(perm_clusters_sums))
            max_cluster_sums[perm_idx] = perm_clusters_sums[idx_max]
        else:
            max_cluster_sums[perm_idx] = 0

        # progress_bar.update(perm_idx + 1) # Uncomment if using a progress bar

    # Save t-values to disk if required
    if save_t_values and t_values_filename is not None:
        np.save(t_values_filename, all_t_values)

    return max_cluster_sums


import os
import numpy as np

'''
improvement: 
I want to test ploting part of the code, and I need T or TFCE values.
read the t values from file if it exists. if not exists, then do permutation test and save t values.
for example, if I have already saved permutation test results,
then I don't need to do permutation test again. I can just read the t values from file.
so here t values means the t values for every permutation or the original t values.

 if the function find existing file, then, it don't need to shuffle and 
 calculate t again(t calculation is the most time consuming),
  we can just loop t values sets to calculate cluster??

considering the tfce calculation is not that time consuming
I should save t values for each permutation.

the permutation function include:

shuffle data 
calculate t 
correct data (tfce)  - also correct by find_cluster function..
test every original tfce

shuffle data
calculate t
find cluster and calculate cluster_T
test every original cluster_T

the common part of these two processes is first two process.
so I need to save t values before the third step..




'''


def _do_1samp_permutations(
        X,
        slices,
        threshold,
        tail,
        adjacency,
        stat_fun,
        max_step,
        include,
        partitions,
        t_power,
        sample_shape,
        buffer_size,
        n_permutations=1,
        save_t_values=True,  # Parameter to decide whether to save t-values
        t_values_filename='./results_permutation1.npy'  # Filename to save/load t-values
):
    """
    Perform a one-sample permutation test with the option to save and load t-values.

    Parameters:
    - X: The data to be permuted and tested.
    - slices: Should be None for the 1-sample case.
    - threshold, tail, adjacency, stat_fun, max_step, include, partitions, t_power: Various parameters for cluster testing.
    - sample_shape: The shape of the sample to ensure proper reshaping.
    - buffer_size: Buffer size for memory efficiency.
    - n_permutations: The number of permutations to perform.
    - save_t_values: Boolean flag to decide whether to save t-values.
    - t_values_filename: Filename to save/load the t-values if save_t_values is True.

    Returns:
    - max_cluster_sums: The maximum cluster sums for each permutation.
    - all_t_values: The t-values for each permutation if save_t_values is True.
    """

    print("Permutation test for mixed model")

    # t_values_filename = './results_permutation50_28subs.npy' # for test, change filename here..
    t_values_filename = t_values_filename

    # Check if t-values file exists
    # save_t_values=True,
    if save_t_values and os.path.exists(t_values_filename):
        print(f"Found existing results in {t_values_filename}. Loading saved t-values.")
        all_t_values = np.load(t_values_filename, allow_pickle=True)

        # Initialize max_cluster_sums based on loaded t-values
        max_cluster_sums = np.empty(len(all_t_values), dtype=np.double)

        # Calculate clusters using loaded t-values
        for perm_idx, t_obs_surr in enumerate(all_t_values):
            # Ensure t_obs_surr has the correct shape
            if adjacency is None:
                t_obs_surr.shape = sample_shape

            # Find clusters on saved t-values
            # if threshold is a dict,
            # the perm_clusters_sums is tfce corrected t-values in every point.
            out = _find_clusters(
                t_obs_surr,
                threshold=threshold,
                tail=tail,
                max_step=max_step,
                adjacency=adjacency,
                partitions=partitions,
                include=include,
                t_power=t_power,
            )

            # if threshold is a dict, the perm_clusters_sums is tfce corrected t-values in every point.
            # not the cluster_T.
            perm_clusters_sums = out[1]
            if len(perm_clusters_sums) > 0:
                # Get max with sign info
                idx_max = np.argmax(np.abs(perm_clusters_sums))
                max_cluster_sums[perm_idx] = perm_clusters_sums[idx_max]
            else:
                max_cluster_sums[perm_idx] = 0
        # return means that we don't need to do permutation test again and jump out the
        # permutation function.
        # if tfce, return max tfce corrected t-values for every permutation.
        return max_cluster_sums

    print("No saved results found. Starting permutation test and saving t-values.")

    n_samp, n_vars = X.shape
    assert slices is None  # should be None for the 1 sample case

    # Allocate space for output
    max_cluster_sums = np.empty(n_permutations, dtype=np.double)
    all_t_values = [] if save_t_values else None

    for perm_idx in range(n_permutations):
        # Shuffle the data using the provided shuffle function
        X_shuffled = shuffle_distance_within_subId(X)

        if buffer_size is None:
            # Recompute statistic on shuffled data
            t_obs_surr = stat_fun(X_shuffled)

            # Save the t-values if required
            if save_t_values:
                all_t_values.append(t_obs_surr)

        # Ensure t_obs_surr has the correct shape
        if adjacency is None:
            t_obs_surr.shape = sample_shape

        # Find clusters on randomized stats
        out = _find_clusters(
            t_obs_surr,
            threshold=threshold,
            tail=tail,
            max_step=max_step,
            adjacency=adjacency,
            partitions=partitions,
            include=include,
            t_power=t_power,
        )
        perm_clusters_sums = out[1]
        if len(perm_clusters_sums) > 0:
            # Get max with sign info
            idx_max = np.argmax(np.abs(perm_clusters_sums))
            max_cluster_sums[perm_idx] = perm_clusters_sums[idx_max]
        else:
            max_cluster_sums[perm_idx] = 0

    # Save permutations' t-values to disk if required
    print('save results_permutation.')
    print(t_values_filename)
    if save_t_values and t_values_filename is not None:
        np.save(t_values_filename, all_t_values)

    return max_cluster_sums


def bin_perm_rep(ndim, a=0, b=1):
    """Ndim permutations with repetitions of (a,b).

    Returns an array with all the possible permutations with repetitions of
    (0,1) in ndim dimensions.  The array is shaped as (2**ndim,ndim), and is
    ordered with the last index changing fastest.  For examble, for ndim=3:

    Examples
    --------
    >>> bin_perm_rep(3)
    array([[0, 0, 0],
           [0, 0, 1],
           [0, 1, 0],
           [0, 1, 1],
           [1, 0, 0],
           [1, 0, 1],
           [1, 1, 0],
           [1, 1, 1]])
    """
    # Create the leftmost column as 0,0,...,1,1,...
    nperms = 2 ** ndim
    perms = np.empty((nperms, ndim), type(a))
    perms.fill(a)
    half_point = nperms // 2
    perms[half_point:, 0] = b
    # Fill the rest of the table by sampling the previous column every 2 items
    for j in range(1, ndim):
        half_col = perms[::2, j - 1]
        perms[:half_point, j] = half_col
        perms[half_point:, j] = half_col
    # This is equivalent to something like:
    # orders = [np.fromiter(np.binary_repr(s + 1, ndim), dtype=int)
    #           for s in np.arange(2 ** ndim)]
    return perms


def _get_1samp_orders(n_samples, n_permutations, tail, rng):
    """Get the 1samp orders."""
    max_perms = 2 ** (n_samples - (tail == 0)) - 1
    extra = ""
    if isinstance(n_permutations, str):
        if n_permutations != "all":
            raise ValueError('n_permutations as a string must be "all"')
        n_permutations = max_perms
    n_permutations = int(n_permutations)
    if max_perms < n_permutations:
        # omit first perm b/c accounted for in H0.append() later;
        # convert to binary array representation
        extra = " (exact test)"
        orders = bin_perm_rep(n_samples)[1: max_perms + 1]
    elif n_samples <= 20:  # fast way to do it for small(ish) n_samples
        orders = rng.choice(max_perms, n_permutations - 1, replace=False)
        orders = [
            np.fromiter(np.binary_repr(s + 1, n_samples), dtype=int) for s in orders
        ]
    else:  # n_samples >= 64
        # Here we can just use the hash-table (w/collision detection)
        # functionality of a dict to ensure uniqueness
        orders = np.zeros((n_permutations - 1, n_samples), int)
        hashes = {}
        ii = 0
        # in the symmetric case, we should never flip one of the subjects
        # to prevent positive/negative equivalent collisions
        use_samples = n_samples - (tail == 0)
        while ii < n_permutations - 1:
            signs = tuple((rng.uniform(size=use_samples) < 0.5).astype(int))
            if signs not in hashes:
                orders[ii, :use_samples] = signs
                if tail == 0 and rng.uniform() < 0.5:
                    # To undo the non-flipping of the last subject in the
                    # tail == 0 case, half the time we use the positive
                    # last subject, half the time negative last subject
                    orders[ii] = 1 - orders[ii]
                hashes[signs] = None
                ii += 1
    return orders, n_permutations, extra


from tqdm import tqdm

print("reload ok")
# change the code in the following function
'''
this function is really complex
it include t calculation
find cluster OR TFCE correction
permutation test(iterate t calculation and tfce correction)

'''


def _permutation_cluster_test(
        X,
        threshold,
        n_permutations,
        tail,
        stat_fun,
        adjacency,
        n_jobs,
        seed,
        max_step,
        exclude,
        step_down_p,
        t_power,
        out_type,
        check_disjoint,
        buffer_size,
):
    """Aux Function.

    Note. X is required to be a list. Depending on the length of X
    either a 1 sample t-test or an F test / more sample permutation scheme
    is elicited.
    """
    _check_option("out_type", out_type, ["mask", "indices"])
    _check_option("tail", tail, [-1, 0, 1])
    if not isinstance(threshold, dict):
        threshold = float(threshold)
        if (
                tail < 0
                and threshold > 0
                or tail > 0
                and threshold < 0
                or tail == 0
                and threshold < 0
        ):
            raise ValueError(
                "incompatible tail and threshold signs, got "
                "%s and %s" % (tail, threshold)
            )

    # check dimensions for each group in X (a list at this stage).
    # check dimensions for each group in X (a list at this stage).
    # I don't know why this decrease the dimensionality of X.from (1,91,63) to (91,63)
    X = [x[:, np.newaxis] if x.ndim == 1 else x for x in X]
    # X = [np.expand_dims(x, axis=0) for x in X] # Adds a new axis at the beginning
    # Assuming X is a list of numpy arrays of various dimensions
    # for convinience to your code and official code.
    # in your code, the dimensionality of X is (91,63) here, neeed to add a new axis at the beginning.
    X = [np.expand_dims(x, axis=0) if x.ndim < 3 else x for x in X]
    n_samples = X[0].shape[0]
    n_times = X[0].shape[1]

    sample_shape = X[0].shape[1:]
    for x in X:
        if x.shape[1:] != sample_shape:
            raise ValueError("All samples mush have the same size")

    # flatten the last dimensions in case the data is high dimensional
    X = [np.reshape(x, (x.shape[0], -1)) for x in X]
    n_tests = X[0].shape[1]

    if adjacency is not None and adjacency is not False:
        adjacency = _setup_adjacency(adjacency, n_tests, n_times)

    if (exclude is not None) and not exclude.size == n_tests:
        raise ValueError("exclude must be the same shape as X[0]")

    # Step 1: Calculate t-stat for original data
    # -----------------------------------------------------------——--
    # *X means unpacking the list X, which is a list of numpy arrays,
    # in your condition, X is a list of numpy arrays with shape (1, 91,63)
    # just got one element in the list， so in your case, X is the data. you  don't neee *
    t_obs = stat_fun(*X)
    _validate_type(t_obs, np.ndarray, "return value of stat_fun")
    logger.info("stat_fun(H1): min=%f max=%f" % (np.min(t_obs), np.max(t_obs)))

    # test if stat_fun treats variables independently
    # buffer_size if helpful for large data.
    if buffer_size is not None:
        t_obs_buffer = np.zeros_like(t_obs)
        for pos in range(0, n_tests, buffer_size):
            t_obs_buffer[pos: pos + buffer_size] = stat_fun(
                *[x[:, pos: pos + buffer_size] for x in X]
            )

        if not np.all(t_obs == t_obs_buffer):
            warn(
                "Provided stat_fun does not treat variables independently. "
                "Setting buffer_size to None."
            )
            buffer_size = None

    # The stat should have the same shape as the samples for no adj.
    if t_obs.size != np.prod(sample_shape):
        raise ValueError(
            "t_obs.shape %s provided by stat_fun %s is not "
            "compatible with the sample shape %s"
            % (t_obs.shape, stat_fun, sample_shape)
        )
    if adjacency is None or adjacency is False:
        t_obs.shape = sample_shape

    if exclude is not None:
        include = np.logical_not(exclude)
    else:
        include = None

    # determine if adjacency itself can be separated into disjoint sets
    if check_disjoint is True and (adjacency is not None and adjacency is not False):
        partitions = _get_partitions_from_adjacency(adjacency, n_times)
    else:
        partitions = None
    logger.info("Running initial clustering …")
    # maybe I should find out the principle of this function
    # I just need to call this function in your own permutation function..
    # if threshold is a dict,
    # the clusters is tfce corrected t-values in every point. not cluster!
    out = _find_clusters(
        t_obs,
        threshold,
        tail,
        adjacency,
        max_step=max_step,
        include=include,
        partitions=partitions,
        t_power=t_power,
        show_info=True,
    )
    clusters, cluster_stats = out

    # The stat should have the same shape as the samples
    t_obs.shape = sample_shape

    # keep the original raw t
    t_obs_orig = t_obs.copy()

    # For TFCE, return the "adjusted" statistic instead of raw scores
    if isinstance(threshold, dict):
        t_obs = cluster_stats.reshape(t_obs.shape) * np.sign(t_obs)

    logger.info(f"Found {len(clusters)} cluster{_pl(clusters)}")

    # convert clusters to old format
    if adjacency is not None and adjacency is not False:
        # our algorithms output lists of indices by default
        if out_type == "mask":
            clusters = _cluster_indices_to_mask(clusters, n_tests)
    else:
        # ndimage outputs slices or boolean masks by default
        if out_type == "indices":
            clusters = _cluster_mask_to_indices(clusters, t_obs.shape)

    # convert our seed to orders
    # check to see if we can do an exact test
    # (for a two-tailed test, we can exploit symmetry to just do half)
    extra = ""
    rng = check_random_state(seed)
    del seed
    if len(X) == 1:  # 1-sample test
        do_perm_func = _do_1samp_permutations
        X_full = X[0]
        slices = None
        orders, n_permutations, extra = _get_1samp_orders(
            n_samples, n_permutations, tail, rng
        )
    else:
        n_permutations = int(n_permutations)
        do_perm_func = _do_permutations
        X_full = np.concatenate(X, axis=0)
        n_samples_per_condition = [x.shape[0] for x in X]
        splits_idx = np.append([0], np.cumsum(n_samples_per_condition))
        slices = [slice(splits_idx[k], splits_idx[k + 1]) for k in range(len(X))]
        orders = [rng.permutation(len(X_full)) for _ in range(n_permutations - 1)]
    del rng
    parallel, my_do_perm_func, n_jobs = parallel_func(  # parallel permutation process.
        do_perm_func, n_jobs, verbose=False
    )

    if len(clusters) == 0:
        warn("No clusters found, returning empty H0, clusters, and cluster_pv")
        return t_obs, np.array([]), np.array([]), np.array([])

    # Step 2: If we have some clusters, repeat process on permuted data
    # -------------------------------------------------------------------
    # Step 3: repeat permutations for step-down-in-jumps procedure
    n_removed = 1  # number of new clusters added
    total_removed = 0
    step_down_include = None  # start out including all points
    n_step_downs = 0

    # while n_removed > 0:
    #     # actually do the clustering for each partition
    #     if include is not None:
    #         if step_down_include is not None:
    #             this_include = np.logical_and(include, step_down_include)
    #         else:
    #             this_include = include
    #     else:
    #         this_include = step_down_include
    #
    #     # with ProgressBar(
    #     #     iterable=range(len(orders)), mesg=f"Permuting{extra}"
    #     # ) as progress_bar:
    #     #     H0 = parallel(
    #     #         my_do_perm_func(
    #     #             X_full,
    #     #             slices,
    #     #             threshold,
    #     #             tail,
    #     #             adjacency,
    #     #             stat_fun,
    #     #             max_step,
    #     #             this_include,
    #     #             partitions,
    #     #             t_power,
    #     #             order,
    #     #             sample_shape,
    #     #             buffer_size,
    #     #             progress_bar.subset(idx),
    #     #         )
    #     #         for idx, order in split_list(orders, n_jobs, idx=True)
    #     #     )
    #
    #     # Custom progress bar using tqdm
    #     with tqdm(total=n_permutations, desc=f"Permuting{extra}") as progress_bar:
    #         H0 = parallel(
    #              my_do_perm_func(
    #                 X_full,
    #                 slices,
    #                 threshold,
    #                 tail,
    #                 adjacency,
    #                 stat_fun,
    #                 max_step,
    #                 this_include,
    #                 partitions,
    #                 t_power,
    #                 order,
    #                 sample_shape,
    #                 buffer_size,
    #                 progress_bar
    #             )
    #             for order in range(n_permutations)
    #         )
    #         progress_bar.update(n_permutations)
    #
    #
    #
    #
    #     # include original (true) ordering
    #     if tail == -1:  # up tail
    #         orig = cluster_stats.min()
    #     elif tail == 1:
    #         orig = cluster_stats.max()
    #     else:
    #         orig = abs(cluster_stats).max()
    #     H0.insert(0, [orig])
    #     H0 = np.concatenate(H0)
    #     logger.debug("Computing cluster p-values")
    #     cluster_pv = _pval_from_histogram(cluster_stats, H0, tail)
    #
    #     # figure out how many new ones will be removed for step-down
    #     to_remove = np.where(cluster_pv < step_down_p)[0]
    #     n_removed = to_remove.size - total_removed
    #     total_removed = to_remove.size
    #     step_down_include = np.ones(n_tests, dtype=bool)
    #     for ti in to_remove:
    #         step_down_include[clusters[ti]] = False
    #     if adjacency is None and adjacency is not False:
    #         step_down_include.shape = sample_shape
    #     n_step_downs += 1
    #     if step_down_p > 0:
    #         a_text = "additional " if n_step_downs > 1 else ""
    #         logger.info(
    #             "Step-down-in-jumps iteration #%i found %i %s"
    #             "cluster%s to exclude from subsequent iterations"
    #             % (n_step_downs, n_removed, a_text, _pl(n_removed))
    #         )

    # while n_removed > 0:
    # actually do the clustering for each partition
    if include is not None:
        if step_down_include is not None:
            this_include = np.logical_and(include, step_down_include)
        else:
            this_include = include
    else:
        this_include = step_down_include

    # with ProgressBar(
    #     iterable=range(len(orders)), mesg=f"Permuting{extra}"
    # ) as progress_bar:
    #     H0 = parallel(
    #         my_do_perm_func(
    #             X_full,
    #             slices,
    #             threshold,
    #             tail,
    #             adjacency,
    #             stat_fun,
    #             max_step,
    #             this_include,
    #             partitions,
    #             t_power,
    #             order,
    #             sample_shape,
    #             buffer_size,
    #             progress_bar.subset(idx),
    #         )
    #         for idx, order in split_list(orders, n_jobs, idx=True)
    #     )

    # Custom progress bar using tqdm
    # with tqdm(total=n_permutations, desc=f"Permuting{extra}") as progress_bar:
    #     H0 = parallel(
    #          my_do_perm_func(
    #             X_full,
    #             slices,
    #             threshold,
    #             tail,
    #             adjacency,
    #             stat_fun,
    #             max_step,
    #             this_include,
    #             partitions,
    #             t_power,
    #             order,
    #             sample_shape,
    #             buffer_size,
    #             progress_bar
    #         )
    #         for order in range(n_permutations)
    #     )
    #     progress_bar.update(n_permutations)

    H0 = _do_1samp_permutations(
        X_full,
        slices,
        threshold,
        tail,
        adjacency,
        stat_fun,
        max_step,
        this_include,
        partitions,
        t_power,
        # order,
        sample_shape,
        buffer_size,
        # progress_bar
        save_t_values=True,  # New parameter to decide whether to save t-values
        # I prefer change the name in function definition not here.
        # t_values_filename='./results_permutation'  # Filename to save t-values if needed

    )

    cluster_pv = _pval_from_histogram(cluster_stats, H0, tail)

    # The clusters should have the same shape as the samples
    clusters = _reshape_clusters(clusters, sample_shape)
    return t_obs, clusters, cluster_stats, cluster_pv, H0, t_obs_orig


def _check_fun(X, stat_fun, threshold, tail=0, kind="within"):
    """Check the stat_fun and threshold values."""
    if kind == "within":
        if threshold is None:
            if stat_fun is not None and stat_fun is not ttest_1samp_no_p:
                warn(
                    "Automatic threshold is only valid for stat_fun=None "
                    "(or ttest_1samp_no_p), got %s" % (stat_fun,)
                )
            p_thresh = 0.05 / (1 + (tail == 0))
            n_samples = len(X)
            threshold = -tstat.ppf(p_thresh, n_samples - 1)
            if np.sign(tail) < 0:
                threshold = -threshold
            logger.info("Using a threshold of {:.6f}".format(threshold))
        stat_fun = ttest_1samp_no_p if stat_fun is None else stat_fun
    else:
        assert kind == "between"
        if threshold is None:
            if stat_fun is not None and stat_fun is not f_oneway:
                warn(
                    "Automatic threshold is only valid for stat_fun=None "
                    "(or f_oneway), got %s" % (stat_fun,)
                )
            elif tail != 1:
                warn('Ignoring argument "tail", performing 1-tailed F-test')
            p_thresh = 0.05
            dfn = len(X) - 1
            dfd = np.sum([len(x) for x in X]) - len(X)
            threshold = fstat.ppf(1.0 - p_thresh, dfn, dfd)
            logger.info("Using a threshold of {:.6f}".format(threshold))
        stat_fun = f_oneway if stat_fun is None else stat_fun
    return stat_fun, threshold


@verbose
def permutation_cluster_test(
        X,
        threshold=None,
        n_permutations=1024,
        tail=0,
        stat_fun=None,
        adjacency=None,
        n_jobs=None,
        seed=None,
        max_step=1,
        exclude=None,
        step_down_p=0,
        t_power=1,
        out_type="indices",
        check_disjoint=False,
        buffer_size=1000,
        verbose=None,
):
    """Cluster-level statistical permutation test.

    For a list of :class:`NumPy arrays <numpy.ndarray>` of data,
    calculate some statistics corrected for multiple comparisons using
    permutations and cluster-level correction. Each element of the list ``X``
    should contain the data for one group of observations (e.g., 2D arrays for
    time series, 3D arrays for time-frequency power values). Permutations are
    generated with random partitions of the data. For details, see
    :footcite:p:`MarisOostenveld2007,Sassenhagen2019`.

    Parameters
    ----------
    X : list of array, shape (n_observations, p[, q][, r])
        The data to be clustered. Each array in ``X`` should contain the
        observations for one group. The first dimension of each array is the
        number of observations from that group; remaining dimensions comprise
        the size of a single observation. For example if ``X = [X1, X2]``
        with ``X1.shape = (20, 50, 4)`` and ``X2.shape = (17, 50, 4)``, then
        ``X`` has 2 groups with respectively 20 and 17 observations in each,
        and each data point is of shape ``(50, 4)``. Note: that the
        *last dimension* of each element of ``X`` should correspond to the
        dimension represented in the ``adjacency`` parameter
        (e.g., spectral data should be provided as
        ``(observations, frequencies, channels/vertices)``).
    %(threshold_clust_f)s
    %(n_permutations_clust_int)s
    %(tail_clust)s
    %(stat_fun_clust_f)s
    %(adjacency_clust_n)s
    %(n_jobs)s
    %(seed)s
    %(max_step_clust)s
    %(exclude_clust)s
    %(step_down_p_clust)s
    %(f_power_clust)s
    %(out_type_clust)s
    %(check_disjoint_clust)s
    %(buffer_size_clust)s
    %(verbose)s

    Returns
    -------
    F_obs : array, shape (p[, q][, r])
        Statistic (F by default) observed for all variables.
    clusters : list
        List type defined by out_type above.
    cluster_pv : array
        P-value for each cluster.
    H0 : array, shape (n_permutations,)
        Max cluster level stats observed under permutation.

    Notes
    -----
    %(threshold_clust_f_notes)s

    References
    ----------
    .. footbibliography::
    """
    stat_fun, threshold = _check_fun(X, stat_fun, threshold, tail, "between")
    return _permutation_cluster_test(
        X=X,
        threshold=threshold,
        n_permutations=n_permutations,
        tail=tail,
        stat_fun=stat_fun,
        adjacency=adjacency,
        n_jobs=n_jobs,
        seed=seed,
        max_step=max_step,
        exclude=exclude,
        step_down_p=step_down_p,
        t_power=t_power,
        out_type=out_type,
        check_disjoint=check_disjoint,
        buffer_size=buffer_size,
    )


@verbose
def permutation_cluster_1samp_test(
        X,
        threshold=None,
        n_permutations=1024,
        tail=0,
        stat_fun=None,
        adjacency=None,
        n_jobs=None,
        seed=None,
        max_step=1,
        exclude=None,
        step_down_p=0,
        t_power=1,
        out_type="indices",
        check_disjoint=False,
        buffer_size=1000,
        verbose=None,
):
    """Non-parametric cluster-level paired t-test.

    For details, see :footcite:p:`MarisOostenveld2007,Sassenhagen2019`.

    Parameters
    ----------
    X : array, shape (n_observations, p[, q][, r])
        The data to be clustered. The first dimension should correspond to the
        difference between paired samples (observations) in two conditions.
        The subarrays ``X[k]`` can be 1D (e.g., time series), 2D (e.g.,
        time series over channels), or 3D (e.g., time-frequencies over
        channels) associated with the kth observation. For spatiotemporal data,
        see also :func:`mne.stats.spatio_temporal_cluster_1samp_test`.
    %(threshold_clust_t)s
    %(n_permutations_clust_all)s
    %(tail_clust)s
    %(stat_fun_clust_t)s
    %(adjacency_clust_1)s
    %(n_jobs)s
    %(seed)s
    %(max_step_clust)s
    %(exclude_clust)s
    %(step_down_p_clust)s
    %(t_power_clust)s
    %(out_type_clust)s
    %(check_disjoint_clust)s
    %(buffer_size_clust)s
    %(verbose)s

    Returns
    -------
    t_obs : array, shape (p[, q][, r])
        T-statistic observed for all variables.
    clusters : list
        List type defined by out_type above.
    cluster_pv : array
        P-value for each cluster.
    H0 : array, shape (n_permutations,)
        Max cluster level stats observed under permutation.

    Notes
    -----
    From an array of paired observations, e.g. a difference in signal
    amplitudes or power spectra in two conditions, calculate if the data
    distributions in the two conditions are significantly different.
    The procedure uses a cluster analysis with permutation test
    for calculating corrected p-values. Randomized data are generated with
    random sign flips. See :footcite:`MarisOostenveld2007` for more
    information.

    Because a 1-sample t-test on the difference in observations is
    mathematically equivalent to a paired t-test, internally this function
    computes a 1-sample t-test (by default) and uses sign flipping (always)
    to perform permutations. This might not be suitable for the case where
    there is truly a single observation under test; see :ref:`disc-stats`.
    %(threshold_clust_t_notes)s

    If ``n_permutations`` exceeds the maximum number of possible permutations
    given the number of observations, then ``n_permutations`` and ``seed``
    will be ignored since an exact test (full permutation test) will be
    performed (this is the case when
    ``n_permutations >= 2 ** (n_observations - (tail == 0))``).

    If no initial clusters are found because all points in the true
    distribution are below the threshold, then ``clusters``, ``cluster_pv``,
    and ``H0`` will all be empty arrays.

    References
    ----------
    .. footbibliography::
    """
    stat_fun, threshold = _check_fun(X, stat_fun, threshold, tail)
    return _permutation_cluster_test(
        X=[X],
        threshold=threshold,
        n_permutations=n_permutations,
        tail=tail,
        stat_fun=stat_fun,
        adjacency=adjacency,
        n_jobs=n_jobs,
        seed=seed,
        max_step=max_step,
        exclude=exclude,
        step_down_p=step_down_p,
        t_power=t_power,
        out_type=out_type,
        check_disjoint=check_disjoint,
        buffer_size=buffer_size,
    )


@verbose
def spatio_temporal_cluster_1samp_test(
        X,
        threshold=None,
        n_permutations=1024,
        tail=0,
        stat_fun=None,
        adjacency=None,
        n_jobs=None,
        seed=None,
        max_step=1,
        spatial_exclude=None,
        step_down_p=0,
        t_power=1,
        out_type="indices",
        check_disjoint=False,
        buffer_size=1000,
        verbose=None,
):
    """Non-parametric cluster-level paired t-test for spatio-temporal data.

    This function provides a convenient wrapper for
    :func:`mne.stats.permutation_cluster_1samp_test`, for use with data
    organized in the form (observations × time × space),
    (observations × frequencies × space), or optionally
    (observations × time × frequencies × space). For details, see
    :footcite:p:`MarisOostenveld2007,Sassenhagen2019`.

    Parameters
    ----------
    X : array, shape (n_observations, p[, q], n_vertices)
        The data to be clustered. The first dimension should correspond to the
        difference between paired samples (observations) in two conditions.
        The second, and optionally third, dimensions correspond to the
        time or time-frequency data. And, the last dimension should be spatial.
    %(threshold_clust_t)s
    %(n_permutations_clust_all)s
    %(tail_clust)s
    %(stat_fun_clust_t)s
    %(adjacency_clust_st1)s
    %(n_jobs)s
    %(seed)s
    %(max_step_clust)s
    spatial_exclude : list of int or None
        List of spatial indices to exclude from clustering.
    %(step_down_p_clust)s
    %(t_power_clust)s
    %(out_type_clust)s
    %(check_disjoint_clust)s
    %(buffer_size_clust)s
    %(verbose)s

    Returns
    -------
    t_obs : array, shape (p[, q], n_vertices)
        T-statistic observed for all variables.
    clusters : list
        List type defined by out_type above.
    cluster_pv : array
        P-value for each cluster.
    H0 : array, shape (n_permutations,)
        Max cluster level stats observed under permutation.

    Notes
    -----
    %(threshold_clust_t_notes)s

    References
    ----------
    .. footbibliography::
    """
    # convert spatial_exclude before passing on if necessary
    if spatial_exclude is not None:
        exclude = _st_mask_from_s_inds(
            np.prod(X.shape[1:-1]), X.shape[-1], spatial_exclude, True
        )
    else:
        exclude = None
    return permutation_cluster_1samp_test(
        X,
        threshold=threshold,
        stat_fun=stat_fun,
        tail=tail,
        n_permutations=n_permutations,
        adjacency=adjacency,
        n_jobs=n_jobs,
        seed=seed,
        max_step=max_step,
        exclude=exclude,
        step_down_p=step_down_p,
        t_power=t_power,
        out_type=out_type,
        check_disjoint=check_disjoint,
        buffer_size=buffer_size,
    )


@verbose
def spatio_temporal_cluster_test(
        X,
        threshold=None,
        n_permutations=1024,
        tail=0,
        stat_fun=None,
        adjacency=None,
        n_jobs=None,
        seed=None,
        max_step=1,
        spatial_exclude=None,
        step_down_p=0,
        t_power=1,
        out_type="indices",
        check_disjoint=False,
        buffer_size=1000,
        verbose=None,
):
    """Non-parametric cluster-level test for spatio-temporal data.

    This function provides a convenient wrapper for
    :func:`mne.stats.permutation_cluster_test`, for use with data
    organized in the form (observations × time × space),
    (observations × time × space), or optionally
    (observations × time × frequencies × space). For more information,
    see :footcite:p:`MarisOostenveld2007,Sassenhagen2019`.

    Parameters
    ----------
    X : list of array, shape (n_observations, p[, q], n_vertices)
        The data to be clustered. Each array in ``X`` should contain the
        observations for one group. The first dimension of each array is the
        number of observations from that group (and may vary between groups).
        The second, and optionally third, dimensions correspond to the
        time or time-frequency data. And, the last dimension should be spatial.
        All dimensions except the first should match across all groups.
    %(threshold_clust_f)s
    %(n_permutations_clust_int)s
    %(tail_clust)s
    %(stat_fun_clust_f)s
    %(adjacency_clust_stn)s
    %(n_jobs)s
    %(seed)s
    %(max_step_clust)s
    spatial_exclude : list of int or None
        List of spatial indices to exclude from clustering.
    %(step_down_p_clust)s
    %(f_power_clust)s
    %(out_type_clust)s
    %(check_disjoint_clust)s
    %(buffer_size_clust)s
    %(verbose)s

    Returns
    -------
    F_obs : array, shape (p[, q], n_vertices)
        Statistic (F by default) observed for all variables.
    clusters : list
        List type defined by out_type above.
    cluster_pv: array
        P-value for each cluster.
    H0 : array, shape (n_permutations,)
        Max cluster level stats observed under permutation.

    Notes
    -----
    %(threshold_clust_f_notes)s

    References
    ----------
    .. footbibliography::
    """
    # convert spatial_exclude before passing on if necessary
    if spatial_exclude is not None:
        exclude = _st_mask_from_s_inds(
            np.prod(X[0].shape[1:-1]), X[0].shape[-1], spatial_exclude, True
        )
    else:
        exclude = None
    return permutation_cluster_test(
        X,
        threshold=threshold,
        stat_fun=stat_fun,
        tail=tail,
        n_permutations=n_permutations,
        adjacency=adjacency,
        n_jobs=n_jobs,
        seed=seed,
        max_step=max_step,
        exclude=exclude,
        step_down_p=step_down_p,
        t_power=t_power,
        out_type=out_type,
        check_disjoint=check_disjoint,
        buffer_size=buffer_size,
    )


def _st_mask_from_s_inds(n_times, n_vertices, vertices, set_as=True):
    """Compute mask to apply to a spatio-temporal adjacency matrix.

    This can be used to include (or exclude) certain spatial coordinates.
    This is useful for excluding certain regions from analysis (e.g.,
    medial wall vertices).

    Parameters
    ----------
    n_times : int
        Number of time points.
    n_vertices : int
        Number of spatial points.
    vertices : list or array of int
        Vertex numbers to set.
    set_as : bool
        If True, all points except "vertices" are set to False (inclusion).
        If False, all points except "vertices" are set to True (exclusion).

    Returns
    -------
    mask : array of bool
        A (n_times * n_vertices) array of boolean values for masking
    """
    mask = np.zeros((n_times, n_vertices), dtype=bool)
    mask[:, vertices] = True
    mask = mask.ravel()
    if set_as is False:
        mask = np.logical_not(mask)
    return mask


@verbose
def _get_partitions_from_adjacency(adjacency, n_times, verbose=None):
    """Specify disjoint subsets (e.g., hemispheres) based on adjacency."""
    if isinstance(adjacency, list):
        test = np.ones(len(adjacency))
        test_adj = np.zeros((len(adjacency), len(adjacency)), dtype="bool")
        for vi in range(len(adjacency)):
            test_adj[adjacency[vi], vi] = True
        test_adj = sparse.coo_matrix(test_adj, dtype="float")
    else:
        test = np.ones(adjacency.shape[0])
        test_adj = adjacency

    part_clusts = _find_clusters(test, 0, 1, test_adj)[0]
    if len(part_clusts) > 1:
        logger.info("%i disjoint adjacency sets found" % len(part_clusts))
        partitions = np.zeros(len(test), dtype="int")
        for ii, pc in enumerate(part_clusts):
            partitions[pc] = ii
        if isinstance(adjacency, list):
            partitions = np.tile(partitions, n_times)
    else:
        logger.info("No disjoint adjacency sets found")
        partitions = None

    return partitions


def _reshape_clusters(clusters, sample_shape):
    """Reshape cluster masks or indices to be of the correct shape."""
    # format of the bool mask and indices are ndarrays
    if len(clusters) > 0 and isinstance(clusters[0], np.ndarray):
        if clusters[0].dtype == bool:  # format of mask
            clusters = [c.reshape(sample_shape) for c in clusters]
        else:  # format of indices
            clusters = [np.unravel_index(c, sample_shape) for c in clusters]
    return clusters


def summarize_clusters_stc(
        clu, p_thresh=0.05, tstep=1.0, tmin=0, subject="fsaverage", vertices=None
):
    """Assemble summary SourceEstimate from spatiotemporal cluster results.

    This helps visualizing results from spatio-temporal-clustering
    permutation tests.

    Parameters
    ----------
    clu : tuple
        The output from clustering permutation tests.
    p_thresh : float
        The significance threshold for inclusion of clusters.
    tstep : float
        The time step between samples of the original :class:`STC
        <mne.SourceEstimate>`, in seconds (i.e., ``1 / stc.sfreq``). Defaults
        to ``1``, which will yield a colormap indicating cluster duration
        measured in *samples* rather than *seconds*.
    tmin : float | int
        The time of the first sample.
    subject : str
        The name of the subject.
    vertices : list of array | instance of SourceSpaces | None
        The vertex numbers associated with the source space locations. Defaults
        to None. If None, equals ``[np.arange(10242), np.arange(10242)]``.
        Can also be an instance of SourceSpaces to get vertex numbers from.

        .. versionchanged:: 0.21
           Added support for SourceSpaces.

    Returns
    -------
    out : instance of SourceEstimate
        A summary of the clusters. The first time point in this SourceEstimate
        object is the summation of all the clusters. Subsequent time points
        contain each individual cluster. The magnitude of the activity
        corresponds to the duration spanned by the cluster (duration units are
        determined by ``tstep``).

        .. versionchanged:: 0.21
           Added support for volume and mixed source estimates.
    """
    _validate_type(vertices, (None, list, SourceSpaces), "vertices")
    if vertices is None:
        vertices = [np.arange(10242), np.arange(10242)]
        klass = SourceEstimate
    elif isinstance(vertices, SourceSpaces):
        klass = dict(
            surface=SourceEstimate, volume=VolSourceEstimate, mixed=MixedSourceEstimate
        )[vertices.kind]
        vertices = [s["vertno"] for s in vertices]
    else:
        klass = {1: VolSourceEstimate, 2: SourceEstimate}.get(
            len(vertices), MixedSourceEstimate
        )
    n_vertices_need = sum(len(v) for v in vertices)

    t_obs, clusters, clu_pvals, _ = clu
    n_times, n_vertices = t_obs.shape
    if n_vertices != n_vertices_need:
        raise ValueError(
            f"Number of cluster vertices ({n_vertices}) did not match the "
            f"provided vertices ({n_vertices_need})"
        )
    good_cluster_inds = np.where(clu_pvals < p_thresh)[0]

    #  Build a convenient representation of each cluster, where each
    #  cluster becomes a "time point" in the SourceEstimate
    if len(good_cluster_inds) == 0:
        raise RuntimeError(
            "No significant clusters available. Please adjust "
            "your threshold or check your statistical "
            "analysis."
        )
    data = np.zeros((n_vertices, n_times))
    data_summary = np.zeros((n_vertices, len(good_cluster_inds) + 1))
    for ii, cluster_ind in enumerate(good_cluster_inds):
        data.fill(0)
        t_inds, v_inds = clusters[cluster_ind]
        data[v_inds, t_inds] = t_obs[t_inds, v_inds]
        # Store a nice visualization of the cluster by summing across time
        data_summary[:, ii + 1] = np.sum(_sum_cluster_data(data, tstep), axis=1)
        # Make the first "time point" a sum across all clusters for easy
        # visualization
    data_summary[:, 0] = np.sum(data_summary, axis=1)

    return klass(data_summary, vertices, tmin, tstep, subject)