from typing import Tuple, Union

from torch import Tensor, Size


def fold_batch(x: Tensor, nonbatch_ndims: int = 1) -> Tuple[Tensor, Size]:
    """
    Overview:
        :math:`(T, B, X) \leftarrow (T*B, X)`\
        Fold the first (ndim - nonbatch_ndims) dimensions of a tensor as batch dimension.\
        This operation is similar to `torch.flatten` but provides an inverse function
        `unfold_batch` to restore the folded dimensions.

    Arguments:
        - x (:obj:`torch.Tensor`): the tensor to fold
        - nonbatch_ndims (:obj:`int`): the number of dimensions that is not folded as
            batch dimension.

    Returns:
        - x (:obj:`torch.Tensor`): the folded tensor
        - batch_dims: the folded dimensions of the original tensor, which can be used to
             reverse the operation

    Examples:
        >>> x = torch.ones(10, 20, 5, 4, 8)
        >>> x, batch_dim = fold_batch(x, 2)
        >>> x.shape == (1000, 4, 8)
        >>> batch_dim == (10, 20, 5)

    """
    if nonbatch_ndims > 0:
        batch_dims = x.shape[:-nonbatch_ndims]
        x = x.view(-1, *(x.shape[-nonbatch_ndims:]))
        return x, batch_dims
    else:
        batch_dims = x.shape
        x = x.view(-1)
        return x, batch_dims


def unfold_batch(x: Tensor, batch_dims: Union[Size, Tuple]) -> Tensor:
    """
    Overview:
        Unfold the batch dimension of a tensor.

    Arguments:
        - x (:obj:`torch.Tensor`): the tensor to unfold
        - batch_dims (:obj:`torch.Size`): the dimensions that are folded

    Returns:
        - x (:obj:`torch.Tensor`): the original unfolded tensor

    Examples:
        >>> x = torch.ones(10, 20, 5, 4, 8)
        >>> x, batch_dim = fold_batch(x, 2)
        >>> x.shape == (1000, 4, 8)
        >>> batch_dim == (10, 20, 5)
        >>> x = unfold_batch(x, batch_dim)
        >>> x.shape == (10, 20, 5, 4, 8)
    """
    return x.view(*batch_dims, *x.shape[1:])


def unsqueeze_repeat(x: Tensor, repeat_times: int, unsqueeze_dim: int = 0) -> Tensor:
    """
    Overview:
        Squeeze the tensor on `unsqueeze_dim` and then repeat in this dimension for `repeat_times` times.\
        This is useful for preproprocessing the input to an model ensemble.

    Arguments:
        - x (:obj:`torch.Tensor`): the tensor to squeeze and repeat
        - repeat_times (:obj:`int`): the times that the tensor is repeatd
        - unsqueeze_dim (:obj:`int`): the unsqueezed dimension

    Returns:
        - x (:obj:`torch.Tensor`): the unsqueezed and repeated tensor

    Examples:
        >>> x = torch.ones(64, 6)
        >>> x = unsqueeze_repeat(x, 4)
        >>> x.shape == (4, 64, 6)

        >>> x = torch.ones(64, 6)
        >>> x = unsqueeze_repeat(x, 4, -1)
        >>> x.shape == (64, 6, 4)
    """
    assert -1 <= unsqueeze_dim <= len(x.shape), f'unsqueeze_dim should be from {-1} to {len(x.shape)}'
    x = x.unsqueeze(unsqueeze_dim)
    repeats = [1] * len(x.shape)
    repeats[unsqueeze_dim] *= repeat_times
    return x.repeat(*repeats)