# Ported from ostris/ai-toolkit (MIT License)
# https://github.com/ostris/ai-toolkit/blob/main/toolkit/optimizers/optimizer_utils.py
import torch
from torch import Tensor
from typing import Optional

try:
    from optimum.quanto import QBytesTensor
    _QUANTO_AVAILABLE = True
except ImportError:
    QBytesTensor = None
    _QUANTO_AVAILABLE = False


def compute_scale_for_dtype(tensor, dtype):
    if dtype == torch.int8:
        abs_max = torch.max(torch.abs(tensor))
        return abs_max / 127.0 if abs_max > 0 else 1.0
    elif dtype == torch.uint8:
        max_val = torch.max(tensor)
        min_val = torch.min(tensor)
        range_val = max_val - min_val
        return range_val / 255.0 if range_val > 0 else 1.0
    elif dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
        abs_max = torch.max(torch.abs(tensor))
        max_representable = 448.0 if dtype == torch.float8_e4m3fn else 57344.0
        return abs_max / max_representable if abs_max > 0 else 1.0
    else:
        raise ValueError(f"Unsupported dtype for quantization: {dtype}")


def quantize_tensor(tensor, dtype):
    scale = compute_scale_for_dtype(tensor, dtype)
    if dtype == torch.int8:
        quantized_data = torch.clamp(torch.round(tensor / scale), -128, 127).to(dtype)
    elif dtype == torch.uint8:
        quantized_data = torch.clamp(torch.round(tensor / scale), 0, 255).to(dtype)
    elif dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
        quantized_data = (tensor / scale).to(dtype)
    else:
        raise ValueError(f"Unsupported dtype for quantization: {dtype}")
    return quantized_data, scale


def update_parameter(target, result_float):
    if _QUANTO_AVAILABLE and isinstance(target, QBytesTensor):
        target_dtype = target._data.dtype
        device = target._data.device
        result_float = result_float.to(device)
        quantized_data, new_scale = quantize_tensor(result_float, target_dtype)
        target._data.copy_(quantized_data)
        target._scale.copy_(new_scale)
    else:
        target.copy_(result_float)


def get_format_params(dtype: torch.dtype):
    if dtype == torch.float32:
        return 23, 32
    elif dtype == torch.bfloat16:
        return 7, 16
    elif dtype == torch.float16:
        return 10, 16
    elif dtype == torch.float8_e4m3fn:
        return 3, 8
    elif dtype == torch.float8_e5m2:
        return 2, 8
    elif dtype == torch.int8:
        return 0, 8
    else:
        raise ValueError(f"Unsupported dtype: {dtype}")


def copy_stochastic_bf16(target: torch.Tensor, source: torch.Tensor):
    result = torch.randint_like(source, dtype=torch.int32, low=0, high=(1 << 16))
    result.add_(source.view(dtype=torch.int32))
    result.bitwise_and_(-65536)
    target.copy_(result.view(dtype=torch.float32))
    del result


def copy_stochastic(target: torch.Tensor, source: torch.Tensor, eps: Optional[float] = None) -> None:
    with torch.no_grad():
        assert target.device.type != 'cpu', "Target is on cpu!"
        assert source.device.type != 'cpu', "Source is on cpu!"
        if target.dtype == torch.float32:
            target.copy_(source)
            return
        if target.dtype == torch.bfloat16:
            copy_stochastic_bf16(target, source)
            return
        mantissa_bits, _ = get_format_params(target.dtype)
        round_factor = 2 ** (23 - mantissa_bits)
        noise = torch.rand_like(source, device=source.device) - 0.5
        rounded = torch.round(source * round_factor + noise)
        result_float = rounded / round_factor
        if target.dtype == torch.float8_e4m3fn:
            result_float.clamp_(-448.0, 448.0)
        elif target.dtype == torch.float8_e5m2:
            result_float.clamp_(-57344.0, 57344.0)
        update_parameter(target, result_float)


class Auto8bitTensor:
    def __init__(self, data: Tensor, *args, **kwargs):
        if isinstance(data, dict):
            self._load_from_state_dict(data)
        else:
            abs_max = data.abs().max().item()
            scale = abs_max / 127.0 if abs_max > 0 else 1.0
            self.quantized = (data / scale).round().clamp(-127, 127).to(torch.int8)
            self.scale = scale
            self.orig_dtype = data.dtype

    def dequantize(self) -> Tensor:
        return self.quantized.to(dtype=torch.float32) * self.scale

    def to(self, *args, **kwargs):
        dtype = None
        if args and isinstance(args[0], torch.dtype):
            dtype = args[0]
            args = args[1:]
        elif 'dtype' in kwargs:
            dtype = kwargs['dtype']
            del kwargs['dtype']
        if dtype is not None:
            return self.dequantize().to(dtype=dtype, *args, **kwargs)
        return self.dequantize().to(*args, **kwargs)

    def state_dict(self):
        return {'quantized': self.quantized, 'scale': self.scale, 'orig_dtype': self.orig_dtype}

    def _load_from_state_dict(self, state_dict):
        self.quantized = state_dict['quantized']
        self.scale = state_dict['scale']
        self.orig_dtype = state_dict['orig_dtype']

    def __str__(self):
        return f"Auto8bitTensor({self.dequantize()})"


def stochastic_grad_accummulation(param):
    if hasattr(param, "_accum_grad"):
        grad_fp32 = param._accum_grad.clone().to(torch.float32)
        grad_fp32.add_(param.grad.to(torch.float32))
        copy_stochastic(param._accum_grad, grad_fp32)
        del grad_fp32
        del param.grad
    else:
        param._accum_grad = param.grad.clone()
        del param.grad