import torch


def get_conditional_pixels(batch, mask, num_colors):
    """Returns conditional pixels obtained from masking the data in batch with
    mask and appending the mask. E.g. if the input has size (N, C, H, W)
    then the output will have size (N, C + 1, H, W) i.e. the mask is appended
    as an extra color channel.

    Parameters
    ----------
    batch : torch.Tensor
        Batch of data as returned by a DataLoader, i.e. unnormalized.
        Shape (num_examples, num_channels, width, height)

    mask : torch.Tensor
        Mask as returned by MaskGenerator.get_masks.
        Shape (num_examples, 1, width, height)

    num_colors : int
        Number of colors image is quantized to.
    """
    batch_size, channels, width, height = batch.size()
    # Add extra channel to keep mask
    cond_pixels = torch.zeros((batch_size, channels + 1, height, width))
    # Mask batch to only show visible pixels
    cond_pixels[:, :channels, :, :] = mask * batch.float()
    # Add mask scaled by number of colors in last channel dimension
    cond_pixels[:, -1:, :, :] = mask * (num_colors - 1)
    # Normalize conditional pixels to be in 0 - 1 range
    return cond_pixels / (num_colors - 1)