import torch


def hsv_to_rgb(input_hsv_tensor):
    """
    Differentiable HSV to RGB conversion function.
    :param input_hsv_tensor: Batch of HSV images [batch_size, 3, height, width]
    :return: Batch of RGB images [batch_size, 3, height, width]
    """
    assert len(input_hsv_tensor.shape) == 4 and input_hsv_tensor.shape[1] == 3
    hues = input_hsv_tensor[:, 0, :, :]
    sats = input_hsv_tensor[:, 1, :, :]
    vals = input_hsv_tensor[:, 2, :, :]
    c = sats * vals

    x = c * (1 - torch.abs((hues * 6.0) % 2.0 - 1.0))
    m = vals - c

    # Compute R
    r_hat = torch.zeros_like(hues)
    filter_hues = hues.clone()
    r_hat[filter_hues < 1.0 / 6.0] = c[filter_hues < 1.0 / 6.0]
    filter_hues[filter_hues < 1.0 / 6.0] += 10.0
    r_hat[filter_hues < 2.0 / 6.0] = x[filter_hues < 2.0 / 6.0]
    filter_hues[filter_hues < 2.0 / 6.0] += 10.0
    r_hat[filter_hues < 3.0 / 6.0] = 0
    filter_hues[filter_hues < 3.0 / 6.0] += 10.0
    r_hat[filter_hues < 4.0 / 6.0] = 0
    filter_hues[filter_hues < 4.0 / 6.0] += 10.0
    r_hat[filter_hues < 5.0 / 6.0] = x[filter_hues < 5.0 / 6.0]
    filter_hues[filter_hues < 5.0 / 6.0] += 10.0
    r_hat[filter_hues <= 6.0 / 6.0] = c[filter_hues <= 6.0 / 6.0]
    filter_hues[filter_hues <= 6.0 / 6.0] += 10.0

    # Compute G
    g_hat = torch.zeros_like(hues)
    filter_hues = hues.clone()
    g_hat[filter_hues < 1.0 / 6.0] = x[filter_hues < 1.0 / 6.0]
    filter_hues[filter_hues < 1.0 / 6.0] += 10.0
    g_hat[filter_hues < 2.0 / 6.0] = c[filter_hues < 2.0 / 6.0]
    filter_hues[filter_hues < 2.0 / 6.0] += 10.0
    g_hat[filter_hues < 3.0 / 6.0] = c[filter_hues < 3.0 / 6.0]
    filter_hues[filter_hues < 3.0 / 6.0] += 10.0
    g_hat[filter_hues < 4.0 / 6.0] = x[filter_hues < 4.0 / 6.0]
    filter_hues[filter_hues < 4.0 / 6.0] += 10.0
    g_hat[filter_hues < 5.0 / 6.0] = 0
    filter_hues[filter_hues < 5.0 / 6.0] += 10.0
    g_hat[filter_hues <= 6.0 / 6.0] = 0
    filter_hues[filter_hues <= 6.0 / 6.0] += 10.0

    # Compute B
    b_hat = torch.zeros_like(hues)
    filter_hues = hues.clone()
    b_hat[filter_hues < 1.0 / 6.0] = 0
    filter_hues[filter_hues < 1.0 / 6.0] += 10.0
    b_hat[filter_hues < 2.0 / 6.0] = 0
    filter_hues[filter_hues < 2.0 / 6.0] += 10.0
    b_hat[filter_hues < 3.0 / 6.0] = x[filter_hues < 3.0 / 6.0]
    filter_hues[filter_hues < 3.0 / 6.0] += 10.0
    b_hat[filter_hues < 4.0 / 6.0] = c[filter_hues < 4.0 / 6.0]
    filter_hues[filter_hues < 4.0 / 6.0] += 10.0
    b_hat[filter_hues < 5.0 / 6.0] = c[filter_hues < 5.0 / 6.0]
    filter_hues[filter_hues < 5.0 / 6.0] += 10.0
    b_hat[filter_hues <= 6.0 / 6.0] = x[filter_hues <= 6.0 / 6.0]
    filter_hues[filter_hues <= 6.0 / 6.0] += 10.0

    r = (r_hat + m).view(input_hsv_tensor.shape[0], 1, input_hsv_tensor.shape[2],
                         input_hsv_tensor.shape[3])
    g = (g_hat + m).view(input_hsv_tensor.shape[0], 1, input_hsv_tensor.shape[2],
                         input_hsv_tensor.shape[3])
    b = (b_hat + m).view(input_hsv_tensor.shape[0], 1, input_hsv_tensor.shape[2],
                         input_hsv_tensor.shape[3])

    rgb = torch.cat([r, g, b], dim=1)
    return rgb