import torch
from torch import nn


def get_residual(weights):
    """Get the order of the first significant digit of the tensors"""
    signs = torch.sign(weights)
    exps = torch.round(torch.log2(torch.abs(weights)))
    pow_weights = signs * torch.pow(2, exps)
    return pow_weights, exps


def rf8(model, n=4):
    """Residual Float-Point 8-bit Model Quantization"""
    with torch.no_grad():
        for param in model.parameters():
            data1, exps1 = get_residual(param.data)
            data2, exps2 = get_residual(param.data - data1)
            flags = (exps1-exps2 <= n)
            param.data = data1 + flags * data2
            
            
def rf8_new(model):
    """8-bit Residual Float-pointing Format"""
    with torch.no_grad():
        for param in model.parameters():
            param_ = param.cpu()

            signs, exps = torch.sign(param_), torch.frexp(param_)[1] - 1

            bias = torch.tensor([-4, -3, -2, 1, 0], dtype=int)
            exps_ = exps.unsqueeze(-1).expand(*exps.shape, 5)
            Exponents = torch.exp2(exps)

            res_list = torch.exp2(bias + exps_)
            res_true = torch.abs(param_) - Exponents
            res_true = res_true.unsqueeze(-1).expand(*res_true.shape, 5)

            indices = (res_true - res_list).abs().argmin(-1).unsqueeze(-1)
            Residuals = torch.gather(res_list, -1, indices).squeeze()

            values = signs * (Exponents + Residuals)
            values[values.abs() < 2**-12] = 0
            values[values.abs() > 2**5] = 0
            param.data = values.to(torch.bfloat16).to(param.device)