# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai
# SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai
# SPDX-License-Identifier: Apache-2.0
# Contact: qubitium@modelcloud.ai, x.com/qubitium

import math
import time
from typing import Optional

import torch
import transformers
from torch import nn

from .. import QuantizeConfig
from ..looper.named_module import NamedModule
from ..quantization.config import FailSafeStrategy, SmoothMSE
from ..quantization.quantizer import HF_OPTIMUM
from ..utils import setup_logger
from .failsafe_smooth import mse_optimal_quant, smooth_block
from .gptq import get_number_of_rows_and_cols


DEBUG = False

log = setup_logger()

def quantize(x, scale, zero, maxq, sym, groupsize):
    if maxq < 0:
        return (x > scale / 2).float() * scale + (x < zero / 2).float() * zero
    if groupsize != -1 or not sym:
        q = torch.clamp(torch.round(x / scale) + zero, 0, maxq)
        return scale * (q - zero)
    else:
        q = torch.clamp(torch.round(x / scale), -maxq, maxq)
        return scale * q


class Quantizer(nn.Module):
    def __init__(self, shape=1):
        super(Quantizer, self).__init__()
        self.register_buffer("maxq", torch.tensor(0))
        self.register_buffer("scale", torch.zeros(shape))
        self.register_buffer("zero", torch.zeros(shape))

    def configure(
            self,
            bits,
            perchannel=False,
            sym=True,
            groupsize=-1,
            mse=False,
            norm=2.4,
            grid=100,
            maxshrink=0.8,
            trits=False,
    ):
        if groupsize != -1 or not sym:
            self.maxq = torch.tensor(2 ** bits - 1)
        else:
            self.maxq = torch.tensor(2 ** (bits - 1) - 1)
        self.perchannel = perchannel
        self.groupsize = groupsize
        self.sym = sym
        self.mse = mse
        self.norm = norm
        self.grid = grid
        self.maxshrink = maxshrink
        if trits:
            self.maxq = torch.tensor(-1)

    def find_params(self, x, weight=False):
        dev = x.device
        self.maxq = self.maxq.to(dev)

        shape = x.shape
        if self.perchannel:
            if weight:
                x = x.flatten(1)
            else:
                if len(shape) == 4:
                    x = x.permute([1, 0, 2, 3])
                    x = x.flatten(1)
                if len(shape) == 3:
                    x = x.reshape((-1, shape[-1])).t()
                if len(shape) == 2:
                    x = x.t()
        else:
            x = x.flatten().unsqueeze(0)

        tmp = torch.zeros(x.shape[0], device=dev)
        xmin = torch.minimum(x.min(1)[0], tmp)
        xmax = torch.maximum(x.max(1)[0], tmp)

        if self.sym:
            xmax = torch.maximum(torch.abs(xmin), xmax)
            tmp = xmin < 0
            if torch.any(tmp):
                xmin[tmp] = -xmax[tmp]
        tmp = (xmin == 0) & (xmax == 0)
        xmin[tmp] = -1
        xmax[tmp] = +1

        if self.maxq < 0:
            self.scale = xmax
            self.zero = xmin
        else:
            if self.groupsize != -1 or not self.sym:
                self.scale = (xmax - xmin) / self.maxq
                if self.sym:
                    self.zero = torch.full_like(self.scale, (self.maxq + 1) / 2)
                else:
                    self.zero = torch.round(-xmin / self.scale)
            else:
                self.scale = xmax / self.maxq
                self.zero = torch.zeros_like(self.scale)

        if self.mse:
            best = torch.full([x.shape[0]], float("inf"), device=dev)
            for i in range(int(self.maxshrink * self.grid)):
                p = 1 - i / self.grid
                xmin1 = p * xmin
                xmax1 = p * xmax
                scale1 = (
                    (xmax1 - xmin1) / self.maxq
                    if (self.groupsize != -1 or not self.sym)
                    else xmax1 / self.maxq
                )
                zero1 = torch.round(-xmin1 / scale1) if not self.sym else self.zero
                q = quantize(
                    x,
                    scale1.unsqueeze(1),
                    zero1.unsqueeze(1),
                    self.maxq,
                    self.sym,
                    self.groupsize,
                )
                q -= x
                q.abs_()
                q.pow_(self.norm)
                err = torch.sum(q, 1)
                tmp = err < best
                if torch.any(tmp):
                    best[tmp] = err[tmp]
                    self.scale[tmp] = scale1[tmp]
                    self.zero[tmp] = zero1[tmp]
        if not self.perchannel:
            if weight:
                tmp = shape[0]
            else:
                tmp = shape[1] if len(shape) != 3 else shape[2]
            self.scale = self.scale.repeat(tmp)
            self.zero = self.zero.repeat(tmp)

        if weight:
            shape = [-1] + [1] * (len(shape) - 1)
            self.scale = self.scale.reshape(shape)
            self.zero = self.zero.reshape(shape)
            return
        if len(shape) == 4:
            self.scale = self.scale.reshape((1, -1, 1, 1))
            self.zero = self.zero.reshape((1, -1, 1, 1))
        if len(shape) == 3:
            self.scale = self.scale.reshape((1, 1, -1))
            self.zero = self.zero.reshape((1, 1, -1))
        if len(shape) == 2:
            self.scale = self.scale.unsqueeze(0)
            self.zero = self.zero.unsqueeze(0)

    def quantize(self, x):
        if self.ready():
            return quantize(
                x, self.scale, self.zero, self.maxq, self.sym, self.groupsize
            )
        return x

    def enabled(self):
        return self.maxq > 0

    def ready(self):
        return torch.all(self.scale != 0)


class QQQ:
    def __init__(self, module: nn.Module, qcfg: Optional[QuantizeConfig] = None):
        # self.lock = threading.Lock()

        # self.num_tied_handles = 0
        # if qcfg.tied_gptq_handle is not None:
        #     qcfg.tied_gptq_handle.num_tied_handles += 1

        # Flags indicating issues
        # self.issue_zero_samples = False
        # self.issue_nan_hessian = False
        # self.issue_non_invertible = False

        # self.W = module.weight
        self.rows, self.columns = get_number_of_rows_and_cols(module)
        if isinstance(module, NamedModule):
            self.layer = module.module
            self.name = module.name
            self._named_module = module
        else:
            self.name = HF_OPTIMUM
            self.layer = module
            self._named_module = None

        layer_device = self.layer.weight.device
        setattr(self.layer, "target_device", layer_device)

        self.dev = self.layer.weight.device

        self._validate_module(self.layer)

        self.qcfg = qcfg if qcfg else QuantizeConfig()  # HF compat will not pass qcfg

        self._original_rows = self.rows
        self._original_columns = self.columns
        if self._named_module is not None:
            pad_info = self._named_module.state.get("tp_pad_info")
        else:
            pad_info = getattr(self.layer, "_tp_pad_info", None)

        if isinstance(pad_info, dict):
            pad_cols = int(pad_info.get("pad_cols", 0) or 0)
            pad_cols = max(pad_cols, 0)
        else:
            pad_cols = 0
        self._tp_pad_cols = pad_cols
        if self._tp_pad_cols:
            self.columns += self._tp_pad_cols

        self.module_copy = None

        self.nsamples = 0

        self.quantizer = Quantizer()

        # fwd counter
        self.fwd_counter = 0

        self.failsafe = self.qcfg.failsafe
        self.expected_nsamples: Optional[float] = None

        self.H = torch.zeros((self.columns, self.columns),
                             dtype=torch.float32,
                             device = self.dev)

    @staticmethod
    def _validate_module(module):
        assert isinstance(module, (nn.Linear, nn.Conv1d, nn.Conv2d,
                                   transformers.Conv1D)), f"We supports only linear and convolutional layers. actual = `{module}`"

    @staticmethod
    def _truncate_last_dim(tensor: torch.Tensor, length: int) -> torch.Tensor:
        if tensor.dim() == 0:
            return tensor

        trim = min(length, tensor.shape[-1])
        if trim == tensor.shape[-1]:
            return tensor

        return tensor.narrow(tensor.dim() - 1, 0, trim).contiguous()

    def _failsafe_quantize(self, strategy: FailSafeStrategy):
        maxq = 2 ** self.qcfg.bits - 1
        sigma = 3.0
        group_size = self.qcfg.group_size if self.qcfg.group_size != -1 else self.columns
        start_time = time.time()
        smooth_method = getattr(self.failsafe, "smooth", None)
        mse_steps = 32
        mse_maxshrink = 0.8
        if isinstance(smooth_method, SmoothMSE):
            mse_steps = smooth_method.steps
            mse_maxshrink = smooth_method.maxshrink

        W = self.layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        W = W.float()

        if self._tp_pad_cols:
            pad = torch.zeros(
                (W.shape[0], self._tp_pad_cols),
                dtype=W.dtype,
                device=W.device,
            )
            W = torch.cat((W, pad), dim=1)
        Q = torch.empty_like(W)
        scale_chunks = []
        zero_chunks = []

        for start in range(0, self.columns, group_size):
            end = min(start + group_size, self.columns)
            block = W[:, start:end]

            if isinstance(smooth_method, SmoothMSE):
                dequant, scale, zero = mse_optimal_quant(
                    block,
                    self.qcfg,
                    maxq,
                    steps=mse_steps,
                    maxshrink=mse_maxshrink,
                )
            else:
                block_mod, scale_factor = smooth_block(
                    block,
                    self.failsafe,
                    group_size=self.qcfg.group_size if self.qcfg.group_size != -1 else self.columns,
                )
                if strategy == FailSafeStrategy.MIDPOINT:
                    w_min = block_mod.min(dim=1, keepdim=True).values
                    w_max = block_mod.max(dim=1, keepdim=True).values
                    mid = (w_max + w_min) / 2.0
                    scale = torch.clamp((w_max - w_min) / maxq, min=1e-8)
                    zero_mid = torch.full_like(scale, maxq / 2.0)
                    q = torch.round((block_mod - mid) / scale + zero_mid)
                    q = torch.clamp(q, 0, maxq)
                    zero = torch.round(zero_mid - (mid / scale))
                    zero = torch.clamp(zero, 0, maxq)
                    dequant = (q - zero) * scale
                elif strategy == FailSafeStrategy.MEAN:
                    mean = block_mod.mean(dim=1, keepdim=True)
                    max_dev = torch.max((block_mod - mean).abs(), dim=1, keepdim=True).values
                    max_dev = torch.clamp(max_dev, min=1e-8)
                    scale = (2 * max_dev) / maxq
                    zero_mid = torch.full_like(scale, maxq / 2.0)
                    q = torch.round((block_mod - mean) / scale + zero_mid)
                    q = torch.clamp(q, 0, maxq)
                    zero = torch.round(zero_mid - (mean / scale))
                    zero = torch.clamp(zero, 0, maxq)
                    dequant = (q - zero) * scale
                elif strategy == FailSafeStrategy.MEDIAN:
                    median = block_mod.median(dim=1, keepdim=True).values
                    max_dev = torch.max((block_mod - median).abs(), dim=1, keepdim=True).values
                    max_dev = torch.clamp(max_dev, min=1e-8)
                    scale = (2 * max_dev) / maxq
                    zero_mid = torch.full_like(scale, maxq / 2.0)
                    q = torch.round((block_mod - median) / scale + zero_mid)
                    q = torch.clamp(q, 0, maxq)
                    zero = torch.round(zero_mid - (median / scale))
                    zero = torch.clamp(zero, 0, maxq)
                    dequant = (q - zero) * scale
                elif strategy == FailSafeStrategy.STDCLIP:
                    mean = block_mod.mean(dim=1, keepdim=True)
                    std = block_mod.std(dim=1, keepdim=True, unbiased=False)
                    std = torch.clamp(std, min=1e-8)
                    lo = mean - sigma * std
                    hi = mean + sigma * std
                    scale = torch.clamp((hi - lo) / maxq, min=1e-8)
                    zero = torch.round(-lo / scale)
                    zero = torch.clamp(zero, 0, maxq)
                    q = torch.round(block_mod / scale + zero)
                    q = torch.clamp(q, 0, maxq)
                    dequant = (q - zero) * scale
                elif strategy == FailSafeStrategy.RTN:
                    self.quantizer.find_params(block_mod, weight=True)
                    dequant = self.quantizer.quantize(block_mod)
                    scale = self.quantizer.scale
                    zero = self.quantizer.zero
                else:
                    raise ValueError(f"Unsupported failsafe strategy: {strategy}")

                if scale_factor is not None:
                    scale = scale * scale_factor
                    dequant = dequant * scale_factor

            Q[:, start:end] = dequant

            scale_block = scale if scale.dim() > 1 else scale.unsqueeze(1)
            zero_block = zero if zero.dim() > 1 else zero.unsqueeze(1)
            if scale_block.shape[1] > 1:
                scale_block = scale_block.mean(dim=1, keepdim=True)
            if zero_block.shape[1] > 1:
                zero_block = zero_block.mean(dim=1, keepdim=True)
            scale_chunks.append(scale_block)
            zero_chunks.append(zero_block)

        scale = torch.cat(scale_chunks, dim=1)
        zero = torch.cat(zero_chunks, dim=1)

        if self._tp_pad_cols:
            valid_cols = self._original_columns
            Q = Q[:, :valid_cols]
            scale = self._truncate_last_dim(scale, valid_cols)
            zero = self._truncate_last_dim(zero, valid_cols)
        else:
            valid_cols = self.columns

        g_idx = torch.arange(valid_cols, device=Q.device, dtype=torch.int32) // (group_size if group_size != -1 else self.columns)

        if isinstance(self.layer, transformers.Conv1D):
            Q = Q.t()
        if Q.shape != self.layer.weight.shape:
            Q = Q.reshape(self.layer.weight.shape).to(self.layer.weight.dtype)
        else:
            Q = Q.to(self.layer.weight.dtype)

        scale_extra = None
        if group_size != self.columns:
            quantizer_extra = Quantizer()
            quantizer_extra.configure(
                bits=8,
                perchannel=True,
                groupsize=-1,
                sym=True,
                mse=False,
            )
            quantizer_extra.find_params(self.layer.weight.data.clone(), weight=True)
            scale_extra = quantizer_extra.scale

        duration = time.time() - start_time
        mean_abs_err = (Q - self.layer.weight.data).abs().mean().item()
        avg_loss = f"failsafe({strategy.value}): {mean_abs_err:.7f}"
        damp_percent = 0.0
        self.H = None
        return Q, scale, zero, g_idx, duration, avg_loss, damp_percent, scale_extra, self.nsamples

    def add_batch(self, inp, out):
        if DEBUG:
            self.inp1 = inp
            self.out1 = out
        if len(inp.shape) == 2:
            inp = inp.unsqueeze(0)
        tmp = inp.shape[0]
        if isinstance(self.layer, nn.Linear) or isinstance(
                self.layer, transformers.Conv1D
        ):
            if len(inp.shape) == 3:
                inp = inp.reshape((-1, inp.shape[-1]))
            inp = inp.t()
        if isinstance(self.layer, nn.Conv2d):
            unfold = nn.Unfold(
                self.layer.kernel_size,
                dilation=self.layer.dilation,
                padding=self.layer.padding,
                stride=self.layer.stride,
            )
            inp = unfold(inp)
            inp = inp.permute([1, 0, 2])
            inp = inp.flatten(1)
        self.H *= self.nsamples / (self.nsamples + tmp)
        self.nsamples += tmp
        inp = math.sqrt(2 / self.nsamples) * inp.float()
        if self._tp_pad_cols:
            pad = inp.new_zeros((self._tp_pad_cols, inp.shape[1]))
            inp = torch.cat((inp, pad), dim=0)
        self.H += inp.matmul(inp.t())

    @torch.inference_mode()
    def quantize(
            self,
            blocksize=128,
    ):
        start = time.time()
        from ..utils.failsafe import resolve_failsafe_strategy, resolve_threshold, should_use_failsafe

        resolved_strategy = resolve_failsafe_strategy(self.failsafe)

        percdamp = self.qcfg.damp_percent
        groupsize = self.qcfg.group_size
        actorder = self.qcfg.desc_act
        static_groups = self.qcfg.static_groups

        W = self.layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        W = W.float()

        if self._tp_pad_cols:
            pad = torch.zeros(
                (W.shape[0], self._tp_pad_cols),
                dtype=W.dtype,
                device=W.device,
            )
            W = torch.cat((W, pad), dim=1)

        tick = time.time()

        if not self.quantizer.ready():
            self.quantizer.find_params(W, weight=True)

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        g_idx = []
        scale = []
        zero = []
        now_idx = 1
        if static_groups:
            import copy

            groups = []
            for i in range(0, self.columns, groupsize):
                quantizer = copy.deepcopy(self.quantizer)
                quantizer.find_params(W[:, i: (i + groupsize)], weight=True)
                scale.append(quantizer.scale)
                zero.append(quantizer.zero)
                groups.append(quantizer)

        if actorder:
            perm = torch.argsort(torch.diag(H), descending=True)
            W = W[:, perm]
            H = H[perm][:, perm]
            invperm = torch.argsort(perm)

        Losses = torch.zeros_like(W)
        Q = torch.zeros_like(W)

        damp = percdamp * torch.mean(torch.diag(H))
        diag = torch.arange(self.columns, device=self.dev)
        H[diag, diag] += damp
        threshold_raw, is_percent = resolve_threshold(self.failsafe, self.expected_nsamples)
        failsafe_configured = threshold_raw is not None

        try:
            H = torch.linalg.cholesky(H)
            H = torch.cholesky_inverse(H)
            H = torch.linalg.cholesky(H, upper=True)
            Hinv = H
        except Exception:
            fallback_requested = should_use_failsafe(
                self.failsafe,
                float(self.nsamples),
                self.expected_nsamples,
            )
            if fallback_requested:
                extra = f", threshold_raw={threshold_raw}" if threshold_raw is not None and is_percent else ""
                log.warn(
                    "Quantization: Module `%s` -> Using `%s` failsafe quantization (observed %s samples, threshold=%s%s, max_total=%s).",
                    self.name,
                    resolved_strategy.value,
                    self.nsamples,
                    self.failsafe,
                    extra,
                    self.expected_nsamples,
                )
                if resolved_strategy != FailSafeStrategy.RTN:
                    return self._failsafe_quantize(resolved_strategy)
                Hinv = None
            else:
                raise

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            count = i2 - i1

            W1 = W[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            Losses1 = torch.zeros_like(W1)
            Hinv1 = Hinv[i1:i2, i1:i2]

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]

                if groupsize != -1:
                    if not static_groups:
                        if (i1 + i) % groupsize == 0:
                            self.quantizer.find_params(
                                W[:, (i1 + i): (i1 + i + groupsize)], weight=True
                            )

                        if ((i1 + i) // groupsize) - now_idx == -1:
                            scale.append(self.quantizer.scale)
                            zero.append(self.quantizer.zero)
                            now_idx += 1
                    else:
                        idx = i1 + i
                        if actorder:
                            idx = perm[idx]
                        self.quantizer = groups[idx // groupsize]

                q = quantize(
                    w.unsqueeze(1),
                    self.quantizer.scale,
                    self.quantizer.zero,
                    self.quantizer.maxq,
                    self.quantizer.sym,
                    self.quantizer.groupsize,
                ).flatten()
                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d ** 2

                err1 = (w - q) / d
                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
                Err1[:, i] = err1

            Q[:, i1:i2] = Q1
            Losses[:, i1:i2] = Losses1 / 2

            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

            if DEBUG:
                self.layer.weight.data[:, :i2] = Q[:, :i2]
                self.layer.weight.data[:, i2:] = W[:, i2:]
                print(torch.sum((self.layer(self.inp1) - self.out1) ** 2))
                print(torch.sum(Losses))

        torch.cuda.synchronize()
        print("time %.2f" % (time.time() - tick))
        print("error", torch.sum(Losses).item())

        if Hinv is not None:
            del Hinv
            if self.nsamples != 0:
                avg_loss = torch.sum(Losses).item() / self.nsamples

                if math.isnan(avg_loss):
                    print("Losses sum item:", torch.sum(Losses).item())
                    if failsafe_configured:
                        log.info(f"Quantization: Failed due to `NaN` loss for `{self.name}`, use mock quantization retry for `{self.name}`")
                        self.qcfg.mock_quantization = True
                        return self.quantize(blocksize=blocksize)
                    else:
                        raise ValueError(f"Quantization: Failed due to `NaN` loss for `{self.name}`, please try increasing calibration data samples or enable failsafe=True")
            else:
                if failsafe_configured:
                    log.warn(f"Quantization: Module `{self.name}` -> using fail safe mode. Please check if calibration data is sufficient.")
                else:
                    log.warn(f"Quantization: `{self.name}` is not activated due to model inference logic (MoE)")
                avg_loss = f"{resolved_strategy.value} failsafe" if failsafe_configured else 999999999
        else:
            avg_loss = f"{resolved_strategy.value} failsafe" if failsafe_configured else 999999999

        del Losses

        groupsize = groupsize if groupsize != -1 else self.columns
        if static_groups and actorder:
            g_idx = [perm[i] // groupsize for i in range(self.columns)]
        else:
            g_idx = [i // groupsize for i in range(self.columns)]
        g_idx = torch.tensor(g_idx, dtype=torch.int32, device=Q.device)
        if actorder:
            Q = Q[:, invperm]
            g_idx = g_idx[invperm]

        if self._tp_pad_cols:
            valid_cols = self._original_columns
            Q = Q[:, :valid_cols]
            g_idx = g_idx[:valid_cols]

        if isinstance(self.layer, transformers.Conv1D):
            Q = Q.t()
        if Q.shape != self.layer.weight.shape:
            Q = Q.reshape(self.layer.weight.shape).to(self.layer.weight.dtype)
        else:
            Q = Q.to(self.layer.weight.dtype)
        if DEBUG:
            print(torch.sum((self.layer(self.inp1) - self.out1) ** 2))

        if scale == []:
            scale.append(self.quantizer.scale)
            zero.append(self.quantizer.zero)
        scale = torch.cat(scale, dim=1)
        zero = torch.cat(zero, dim=1)

        if self._tp_pad_cols:
            valid_cols = self._original_columns
            scale = self._truncate_last_dim(scale, valid_cols)
            zero = self._truncate_last_dim(zero, valid_cols)

        Q = Q.to(device=self.layer.weight.data.device, non_blocking=False)

        # post int8 quant
        scale_extra = None
        if groupsize != self.columns:
            quantizer_extra = Quantizer()
            quantizer_extra.configure(
                bits=8,
                perchannel=True,
                groupsize=-1,
                sym=True,
                mse=False,
            )
            quantizer_extra.find_params(self.layer.weight.data.clone(), weight=True)
            scale_extra = quantizer_extra.scale

        duration = time.time() - start

        return Q, scale, zero, g_idx, duration, avg_loss, damp, scale_extra, self.nsamples

    def free(self):
        if DEBUG:
            self.inp1 = None
            self.out1 = None
        self.H = None
        self.Losses = None
        self.Trace = None
        if hasattr(self, "quantizer"):
            del self.quantizer
        if hasattr(self, "module_copy"):
            del self.module_copy
        if self._named_module is not None:
            self._named_module.state.pop("tp_pad_info", None)
        torch.cuda.empty_cache()
