# quantization/cfrn/quantlib/qmodules.py
from __future__ import annotations

import math
from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

from .quant_utils import (
    QParams,
    GroupQParams,
    calc_qparams_symmetric,
    calc_group_qparams_symmetric,
    quantize_symmetric,
    dequantize_symmetric,
    quantize_group_symmetric,
    dequantize_group_symmetric,
    pack_int4,
    unpack_int4,
)


def _get_compute_dtype(name: str) -> torch.dtype:
    name = str(name).lower()
    if name in ("fp16", "float16", "half"):
        return torch.float16
    return torch.float32


class FP32LayerNorm(nn.LayerNorm):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = F.layer_norm(
            x.float(),
            self.normalized_shape,
            self.weight.float() if self.weight is not None else None,
            self.bias.float() if self.bias is not None else None,
            self.eps,
        )
        return y.to(x.dtype)


class FP32GroupNorm(nn.GroupNorm):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = F.group_norm(
            x.float(),
            self.num_groups,
            self.weight.float() if self.weight is not None else None,
            self.bias.float() if self.bias is not None else None,
            self.eps,
        )
        return y.to(x.dtype)


# -------------------------
# Helpers
# -------------------------
def _pad_last_dim(x: torch.Tensor, multiple: int) -> Tuple[torch.Tensor, int]:
    """Pad last dim to be divisible by `multiple`, return (padded, original_C)."""
    C = x.shape[-1]
    if multiple <= 0 or C % multiple == 0:
        return x, C
    pad = multiple - (C % multiple)
    return F.pad(x, (0, pad)), C


def _kthvalue_quantile(x: torch.Tensor, percentile: float, dim: int, keepdim: bool = True) -> torch.Tensor:
    """
    IMPORTANT: keep original semantics: k = ceil(p*n), kthvalue uses 1-indexed k.
    For group_size=32, p=0.99 -> k=32 -> max (effectively no clip).
    """
    p = float(percentile)
    if p >= 1.0:
        return x.amax(dim=dim, keepdim=keepdim)
    if p <= 0.0:
        return x.amin(dim=dim, keepdim=keepdim)
    n = x.shape[dim]
    k = int(math.ceil(p * n))
    k = max(1, min(n, k))
    return torch.kthvalue(x, k=k, dim=dim, keepdim=keepdim).values


def _dstg_clip_value(x: torch.Tensor, dstg_clip_percentile: Optional[float], dim: int = -1) -> torch.Tensor:
    """
    DSTG clip:
      - default: max(|x|)
      - optional: kthvalue percentile on |x|
    (RMS cap 已删除)
    """
    abs_x = x.abs()
    if dstg_clip_percentile is None:
        clip = abs_x.amax(dim=dim, keepdim=True)
    else:
        clip = _kthvalue_quantile(abs_x, float(dstg_clip_percentile), dim=dim, keepdim=True)
    return torch.clamp(clip, min=1e-8)


def _token_clip_value(x: torch.Tensor, token_clip_percentile: Optional[float], dim: int = -1) -> torch.Tensor:
    """Deprecated alias of `_dstg_clip_value` for compatibility."""
    return _dstg_clip_value(x, token_clip_percentile, dim=dim)


def _apply_dcrp_constraints(
    clip: torch.Tensor,
    x: torch.Tensor,
    qmax: int,
    tau: Optional[float],
    zero_ratio: Optional[float],
    dim: int = -1,
) -> torch.Tensor:
    """
    DCRP = joint constraint (不拆)：
      - tau: step/std <= tau  => clip <= qmax * tau * std
      - zero_ratio: approx P(q=0) <= zr => clip <= 2*qmax*Q_zr(|x|)
    只有当 tau 和 zero_ratio 都提供时才启用（避免“拆开”单独生效）。
    """
    if tau is None or zero_ratio is None:
        return torch.clamp(clip, min=1e-8)

    std = x.float().std(dim=dim, keepdim=True, unbiased=False) + 1e-12
    c_tau = float(qmax) * float(tau) * std.to(clip.dtype)

    zr = float(zero_ratio)
    zr = max(0.0, min(1.0, zr))
    thr = _kthvalue_quantile(x.abs(), zr, dim=dim, keepdim=True)
    c_zr = 2.0 * float(qmax) * thr.to(clip.dtype)

    out = torch.minimum(clip, c_tau)
    out = torch.minimum(out, c_zr)
    return torch.clamp(out, min=1e-8)


def _apply_star_constraints(
    clip: torch.Tensor,
    x: torch.Tensor,
    qmax: int,
    tau: Optional[float],
    zero_ratio: Optional[float],
    dim: int = -1,
) -> torch.Tensor:
    """Deprecated alias of `_apply_dcrp_constraints` for compatibility."""
    return _apply_dcrp_constraints(clip=clip, x=x, qmax=qmax, tau=tau, zero_ratio=zero_ratio, dim=dim)


def _keep_topk_by_abs(res: torch.Tensor, k: int, dim: int = -1) -> torch.Tensor:
    """
    EP residual topk：
      - k<=0: keep all
      - k>0: only keep top-k by |res| along `dim`
    """
    if k is None or int(k) <= 0:
        return res
    k = int(k)
    n = res.shape[dim]
    if k >= n:
        return res
    idx = res.abs().topk(k=k, dim=dim, largest=True).indices
    out = torch.zeros_like(res)
    out.scatter_(dim, idx, res.gather(dim, idx))
    return out


# -------------------------
# QuantLinear (core: DSTG + DCRP + optional EP)
# -------------------------
class QuantLinear(nn.Module):
    def __init__(
        self,
        linear: nn.Linear,
        w_bits: int = 4,
        a_bits: int = 4,
        # weight quant
        w_per_channel: bool = True,
        w_group_size: Optional[int] = None,
        w_clip_percentile: Optional[float] = None,
        # activation quant
        act_mode: str = "per_dstg_group",     # per_tensor / per_channel / per_dstg_group
        act_group_size: int = 32,              # DSTG group size
        dstg_group_size: Optional[int] = None,
        act_per_channel: bool = False,         # for per_tensor/per_channel fallback
        a_clip_percentile: Optional[float] = None,
        # DSTG clip
        dstg_clip_percentile: Optional[float] = None,
        # Deprecated alias
        token_clip_percentile: Optional[float] = None,
        # EP
        token_outlier_bypass: bool = False,
        token_outlier_topk: int = 0,           # 0=keep full residual; 4=engineer-friendly
        # DCRP (joint)
        dcrp_tau: Optional[float] = None,
        dcrp_zero_ratio: Optional[float] = None,
        # Deprecated aliases
        token_step_over_std: Optional[float] = None,  # tau
        token_zero_ratio: Optional[float] = None,     # zr
        compute_dtype: str = "fp16",
        full_name: Optional[str] = None,
    ):
        super().__init__()
        if not isinstance(linear, nn.Linear):
            raise TypeError("QuantLinear expects an nn.Linear")

        self.full_name = str(full_name) if full_name is not None else ""

        self.in_features = int(linear.in_features)
        self.out_features = int(linear.out_features)

        self.w_bits = int(w_bits)
        self.a_bits = int(a_bits)

        self.w_per_channel = bool(w_per_channel)
        self.w_group_size = int(w_group_size) if w_group_size is not None else None
        self.w_clip_percentile = w_clip_percentile

        mode = str(act_mode).strip().lower()
        if mode in ("per_token_group", "tokengroup", "token_group", "tgq", "dstg"):
            mode = "per_dstg_group"
        self.act_mode = mode
        if dstg_group_size is not None:
            act_group_size = int(dstg_group_size)
        self.act_group_size = int(act_group_size)
        self.dstg_group_size = int(self.act_group_size)
        self.act_per_channel = bool(act_per_channel)
        self.a_clip_percentile = a_clip_percentile

        if dstg_clip_percentile is None and token_clip_percentile is not None:
            dstg_clip_percentile = token_clip_percentile
        if token_clip_percentile is None and dstg_clip_percentile is not None:
            token_clip_percentile = dstg_clip_percentile
        self.dstg_clip_percentile = dstg_clip_percentile
        self.token_clip_percentile = token_clip_percentile
        self.token_outlier_bypass = bool(token_outlier_bypass)
        self.token_outlier_topk = int(token_outlier_topk)

        if dcrp_tau is None and token_step_over_std is not None:
            dcrp_tau = token_step_over_std
        if dcrp_zero_ratio is None and token_zero_ratio is not None:
            dcrp_zero_ratio = token_zero_ratio
        self.dcrp_tau = dcrp_tau
        self.dcrp_zero_ratio = dcrp_zero_ratio
        self.token_step_over_std = dcrp_tau
        self.token_zero_ratio = dcrp_zero_ratio

        self.compute_dtype = str(compute_dtype)

        # bias (fp32)
        if linear.bias is None:
            self.register_buffer("bias_fp", None)
        else:
            self.register_buffer("bias_fp", linear.bias.detach().to(torch.float32))

        # -------- weight quant (static) --------
        W = linear.weight.detach().to(torch.float32)

        if self.w_group_size is not None:
            gq = calc_group_qparams_symmetric(W, bits=self.w_bits, group_size=self.w_group_size)
            q_int8 = quantize_group_symmetric(W, gq)
            self.register_buffer("w_scale_group", gq.scale.to(torch.float32))
            self.register_buffer(
                "w_group_meta",
                torch.tensor([gq.group_size, gq.in_features, gq.padded_in], dtype=torch.int32, device=W.device),
            )
            if self.w_bits == 4:
                packed, shape = pack_int4(q_int8.to(torch.int8))
                self.register_buffer("w_q_packed", packed)
                self.register_buffer("w_q_shape", torch.tensor(list(shape), dtype=torch.int32, device=W.device))
                self.register_buffer("w_q_int8", None)
            else:
                self.register_buffer("w_q_int8", q_int8.to(torch.int8))
                self.register_buffer("w_q_packed", None)
                self.register_buffer("w_q_shape", None)
        else:
            qparams = calc_qparams_symmetric(
                W,
                bits=self.w_bits,
                per_channel=self.w_per_channel,
                ch_axis=0,
                clip_percentile=self.w_clip_percentile,
            )
            q_int8 = quantize_symmetric(W, qparams).to(torch.int8)
            self.register_buffer("w_scale", qparams.scale.to(torch.float32))
            self.register_buffer(
                "w_qparams_meta",
                torch.tensor([qparams.qmin, qparams.qmax, int(self.w_per_channel)], dtype=torch.int32, device=W.device),
            )
            if self.w_bits == 4:
                packed, shape = pack_int4(q_int8)
                self.register_buffer("w_q_packed", packed)
                self.register_buffer("w_q_shape", torch.tensor(list(shape), dtype=torch.int32, device=W.device))
                self.register_buffer("w_q_int8", None)
            else:
                self.register_buffer("w_q_int8", q_int8)
                self.register_buffer("w_q_packed", None)
                self.register_buffer("w_q_shape", None)

    def _dequant_weight(self) -> torch.Tensor:
        if self.w_group_size is not None:
            meta = self.w_group_meta
            group_size = int(meta[0].item())
            in_features = int(meta[1].item())
            padded_in = int(meta[2].item())

            if self.w_bits == 4:
                shape = tuple(int(x.item()) for x in self.w_q_shape)
                q = unpack_int4(self.w_q_packed, shape).to(torch.int8)
            else:
                q = self.w_q_int8

            gq = GroupQParams(
                scale=self.w_scale_group,
                qmax=(1 << (self.w_bits - 1)) - 1,
                qmin=-(1 << (self.w_bits - 1)),
                bits=self.w_bits,
                group_size=group_size,
                in_features=in_features,
                padded_in=padded_in,
            )
            return dequantize_group_symmetric(q, gq)

        qmin, qmax, per_ch = [int(x.item()) for x in self.w_qparams_meta]
        per_channel = bool(per_ch)
        if self.w_bits == 4:
            shape = tuple(int(x.item()) for x in self.w_q_shape)
            q = unpack_int4(self.w_q_packed, shape).to(torch.int8)
        else:
            q = self.w_q_int8

        qparams = QParams(
            scale=self.w_scale,
            qmax=qmax,
            qmin=qmin,
            bits=self.w_bits,
            per_channel=per_channel,
            ch_axis=0,
        )
        return dequantize_symmetric(q, qparams)

    def _quant_dequant_act_per_tensor(self, x: torch.Tensor) -> torch.Tensor:
        if self.a_bits >= 16:
            return x
        qparams = calc_qparams_symmetric(
            x,
            bits=self.a_bits,
            per_channel=self.act_per_channel,
            ch_axis=-1,
            clip_percentile=self.a_clip_percentile,
        )
        q = quantize_symmetric(x, qparams)
        return dequantize_symmetric(q, qparams).to(x.dtype)

    def _quant_dequant_act_dstg(self, x: torch.Tensor) -> torch.Tensor:
        if self.a_bits >= 16:
            return x
        if x.ndim < 2:
            return self._quant_dequant_act_per_tensor(x)

        qmin = -(1 << (self.a_bits - 1))
        qmax = (1 << (self.a_bits - 1)) - 1

        g = max(1, int(self.act_group_size))
        x_pad, C0 = _pad_last_dim(x, g)
        xg = x_pad.view(*x_pad.shape[:-1], x_pad.shape[-1] // g, g)  # [..., nG, g]

        clip = _dstg_clip_value(xg, self.dstg_clip_percentile, dim=-1)
        clip = _apply_dcrp_constraints(
            clip=clip,
            x=xg,
            qmax=qmax,
            tau=self.dcrp_tau,
            zero_ratio=self.dcrp_zero_ratio,
            dim=-1,
        )

        x_clip = xg.clamp(min=-clip, max=clip)

        if self.token_outlier_bypass:
            x_res = xg - x_clip
            if self.token_outlier_topk > 0:
                x_res = _keep_topk_by_abs(x_res, k=min(self.token_outlier_topk, x_res.shape[-1]), dim=-1)
        else:
            x_res = 0.0

        scale = torch.clamp(clip / float(qmax), min=1e-8)
        q = torch.round(x_clip / scale).clamp(qmin, qmax).to(torch.int8)
        x_hat = (q.to(torch.float32) * scale.to(torch.float32)).to(xg.dtype)

        y = (x_hat + x_res).view(*x_pad.shape)[..., :C0]
        return y

    def _quant_dequant_act(self, x: torch.Tensor) -> torch.Tensor:
        mode = str(self.act_mode).lower()
        if mode in ("per_dstg_group", "per_token_group", "tokengroup", "tgq"):
            return self._quant_dequant_act_dstg(x)
        return self._quant_dequant_act_per_tensor(x)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        xq = self._quant_dequant_act(x)
        W = self._dequant_weight()

        dtype = _get_compute_dtype(self.compute_dtype)
        xq = xq.to(dtype)
        W = W.to(dtype)
        b = None if self.bias_fp is None else self.bias_fp.to(dtype)
        return F.linear(xq, W, b)


# -------------------------
# QuantConv2d / QuantConv1d (simple)
# -------------------------
class QuantConv2d(nn.Module):
    def __init__(
        self,
        conv: nn.Conv2d,
        w_bits: int = 4,
        a_bits: int = 4,
        w_per_channel: bool = True,
        w_clip_percentile: Optional[float] = None,
        act_per_channel: bool = False,
        a_clip_percentile: Optional[float] = None,
        compute_dtype: str = "fp16",
        full_name: Optional[str] = None,
    ):
        super().__init__()
        if not isinstance(conv, nn.Conv2d):
            raise TypeError("QuantConv2d expects nn.Conv2d")

        self.full_name = str(full_name) if full_name is not None else ""

        self.w_bits = int(w_bits)
        self.a_bits = int(a_bits)
        self.w_per_channel = bool(w_per_channel)
        self.w_clip_percentile = w_clip_percentile
        self.act_per_channel = bool(act_per_channel)
        self.a_clip_percentile = a_clip_percentile
        self.compute_dtype = str(compute_dtype)

        self.stride = conv.stride
        self.padding = conv.padding
        self.dilation = conv.dilation
        self.groups = conv.groups

        if conv.bias is None:
            self.register_buffer("bias_fp", None)
        else:
            self.register_buffer("bias_fp", conv.bias.detach().to(torch.float32))

        W = conv.weight.detach().to(torch.float32)
        qparams = calc_qparams_symmetric(
            W, bits=self.w_bits, per_channel=self.w_per_channel, ch_axis=0, clip_percentile=self.w_clip_percentile
        )
        q_int8 = quantize_symmetric(W, qparams).to(torch.int8)

        self.register_buffer("w_scale", qparams.scale.to(torch.float32))
        self.register_buffer(
            "w_qparams_meta",
            torch.tensor([qparams.qmin, qparams.qmax, int(self.w_per_channel)], dtype=torch.int32, device=W.device),
        )

        if self.w_bits == 4:
            packed, shape = pack_int4(q_int8)
            self.register_buffer("w_q_packed", packed)
            self.register_buffer("w_q_shape", torch.tensor(list(shape), dtype=torch.int32, device=W.device))
            self.register_buffer("w_q_int8", None)
        else:
            self.register_buffer("w_q_int8", q_int8)
            self.register_buffer("w_q_packed", None)
            self.register_buffer("w_q_shape", None)

    def _dequant_weight(self) -> torch.Tensor:
        qmin, qmax, per_ch = [int(x.item()) for x in self.w_qparams_meta]
        per_channel = bool(per_ch)
        if self.w_bits == 4:
            shape = tuple(int(x.item()) for x in self.w_q_shape)
            q = unpack_int4(self.w_q_packed, shape).to(torch.int8)
        else:
            q = self.w_q_int8
        qparams = QParams(scale=self.w_scale, qmax=qmax, qmin=qmin, bits=self.w_bits, per_channel=per_channel, ch_axis=0)
        return dequantize_symmetric(q, qparams)

    def _quant_dequant_act(self, x: torch.Tensor) -> torch.Tensor:
        if self.a_bits >= 16:
            return x
        qparams = calc_qparams_symmetric(
            x, bits=self.a_bits, per_channel=self.act_per_channel, ch_axis=1, clip_percentile=self.a_clip_percentile
        )
        q = quantize_symmetric(x, qparams)
        return dequantize_symmetric(q, qparams).to(x.dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        xq = self._quant_dequant_act(x)
        W = self._dequant_weight()

        dtype = _get_compute_dtype(self.compute_dtype)
        xq = xq.to(dtype)
        W = W.to(dtype)
        b = None if self.bias_fp is None else self.bias_fp.to(dtype)

        return F.conv2d(xq, W, b, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=self.groups)


class QuantConv1d(nn.Module):
    def __init__(
        self,
        conv: nn.Conv1d,
        w_bits: int = 4,
        a_bits: int = 4,
        w_per_channel: bool = True,
        w_clip_percentile: Optional[float] = None,
        act_per_channel: bool = False,
        a_clip_percentile: Optional[float] = None,
        compute_dtype: str = "fp16",
        full_name: Optional[str] = None,
    ):
        super().__init__()
        if not isinstance(conv, nn.Conv1d):
            raise TypeError("QuantConv1d expects nn.Conv1d")

        self.full_name = str(full_name) if full_name is not None else ""

        self.w_bits = int(w_bits)
        self.a_bits = int(a_bits)
        self.w_per_channel = bool(w_per_channel)
        self.w_clip_percentile = w_clip_percentile
        self.act_per_channel = bool(act_per_channel)
        self.a_clip_percentile = a_clip_percentile
        self.compute_dtype = str(compute_dtype)

        self.stride = conv.stride
        self.padding = conv.padding
        self.dilation = conv.dilation
        self.groups = conv.groups

        if conv.bias is None:
            self.register_buffer("bias_fp", None)
        else:
            self.register_buffer("bias_fp", conv.bias.detach().to(torch.float32))

        W = conv.weight.detach().to(torch.float32)
        qparams = calc_qparams_symmetric(
            W, bits=self.w_bits, per_channel=self.w_per_channel, ch_axis=0, clip_percentile=self.w_clip_percentile
        )
        q_int8 = quantize_symmetric(W, qparams).to(torch.int8)

        self.register_buffer("w_scale", qparams.scale.to(torch.float32))
        self.register_buffer(
            "w_qparams_meta",
            torch.tensor([qparams.qmin, qparams.qmax, int(self.w_per_channel)], dtype=torch.int32, device=W.device),
        )

        if self.w_bits == 4:
            packed, shape = pack_int4(q_int8)
            self.register_buffer("w_q_packed", packed)
            self.register_buffer("w_q_shape", torch.tensor(list(shape), dtype=torch.int32, device=W.device))
            self.register_buffer("w_q_int8", None)
        else:
            self.register_buffer("w_q_int8", q_int8)
            self.register_buffer("w_q_packed", None)
            self.register_buffer("w_q_shape", None)

    def _dequant_weight(self) -> torch.Tensor:
        qmin, qmax, per_ch = [int(x.item()) for x in self.w_qparams_meta]
        per_channel = bool(per_ch)
        if self.w_bits == 4:
            shape = tuple(int(x.item()) for x in self.w_q_shape)
            q = unpack_int4(self.w_q_packed, shape).to(torch.int8)
        else:
            q = self.w_q_int8
        qparams = QParams(scale=self.w_scale, qmax=qmax, qmin=qmin, bits=self.w_bits, per_channel=per_channel, ch_axis=0)
        return dequantize_symmetric(q, qparams)

    def _quant_dequant_act(self, x: torch.Tensor) -> torch.Tensor:
        if self.a_bits >= 16:
            return x
        qparams = calc_qparams_symmetric(
            x, bits=self.a_bits, per_channel=self.act_per_channel, ch_axis=1, clip_percentile=self.a_clip_percentile
        )
        q = quantize_symmetric(x, qparams)
        return dequantize_symmetric(q, qparams).to(x.dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        xq = self._quant_dequant_act(x)
        W = self._dequant_weight()

        dtype = _get_compute_dtype(self.compute_dtype)
        xq = xq.to(dtype)
        W = W.to(dtype)
        b = None if self.bias_fp is None else self.bias_fp.to(dtype)

        return F.conv1d(xq, W, b, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=self.groups)