from typing import Union
from typing import List
from typing import Tuple
import torch


def get_qubit_indices(
        index: Union[int, List[int], Tuple[int, ...], torch.Tensor],
        state_tensor: torch.Tensor,
        num_qubits: int
):
    """Convert qubit indices 0, ..., n-1 to correct PyTorch tensor indices.

    Consider a state with two qubits in tensor layout and with one batch
    dimension. The size is (7, 2, 2) if the batch length is 7.
    The first dimension (with torch index 0) refers to the batch
    while the last two refer to qubits. If we assign the first and second
    qubits "qubit indices" 0 and 1 respectively, then the torch indices of
    (0 and 1) are (1 and 2) respectively.

    Examples:
        >>> state = torch.rand(2, 2)
        >>> get_qubit_indices(0, state, num_qubits=2)
        0
        >>> state = torch.rand(500, 17, 2, 2, 2)
        >>> get_qubit_indices(0, state, num_qubits=3)
        2
        >>> state = torch.rand(500, 17, 2, 2, 2)
        >>> get_qubit_indices([1, 0], state, num_qubits=3)
        [3, 2]
        >>> # Negative indices are not converted to positive:
        >>> state = torch.rand(500, 17, 2, 2, 2)
        >>> get_qubit_indices([-1, 0], state, num_qubits=3)
        [-1, 2]
    """
    batch_dims = state_tensor.dim() - num_qubits
    range_message = 'Expected index in {-num_qubits, ..., num_qubits - 1}.\n'
    range_message += f'Num_qubits: {num_qubits}, index: {index}.'

    convert = False
    if not isinstance(index, torch.Tensor):
        convert = True
        index = torch.tensor(index)

    if (index >= num_qubits).any() or (index < -num_qubits).any():
        raise ValueError(range_message)
    index = index + (index >= 0) * batch_dims
    if convert:
        return index.tolist()
    else:
        return index