import torch


@torch.jit.script
def cal_scale(amax, fp_max, scale):
    margin = 0
    exp = torch.floor(torch.log2(fp_max / amax)) - margin
    sf = torch.round(torch.pow(2, torch.abs(exp)))
    sf = torch.where(amax > 0.0, sf, scale)
    sf = torch.where(torch.isfinite(amax), sf, scale)
    scale = torch.where(exp < 0, 1 / sf, sf)
    scale_inv = torch.reciprocal(scale)
    return scale, scale_inv


instances = {}


def singleton(cls):
    global instances

    def get_instance(*args, **kwargs):
        if cls not in instances:
            instances[cls] = cls(*args, **kwargs)
        return instances[cls]

    return get_instance


def reset_singleton():
    global instances
    instances = {}


@singleton
class QuantFp8:

    def __init__(self, device):
        self.fp_max = torch.tensor([448.0], device=device)
        self.device = device
        self.scale = torch.tensor([1.0], device=self.device)
        pass

    @staticmethod
    def quantize_v1(weight, bits):
        if bits == 8:
            amax = weight.abs().max()
            fp_max = torch.tensor([448.0]).to(weight.device)
            margin = 0
            scale = torch.tensor([1.0]).to(weight.device)

            exp = torch.floor(torch.log2(fp_max / amax)) - margin
            sf = torch.round(torch.pow(2, torch.abs(exp)))
            sf = torch.where(amax > 0.0, sf, scale)
            sf = torch.where(torch.isfinite(amax), sf, scale)
            scale = torch.where(exp < 0, 1 / sf, sf)

            qweight = (weight.to(torch.float32) * scale).to(
                torch.float8_e4m3fn)
            scale = torch.reciprocal(scale)
            # print(f"amax={amax},scalse={scale}")
        else:
            raise ValueError(f"Unsupported bit width: {bits}")
        return qweight, scale

    def quantize(self, weight, bits, weight_scale, use_offline_input_scales):
        if bits == 8:
            amax = torch.empty(1, dtype=torch.float32, device=self.device)
            scale = torch.tensor([1.0], device=self.device)
            torch.ops.OptimusFp8.abs_max_nan_to_inf(weight, amax)
            if weight_scale is None or not use_offline_input_scales:
                scale, scale_inv = cal_scale(amax, self.fp_max, scale)
            else:
                scale, scale_inv = weight_scale, torch.reciprocal(weight_scale)

            qweight = torch.ops.OptimusFp8.quantize(weight, scale, None,
                                                    torch.float8_e4m3fn)
            # print(f"scale={scale},self.amax={self.amax}")
            return qweight, scale_inv
        else:
            raise ValueError(f"Unsupported bit width: {bits}")

    def get_quant_scale(self, tensor):
        amax = torch.empty(1, dtype=torch.float32, device=tensor.device)
        torch.ops.OptimusFp8.abs_max_nan_to_inf(tensor, amax)
        scale, _ = cal_scale(amax, self.fp_max, self.scale)
        return scale


def quantize(weight, bits, weight_scale=None, use_offline_input_scales=True):
    quant = QuantFp8(weight.device)
    return quant.quantize(weight, bits, weight_scale, use_offline_input_scales)


def dequant(weight, weight_scales):
    return torch.ops.OptimusFp8.dequantize(weight, weight_scales,
                                           torch.bfloat16)


def experts_dequant(weights, weight_scales):
    ret = torch.empty(*weights.shape,
                      device=weights.device,
                      dtype=torch.bfloat16)
    for i in range(weights.shape[0]):
        ret[i] = dequant(weights[i], weight_scales[i])
    return ret


def experts_quantize(weight, bits):
    if bits == 8:
        qweight_experts = torch.empty(*weight.shape,
                                      dtype=torch.float8_e4m3fn,
                                      device=weight.device)
        scales = torch.empty(weight.shape[0],
                             dtype=torch.float32,
                             device=weight.device)
        for idx in range(weight.shape[0]):
            expert_weight = weight[idx]
            qweight, scale = quantize(expert_weight, bits)
            qweight_experts[idx] = qweight
            scales[idx] = scale
        return qweight_experts, scales
    else:
        raise ValueError(f"Unsupported bit width: {bits}")


def dynamic_fp8_pertensor_quantize(tensor):
    # amax = torch.empty(1, dtype=torch.float32, device=tensor.device)
    # scale = torch.tensor([1.0], device=tensor.device)
    # fp_max = torch.tensor([448.0], device=tensor.device)
    # torch.ops.OptimusFp8.abs_max_nan_to_inf(tensor, amax)
    # scale, _ = cal_scale(amax, fp_max, scale)
    # return scale
    quant = QuantFp8(tensor.device)
    return quant.get_quant_scale(tensor)