import torch


def dequantize_data(batch_data, batch_max=None, batch_min=None):
    """
    Dequantize data using random noise and logit function.
    It uses Max min to first change the range of the data to [0, 1].

    Usage:
    dequantized_data, batch_max, batch_min, rand_noise = dequantize_data(batch_data)

    Args:
        batch_data -> Data to dequantize
        batch_max -> Maximum value in the batch (one for each feature) (Calculated automatically if not provided).
        batch_min -> Minimum Value (feature wise) (Calculated automatically if not provided).
    """
    if batch_max is None:
        batch_max = batch_data.max(0)[0]
    if batch_min is None:
        batch_min = batch_data.min(0)[0]
    ret_data = (torch.nan_to_num((batch_data - batch_min) / (batch_max - batch_min)) * 255.0)
    rand_noise = torch.rand_like(ret_data)
    ret_data = ret_data + rand_noise
    ret_data = ret_data / 255.0
    ret_data = torch.logit(ret_data, eps=1e-6)

    return ret_data, batch_max, batch_min, rand_noise