import torch


def _get_label_batch(label,
                     num_timesteps=0,
                     num_classes=0,
                     num_batches=0,
                     timesteps_noise=False):
    """Get label batch. Support get sequeue of label along timesteps.

    We support the following use cases ('bz' denotes ```num_batches`` and 'n'
    denotes ``num_timesteps``):

    If num_classes <= 0, return None.

    If timesteps_noise is True, we output label which dimension is 2.
    - Input is [bz, ]: Expand to [n, bz]
    - Input is [n, ]: Expand to [n, bz]
    - Input is [n*bz, ]: View to [n, bz]
    - Dim of the input is 2: Return the input, ignore ``num_batches`` and
      ``num_timesteps``
    - Callable or None: Generate label shape as [n, bz]
    - Otherwise: Raise error

    If timesteps_noise is False, we output label which dimension is 1 and
    ignore ``num_timesteps``.
    - Dim of the input is 1: Unsqueeze to [1, ], ignore ``num_batches``
    - Dim of the input is 2: Return the input. ignore ``num_batches``
    - Callable or None: Generate label shape as [bz, ]
    - Otherwise: Raise error

    It's to be noted that, we do not move the generated label to target device
    in this function because we can not get which device the noise should move
    to.

    Args:
        label (torch.Tensor | callable | None): You can directly give a
            batch of noise through a ``torch.Tensor`` or offer a callable
            function to sample a batch of noise data. Otherwise, the
            ``None`` indicates to use the default noise sampler.
        num_timesteps (int, optional): Total timestpes of the diffusion and
            denoising process. Defaults to 0.
        num_batches (int, optional): The number of batch size. To be noted that
            this argument only work when the input ``noise`` is callable or
            ``None``. Defaults to 0.
        timesteps_noise (bool, optional): If True, returned noise will shape
            as [n, bz, c, h, w], otherwise shape as [bz, c, h, w].
            Defaults to False.
    Returns:
        torch.Tensor: Generated label with desired shape.
    """
    # no labels output if num_classes is 0
    if num_classes == 0:
        assert label is None, ('\'label\' should be None '
                               'if \'num_classes == 0\'.')
        return None

    # receive label and conduct sanity check.
    if isinstance(label, torch.Tensor):
        if timesteps_noise:
            if label.ndim == 1:
                assert num_batches > 0 and num_timesteps > 0
                # [n, ] to [n, bz]
                if label.shape[0] == num_timesteps:
                    label_batch = label.view(num_timesteps, 1)
                    label_batch = label_batch.expand(-1, num_batches)
                # [bz, ] to [n, bz]
                elif label.shape[0] == num_batches:
                    label_batch = label.view(1, num_batches)
                    label_batch = label_batch.expand(num_timesteps, -1)
                # [n*bz, ] to [n, bz]
                elif label.shape[0] == num_timesteps * num_batches:
                    label_batch = label.view(num_timesteps, -1)
                else:
                    raise ValueError(
                        'The timesteps label should be in shape of '
                        '(n, ), (bz,), (n*bz, ) or (n, bz, ). But receive '
                        f'{label.shape}.')

            elif label.ndim == 2:
                # dimension is 2, direct return
                label_batch = label
            else:
                raise ValueError(
                    'The timesteps label should be in shape of '
                    '(n, ), (bz,), (n*bz, ) or (n, bz, ). But receive '
                    f'{label.shape}.')
        else:
            # dimension is 0, expand to [1, ]
            if label.ndim == 0:
                label_batch = label[None, ...]
            # dimension is 1, do nothing
            elif label.ndim == 1:
                label_batch = label
            else:
                raise ValueError(
                    'The label should be in shape of (bz, ) or'
                    f'zero-dimension tensor, but got {label.shape}')
    # receive a noise generator and sample noise.
    elif callable(label):
        assert num_batches > 0
        label_generator = label
        if timesteps_noise:
            assert num_timesteps > 0
            # generate label shape as [n, bz]
            label_batch = label_generator((num_timesteps, num_batches))
        else:
            # generate label shape as [bz, ]
            label_batch = label_generator((num_batches, ))
    # otherwise, we will adopt default label sampler.
    else:
        assert num_batches > 0
        if timesteps_noise:
            assert num_timesteps > 0
            # generate label shape as [n, bz]
            label_batch = torch.randint(0, num_classes,
                                        (num_timesteps, num_batches))
        else:
            # generate label shape as [bz, ]
            label_batch = torch.randint(0, num_classes, (num_batches, ))

    return label_batch