import torch


def float_2_label(x, bits):
    """output = float_2_label(x, bits)
    
    Assume x is a float value, do N-bits quantization and 
    return the code index.

    input
    -----
       x: data to be converted, any shape
          x value should be [-1, 1]
       bits: number of bits, int
    
    output
    ------
       output: tensor.float, [0, 2**bits-1]
    
    Although output is quantized, we use torch.float to save
    the quantized values
    """
    #assert abs(x).max() <= 1.0
    # scale the peaks
    peak = torch.abs(x).max()
    if peak > 1.0:
        x /= peak
    # quantize
    x = (x + 1.) * (2**bits - 1) / 2
    return torch.clamp(x, 0, 2**bits - 1)