import torch


def weighted_loss(inputs, outputs, alpha):
    # Calculate RGB Norm (Perceptual Difference)
    rgb_norm = torch.norm(inputs[:, None, :] - inputs[None, :, :], dim=-1)

    # Calculate 1D Space Norm
    transformed_norm = torch.norm(outputs[:, None] - outputs[None, :], dim=-1)

    # Weighted Loss
    loss = alpha * rgb_norm + (1 - alpha) * transformed_norm
    return torch.mean(loss)


def enhanced_loss(inputs, outputs, alpha, distinct_threshold):
    # Calculate RGB Norm
    rgb_norm = torch.norm(inputs[:, None, :] - inputs[None, :, :], dim=-1)

    # Calculate 1D Space Norm
    transformed_norm = torch.norm(outputs[:, None] - outputs[None, :], dim=-1)

    # Identify Distinct Colors (based on a threshold in RGB space)
    distinct_colors = rgb_norm > distinct_threshold

    # Penalty for Distinct Colors being too close in the transformed space
    # Here we do not take the mean yet, to avoid double averaging
    distinct_penalty = (1.0 / (transformed_norm + 1e-6)) * distinct_colors

    # Combined Loss
    # The mean is taken here, once, after all components are combined
    loss = alpha * rgb_norm + (1 - alpha) * transformed_norm + distinct_penalty
    return torch.mean(loss)