import torch


def augment_light_colour(batch_size, device,
                         ambient_intensity_range=(0.2, 0.8),
                         diffuse_intensity_range=(0.2, 0.8),
                         specular_intensity_range=(0.2, 0.8)):
    """
    Samples batch of random light INTENSITIES (not colours because
    I am forcing RGB components to be equal i.e. white lights).
    :param ambient_intensity_range: ambient component intensity range
    :param diffuse_intensity_range: diffuse component intensity range
    :param specular_intensity_range: specular component intensity range
    :return:
    """

    l, h = ambient_intensity_range
    ambient = (h - l) * torch.rand(batch_size, device=device) + l
    ambient = ambient[:, None].expand(-1, 3)

    l, h = diffuse_intensity_range
    diffuse = (h - l) * torch.rand(batch_size, device=device) + l
    diffuse = diffuse[:, None].expand(-1, 3)

    l, h = specular_intensity_range
    specular = (h - l) * torch.rand(batch_size, device=device) + l
    specular = specular[:, None].expand(-1, 3)

    return ambient, diffuse, specular