# quantization/cfrn/quantlib/quant_utils.py

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional, Tuple

import torch

# ---------------------------
# Helper: int4 pack / unpack
# ---------------------------
def pack_int4(q_int8: torch.Tensor) -> Tuple[torch.Tensor, torch.Size]:
    """
    Pack int4 values stored in int8 tensor (expected range [-8,7]) into uint8.
    Two int4 values per byte: low nibble is first, high nibble is second.

    Returns: packed_uint8, original_shape
    """
    if q_int8.dtype != torch.int8:
        raise TypeError(f"pack_int4 expects int8, got {q_int8.dtype}")
    shape = q_int8.shape
    flat = q_int8.reshape(-1).to(torch.int16)  # avoid overflow
    # map [-8,7] -> [0,15]
    flat = flat + 8
    flat = torch.clamp(flat, 0, 15).to(torch.uint8)

    if flat.numel() % 2 == 1:
        flat = torch.cat([flat, torch.zeros(1, device=flat.device, dtype=flat.dtype)], dim=0)

    lo = flat[0::2]
    hi = flat[1::2]
    packed = lo | (hi << 4)
    return packed, shape


def unpack_int4(packed_uint8: torch.Tensor, shape: torch.Size) -> torch.Tensor:
    """
    Unpack uint8 -> int8 int4 values in [-8,7], restoring given shape.
    """
    if packed_uint8.dtype != torch.uint8:
        raise TypeError(f"unpack_int4 expects uint8, got {packed_uint8.dtype}")
    flat = packed_uint8.reshape(-1)
    lo = flat & 0x0F
    hi = (flat >> 4) & 0x0F
    out = torch.empty((flat.numel() * 2,), device=flat.device, dtype=torch.int16)
    out[0::2] = lo.to(torch.int16)
    out[1::2] = hi.to(torch.int16)
    out = out[: int(torch.tensor(shape).prod().item())]
    out = out.to(torch.int16) - 8
    return out.to(torch.int8).reshape(shape)


# ---------------------------
# Symmetric quant/dequant
# ---------------------------
@dataclass
class QParams:
    scale: torch.Tensor  # float tensor, can be scalar or per-channel/group
    qmax: int
    qmin: int
    bits: int
    per_channel: bool = False
    ch_axis: int = 0
    group_size: Optional[int] = None  # for group-wise along last dim (Linear style)


def _signed_qrange(bits: int) -> Tuple[int, int]:
    qmax = (1 << (bits - 1)) - 1
    qmin = - (1 << (bits - 1))
    return qmin, qmax


@torch.no_grad()
def calc_qparams_symmetric(
    x: torch.Tensor,
    bits: int,
    per_channel: bool,
    ch_axis: int,
    clip_percentile: Optional[float] = None,
    sample: Optional[torch.Tensor] = None,
    eps: float = 1e-8,
) -> QParams:
    """
    Compute symmetric scale such that x/scale is in [qmin,qmax].
    If clip_percentile is provided, use percentile of |x| instead of max(|x|).
    For percentile: pass a 1D sample of abs(x) values (optional); otherwise uses abs(x).flatten() which can be heavy.
    """
    qmin, qmax = _signed_qrange(bits)

    if per_channel:
        # reduce all dims except ch_axis
        dims = list(range(x.ndim))
        dims.pop(ch_axis)
        if clip_percentile is None:
            max_abs = x.abs().amax(dim=dims)
        else:
            # approximate percentile per-channel by using max as fallback (cheap)
            # For exact percentile per-channel, provide 'sample' externally.
            max_abs = x.abs().amax(dim=dims)
        scale = max_abs / float(qmax)
    else:
        if clip_percentile is None:
            max_abs = x.abs().max()
        else:
            if sample is None:
                # approximate: take a smaller random sample
                flat = x.abs().flatten()
                if flat.numel() > 4096:
                    idx = torch.randint(0, flat.numel(), (4096,), device=flat.device)
                    flat = flat[idx]
            else:
                flat = sample
            k = max(0.0, min(1.0, float(clip_percentile)))
            max_abs = torch.quantile(flat, k)
        scale = max_abs / float(qmax)

    scale = torch.clamp(scale, min=eps)
    return QParams(scale=scale, qmax=qmax, qmin=qmin, bits=bits, per_channel=per_channel, ch_axis=ch_axis, group_size=None)


@torch.no_grad()
def quantize_symmetric(x: torch.Tensor, qparams: QParams) -> torch.Tensor:
    """
    Return int8 tensor containing quantized values (still stored as int8 even for int4).
    """
    scale = qparams.scale
    if qparams.per_channel:
        view_shape = [1] * x.ndim
        view_shape[qparams.ch_axis] = -1
        scale_b = scale.view(view_shape)
    else:
        scale_b = scale
    q = torch.round(x / scale_b)
    q = torch.clamp(q, qparams.qmin, qparams.qmax)
    return q.to(torch.int8)


@torch.no_grad()
def dequantize_symmetric(q: torch.Tensor, qparams: QParams) -> torch.Tensor:
    scale = qparams.scale
    if qparams.per_channel:
        view_shape = [1] * q.ndim
        view_shape[qparams.ch_axis] = -1
        scale_b = scale.view(view_shape)
    else:
        scale_b = scale
    return q.to(torch.float32) * scale_b.to(torch.float32)


# ---------------------------
# Group-wise weight quant (LLM-style, along last dim of weight row)
# ---------------------------
@dataclass
class GroupQParams:
    scale: torch.Tensor  # [out, n_groups]
    qmax: int
    qmin: int
    bits: int
    group_size: int
    in_features: int
    padded_in: int


@torch.no_grad()
def calc_group_qparams_symmetric(W_2d: torch.Tensor, bits: int, group_size: int, eps: float = 1e-8) -> GroupQParams:
    """
    W_2d shape [out, in]. Group along in dimension for each out row.
    """
    if W_2d.ndim != 2:
        raise ValueError("calc_group_qparams_symmetric expects 2D weight")
    out, in_features = W_2d.shape
    qmin, qmax = _signed_qrange(bits)
    padded_in = int(math.ceil(in_features / group_size) * group_size)
    if padded_in != in_features:
        pad = padded_in - in_features
        Wp = torch.cat([W_2d, torch.zeros((out, pad), device=W_2d.device, dtype=W_2d.dtype)], dim=1)
    else:
        Wp = W_2d
    Wg = Wp.view(out, padded_in // group_size, group_size)
    max_abs = Wg.abs().amax(dim=2)  # [out, n_groups]
    scale = torch.clamp(max_abs / float(qmax), min=eps)
    return GroupQParams(scale=scale, qmax=qmax, qmin=qmin, bits=bits, group_size=group_size, in_features=in_features, padded_in=padded_in)


@torch.no_grad()
def quantize_group_symmetric(W_2d: torch.Tensor, qparams: GroupQParams) -> torch.Tensor:
    out, in_features = W_2d.shape
    if in_features != qparams.in_features:
        raise ValueError("in_features mismatch")
    padded_in = qparams.padded_in
    if padded_in != in_features:
        pad = padded_in - in_features
        Wp = torch.cat([W_2d, torch.zeros((out, pad), device=W_2d.device, dtype=W_2d.dtype)], dim=1)
    else:
        Wp = W_2d
    Wg = Wp.view(out, padded_in // qparams.group_size, qparams.group_size)
    scale = qparams.scale.unsqueeze(-1)  # [out, n_groups, 1]
    q = torch.round(Wg / scale)
    q = torch.clamp(q, qparams.qmin, qparams.qmax).to(torch.int8)
    return q.view(out, padded_in)


@torch.no_grad()
def dequantize_group_symmetric(q_int8: torch.Tensor, qparams: GroupQParams) -> torch.Tensor:
    out, padded_in = q_int8.shape
    if padded_in != qparams.padded_in:
        raise ValueError("padded_in mismatch")
    qg = q_int8.view(out, padded_in // qparams.group_size, qparams.group_size).to(torch.float32)
    scale = qparams.scale.unsqueeze(-1).to(torch.float32)
    W = qg * scale
    W = W.view(out, padded_in)[:, : qparams.in_features]
    return W


# ---------------------------
# Prob/Softmax power-of-two quantization
# ---------------------------
@torch.no_grad()
def quantize_prob_power2(p: torch.Tensor, k: int) -> torch.Tensor:
    """
    Quantize probabilities p in [0,1] to multiples of 1/2^k.
    """
    k = int(k)
    base = 1 << k
    q = torch.round(p * base) / base
    return torch.clamp(q, 0.0, 1.0)