from typing import Iterable
import torch


def permute_qubits_tensor(
        permutation: Iterable[int],
        state_tensor: torch.Tensor,
        num_qubits: int,
        contiguous_output: bool = False
):
    """Permute qubits for a state in tensor layout.

    If `state_tensor` has batch dimensions, then the same permutation is
    applied to all batch entries. In particular, this function does not
    permute batch indices (use torch.permute for that).

    Args:
        permutation: Sequence of integers defining a permutation among qubits.
        state_tensor: State in vector layout.
        num_qubits: Number of qubits for the state.
        contiguous_output: When True, returns a contiguous copy of the result.
    """
    # non_qubit_dims includes the complex dimension in the complex case.
    non_qubit_dims = state_tensor.dim() - num_qubits
    identity_perm = tuple(i for i in range(non_qubit_dims))

    permutation = tuple(i + non_qubit_dims for i in permutation)

    state_tensor = torch.permute(state_tensor, identity_perm + permutation)
    if not contiguous_output:
        return state_tensor
    else:
        return state_tensor.contiguous()