import torch
from typing import Tuple


def colorize_noise(img: torch.Tensor, color_min: Tuple[int, int, int] = (-255, -255, -255),
                   color_max: Tuple[int, int, int] = (255, 255, 255), p: float = 1) -> torch.Tensor:
    """
    Colorizes given noise images by asserting random color values to pixels that are not black (zero).
    :param img: torch tensor (n x c x h x w)
    :param color_min: limit the random color to be greater than this rgb value
    :param color_max: limit the random color to be less than this rgb value
    :param p: the chance to change the color of a pixel, on average changes p * h * w many pixels.
    :return: colorized images
    """
    assert 0 <= p <= 1
    orig_img = img.clone()
    if len(set(color_min)) == 1 and len(set(color_max)) == 1:
        cmin, cmax = color_min[0], color_max[0]
        img[img != 0] = torch.randint(cmin, cmax+1, img[img != 0].shape).type(img.dtype)
    else:
        img = img.transpose(0, 1)
        for ch, (cmin, cmax) in enumerate(zip(color_min, color_max)):
            img[ch][img[ch] != 0] = torch.randint(cmin, cmax+1, img[ch][img[ch] != 0].shape).type(img.dtype)
    if p < 1:
        pmask = torch.rand(img[img != 0].shape) >= p
        tar = img[img != 0]
        tar[pmask] = orig_img[img != 0][pmask]
        img[img != 0] = tar
    return img