import itertools
from typing import Iterator

import torch


def permutation_group(x: torch.Tensor) -> torch.tensor:
    indices = range(x.shape[-1])
    permuted_indices = [list(p) for p in itertools.permutations(indices)]
    permuted_x = x[..., permuted_indices]
    # permuted_x is a tensor of shape (..., n, G, d)
    # Reorder the dimensions to (..., G, n, d)
    dim_indices = list(range(permuted_x.dim()))
    dim_indices[-2], dim_indices[-3] = dim_indices[-3], dim_indices[-2]
    return permuted_x.permute(*dim_indices)


def block_permutation_group(x: torch.Tensor, block_size: int) -> torch.tensor:
    if x.shape[-1] % block_size != 0:
        raise ValueError(
            "Last dimension of x must be a multiple of block size"
            f" (got {x.shape[-1]} and {block_size} respectively)."
        )

    block_indices = [
        range(i, i + block_size) for i in range(0, x.shape[-1], block_size)
    ]
    # Use itertools.chain to flatten the list of lists generated by the permutation
    permuted_indices = [
        list(itertools.chain.from_iterable(p))
        for p in itertools.permutations(block_indices)
    ]
    permuted_x = x[..., permuted_indices]
    # permuted_x is a tensor of shape (..., n, G, d)
    # Reorder the dimensions to (..., G, n, d)
    dim_indices = list(range(permuted_x.dim()))
    dim_indices[-2], dim_indices[-3] = dim_indices[-3], dim_indices[-2]
    return permuted_x.permute(*dim_indices)
