# quantization/escnet/quant/quantizers.py
from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Tuple

import torch

from .config import QuantConfig
from .observers import MinMaxObserver
from .dstg_dcrp import dstg_dcrp_qdq


def _qmax(bits: int) -> int:
    return (1 << (bits - 1)) - 1


def symmetric_qdq(x: torch.Tensor, scale: torch.Tensor, qmax: int) -> torch.Tensor:
    """Symmetric QDQ: round(x/scale) * scale."""
    orig_dtype = x.dtype
    xf = x.float()
    sf = scale.float()
    # broadcast
    q = torch.round(xf / sf).clamp(-qmax, qmax)
    return (q * sf).to(orig_dtype)


@dataclass
class StaticActQuant:
    """Static activation QDQ (needs calibration)."""
    bits: int
    mode: str  # per_tensor | per_channel
    ch_axis: int

    observer: MinMaxObserver
    scale: Optional[torch.Tensor] = None  # on CPU by default

    def __post_init__(self) -> None:
        if self.mode not in ("per_tensor", "per_channel"):
            raise ValueError(f"Unsupported static act mode: {self.mode}")

    @torch.no_grad()
    def observe(self, x: torch.Tensor) -> None:
        self.observer.update(x)

    @torch.no_grad()
    def finalize(self, device: torch.device) -> None:
        absmax = self.observer.get_absmax(device=device).clamp(min=1e-8)
        self.scale = (absmax / float(_qmax(self.bits))).detach()

    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        if self.scale is None:
            # pre-finalize fallback: identity
            return x
        scale = self.scale.to(x.device)
        return symmetric_qdq(x, scale=scale, qmax=_qmax(self.bits))


@dataclass
class DynamicDSTGDCRPActQuant:
    """Dynamic per-DSTG-group + DCRP (no calibration needed)."""
    bits: int
    group_size: int
    dstg_clip_percentile: Optional[float]
    dcrp_tau: Optional[float]
    dcrp_zero_ratio: Optional[float]
    residual_topk: int = 0

    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        return dstg_dcrp_qdq(
            x,
            qmax=_qmax(self.bits),
            group_size=self.group_size,
            dstg_clip_percentile=self.dstg_clip_percentile,
            tau=self.dcrp_tau,
            zero_ratio=self.dcrp_zero_ratio,
            residual_topk=self.residual_topk,
        )


# Backward compatibility alias
DynamicTGQStarActQuant = DynamicDSTGDCRPActQuant


@torch.no_grad()
def quantize_weight_symmetric(
    w: torch.Tensor, bits: int, mode: str = "per_channel", ch_axis: int = 0
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Return (w_int8, w_scale). Symmetric quant, scale is float32."""
    qmax = _qmax(bits)
    wf = w.detach().float()
    if mode == "per_tensor":
        absmax = wf.abs().max().clamp(min=1e-8)
        scale = absmax / float(qmax)
        wq = torch.round(wf / scale).clamp(-qmax, qmax).to(torch.int8)
        return wq, torch.tensor(scale, dtype=torch.float32)
    if mode == "per_channel":
        # reduce all dims except ch_axis
        dims = [d for d in range(wf.dim()) if d != ch_axis]
        absmax = wf.abs().amax(dim=dims).clamp(min=1e-8)
        scale = absmax / float(qmax)  # [Cout]
        # reshape for broadcasting
        shape = [1] * wf.dim()
        shape[ch_axis] = -1
        scale_bc = scale.view(*shape)
        wq = torch.round(wf / scale_bc).clamp(-qmax, qmax).to(torch.int8)
        return wq, scale.to(torch.float32)
    raise ValueError(f"Unsupported weight mode: {mode}")


def make_activation_quantizer(qcfg: QuantConfig, x_example: Optional[torch.Tensor] = None):
    """Factory for layer-local activation quantizer."""
    if qcfg.a_mode in ("per_dstg_group", "per_token_group"):
        return DynamicDSTGDCRPActQuant(
            bits=qcfg.a_bits,
            group_size=qcfg.group_size,
            dstg_clip_percentile=qcfg.dstg_clip_percentile,
            dcrp_tau=qcfg.dcrp_tau,
            dcrp_zero_ratio=qcfg.dcrp_zero_ratio,
            residual_topk=qcfg.ep_residual_topk,
        )

    if qcfg.a_mode == "per_tensor":
        return StaticActQuant(
            bits=qcfg.a_bits,
            mode="per_tensor",
            ch_axis=1,
            observer=MinMaxObserver(mode="per_tensor", ch_axis=1),
        )

    if qcfg.a_mode == "per_channel":
        # assume channel axis: for Conv it's dim=1, for Linear inputs it's last dim
        ch_axis = 1
        if x_example is not None and x_example.dim() == 2:
            ch_axis = -1
        return StaticActQuant(
            bits=qcfg.a_bits,
            mode="per_channel",
            ch_axis=ch_axis,
            observer=MinMaxObserver(mode="per_channel", ch_axis=ch_axis),
        )

    raise ValueError(f"Unknown a_mode: {qcfg.a_mode}")