"""SOFIA v7.1 quantization utilities.

Two reusable building blocks for INT8 deployment:

1. **OffsetClampSTE** (DCN-M2): hard clamp DCN offsets to physical range
   ``[-k, +k]`` with straight-through gradient. Bounded distribution =
   stable percentile clipping during PTQ + cleaner offset semantics.

2. **KScaledFakeQuant** (R5-style): multi-bin fake quantization. Bins are
   chosen by magnitude percentiles; each bin has its own scale. Adapts
   quantization resolution to long-tail distributions (Mamba Δ, output y).

Plus a drop-in:

3. **KScaledMamba2Block**: ``Mamba2Block`` subclass with k-scaled fake-quant
   nodes on ``x_main``, ``delta``, and per-token ``y`` outputs. Implements
   R5 reparam principle (state ``h_t`` kept in model dtype; only observable
   tensors quantized).

All components are PTQ/QAT-compatible (STE backward, calibration mode).
"""

from __future__ import annotations

from typing import List, Optional, Sequence, Tuple

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

from .blocks import Mamba2Block


# ============================================================
# DCN-M2: Offset clamp with STE
# ============================================================

class _OffsetClampSTE(torch.autograd.Function):
    """Hard clamp [low, high] in forward, identity gradient in backward."""

    @staticmethod
    def forward(ctx, x: torch.Tensor, low: float, high: float) -> torch.Tensor:
        return torch.clamp(x, low, high)

    @staticmethod
    def backward(ctx, grad: torch.Tensor):
        # STE: pass gradient unmodified
        return grad, None, None


def offset_clamp_ste(offsets: torch.Tensor, kernel_size: int) -> torch.Tensor:
    """Functional DCN-M2: clamp offsets to ``[-k, +k]`` with STE backward.

    Args:
        offsets: arbitrary-shape tensor of DCN offsets.
        kernel_size: physical receptive-field bound k. Output is clamped to
            absolute values <= k. Standard choice equals the DW kernel size.

    Returns:
        Clamped tensor (same shape as ``offsets``). Backward propagates
        gradient as if no clamp was applied, so training can still adjust
        weights even when an offset hits the boundary.
    """
    k = float(kernel_size)
    return _OffsetClampSTE.apply(offsets, -k, k)


class OffsetClampSTE(nn.Module):
    """Module wrapper for :func:`offset_clamp_ste`.

    Use as a drop-in layer between ``offset_predictor`` and ``deform_conv2d``::

        self.clamp = OffsetClampSTE(kernel_size=7)
        offsets = self.clamp(self.offset_predictor(x))
    """

    def __init__(self, kernel_size: int) -> None:
        super().__init__()
        self.kernel_size = int(kernel_size)

    def forward(self, offsets: torch.Tensor) -> torch.Tensor:
        return offset_clamp_ste(offsets, self.kernel_size)

    def extra_repr(self) -> str:
        return f"kernel_size={self.kernel_size}, range=[-{self.kernel_size}, +{self.kernel_size}]"


# ============================================================
# K-Scaled Fake Quantization (R5 style, multi-bin)
# ============================================================

class _STEFakeQuant(torch.autograd.Function):
    """Fake-quantize: ``round(x / s) * s`` clamped to [qmin, qmax], STE backward.

    `scale` is a tensor broadcastable with `x` (per-tensor scalar or
    per-element via gathered per-bin scales).
    """

    @staticmethod
    def forward(
        ctx,
        x: torch.Tensor,
        scale: torch.Tensor,
        qmin: int,
        qmax: int,
    ) -> torch.Tensor:
        x_int = torch.round(x / scale).clamp(qmin, qmax)
        return x_int * scale

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        return grad_output, None, None, None


class KScaledFakeQuant(nn.Module):
    """k-scaled (multi-bin magnitude-bucketed) fake quantization.

    Algorithm:
        1. ``g(x) = bin index in [0..k-1] based on |x|`` using thresholds
           ``t_0 < t_1 < ... < t_{k-2}``.
        2. ``s(x) = scales[g(x)]`` per-element scale.
        3. ``x_q = round(x / s) * s`` clamped to ``[qmin, qmax]``.
        4. STE backward.

    Setting ``num_bins=1`` reduces to standard single-scale fake quant.

    Calibration workflow::

        fq = KScaledFakeQuant(num_bins=3)
        fq.start_calibration()
        for batch in cal_loader:
            _ = model(batch)        # collect stats
        fq.finalize_calibration(percentile=99.9)
        # fq is now active (calibrating=False) and applies fake-quant in forward
    """

    def __init__(
        self,
        num_bins: int = 3,
        qmin: int = -128,
        qmax: int = 127,
        max_calibration_samples: int = 200_000,
    ) -> None:
        super().__init__()
        if num_bins < 1:
            raise ValueError(f"num_bins must be >= 1, got {num_bins}")
        if qmax <= qmin:
            raise ValueError(f"qmax ({qmax}) must be > qmin ({qmin})")

        self.num_bins = num_bins
        self.qmin = qmin
        self.qmax = qmax
        self.max_calibration_samples = max_calibration_samples

        # Thresholds: (num_bins - 1,) — bin boundaries on |x|. For num_bins=1, empty.
        n_thr = max(0, num_bins - 1)
        self.register_buffer("thresholds", torch.zeros(n_thr))

        # Scales: (num_bins,) — one scale per bin. Default 1.0 = no quant effect.
        self.register_buffer("scales", torch.ones(num_bins))

        # Per-instance state (not buffers — runtime-only)
        self.calibrating: bool = False
        self.enabled: bool = True
        self._calib_buffer: List[torch.Tensor] = []
        self._n_collected: int = 0

    # -------------------- Calibration API --------------------

    def start_calibration(self) -> None:
        """Begin collecting input distribution stats."""
        self.calibrating = True
        self._calib_buffer = []
        self._n_collected = 0

    def stop_calibration(self) -> None:
        """End calibration without finalizing scales (e.g. abort)."""
        self.calibrating = False
        self._calib_buffer = []
        self._n_collected = 0

    def finalize_calibration(self, percentile: float = 99.9) -> None:
        """Compute thresholds and per-bin scales from collected stats.

        Args:
            percentile: per-bin tail percentile for scale computation
                (typical 99.9). Robust to extreme outliers.
        """
        if not self._calib_buffer:
            self.calibrating = False
            return

        all_data = torch.cat(self._calib_buffer)
        abs_data = all_data.abs()

        if self.num_bins == 1:
            # Single scale: percentile of |x|
            tail = torch.quantile(abs_data, percentile / 100.0).item()
            self.scales[0] = max(tail / float(self.qmax), 1e-8)
        else:
            # Threshold positions: equal-mass quantile splits
            qpts = torch.linspace(0.0, 1.0, self.num_bins + 1, device=abs_data.device)[1:-1]
            thr_vals = torch.quantile(abs_data, qpts)
            self.thresholds.copy_(thr_vals.to(self.thresholds.device))

            # Per-bin scale: per-bin tail percentile / qmax
            for g in range(self.num_bins):
                if g == 0:
                    lo = 0.0
                    hi = thr_vals[0].item()
                elif g == self.num_bins - 1:
                    lo = thr_vals[-1].item()
                    hi = float("inf")
                else:
                    lo = thr_vals[g - 1].item()
                    hi = thr_vals[g].item()

                mask = (abs_data >= lo) & (abs_data < hi)
                if mask.any():
                    bin_tail = torch.quantile(
                        abs_data[mask], percentile / 100.0
                    ).item()
                    self.scales[g] = max(bin_tail / float(self.qmax), 1e-8)
                else:
                    self.scales[g] = 1.0

        self._calib_buffer = []
        self._n_collected = 0
        self.calibrating = False

    # -------------------- Forward --------------------

    def _bin_index(self, x: torch.Tensor) -> torch.Tensor:
        """Per-element bin index (long tensor, shape == x.shape)."""
        if self.num_bins == 1:
            return torch.zeros_like(x, dtype=torch.long)

        abs_x = x.abs()
        idx = torch.zeros_like(x, dtype=torch.long)
        for i in range(self.num_bins - 1):
            t = self.thresholds[i]
            idx = torch.where(abs_x >= t, idx + 1, idx)
        return idx

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not self.enabled:
            return x

        if self.calibrating:
            # Subsample to bound memory
            with torch.no_grad():
                flat = x.detach().flatten()
                budget = self.max_calibration_samples - self._n_collected
                if budget > 0:
                    take = min(flat.numel(), budget)
                    if take < flat.numel():
                        # Random subsample for representativeness
                        idx = torch.randperm(flat.numel(), device=flat.device)[:take]
                        sample = flat[idx]
                    else:
                        sample = flat
                    self._calib_buffer.append(sample.cpu())
                    self._n_collected += sample.numel()
            return x

        # Apply k-scaled fake-quant
        bin_idx = self._bin_index(x)
        scale = self.scales[bin_idx]    # shape == x.shape
        return _STEFakeQuant.apply(x, scale, self.qmin, self.qmax)

    def extra_repr(self) -> str:
        return (
            f"num_bins={self.num_bins}, qmin={self.qmin}, qmax={self.qmax}, "
            f"calibrating={self.calibrating}, enabled={self.enabled}"
        )


# ============================================================
# K-Scaled Mamba-2 drop-in
# ============================================================

class KScaledMamba2Block(Mamba2Block):
    """Drop-in replacement for :class:`Mamba2Block` with k-scaled fake-quant.

    Adds ``KScaledFakeQuant`` nodes on three observable tensors inside the
    Mamba-2 forward pass:

    - ``x_main``: scan input after conv1d + SiLU
    - ``delta``: time-step ``Δ`` after softplus (addresses MF2 long tail)
    - ``y``:      per-token scan output (addresses MF4 state propagation)

    The recurrent state ``h_t`` is **not** quantized — kept in model dtype
    inside the scan loop. This implements R5's reparam principle: only
    observable I/O is quantized; internal state stays high-precision.

    Backend constraint: forces ``backend='torch'``. The mamba_ssm CUDA
    backend would require k-scaled support inside the custom kernel (out of
    scope here). For a fast deploy path, train+QAT in torch, then export to
    a custom TRT plugin that does the same INT8 scan.

    Args:
        channels: model dim ``C``.
        d_state: state size ``N`` (Mamba-2 default 64).
        headdim: head dim (must divide channels).
        d_conv: local conv kernel.
        expand: inner expansion factor.
        backend: must be ``'torch'`` (or ``'auto'``, which resolves to torch).
        num_bins: ``k`` for k-scaled (typical 2-4).
        targets: which tensors to quantize. Subset of
            ``('x_main', 'delta', 'y')``.

    Example::

        mamba = KScaledMamba2Block(channels=192, num_bins=3,
                                   targets=('delta', 'y'))
        mamba.start_calibration()
        for batch in cal_loader:
            _ = model_with_mamba(batch)
        mamba.finalize_calibration()
        # Now in PTQ mode — forward applies fake-quant on chosen targets.
    """

    SUPPORTED_TARGETS: Tuple[str, ...] = ("x_main", "delta", "y")

    def __init__(
        self,
        channels: int,
        d_state: int = 64,
        headdim: int = 64,
        d_conv: int = 4,
        expand: int = 2,
        backend: str = "torch",
        num_bins: int = 3,
        targets: Sequence[str] = ("x_main", "delta", "y"),
    ) -> None:
        # Force torch — k-scaled requires explicit forward-pass control
        if backend == "auto":
            backend = "torch"
        if backend != "torch":
            raise ValueError(
                "KScaledMamba2Block currently supports only backend='torch'. "
                f"Got backend='{backend}'. The mamba_ssm CUDA kernel does not "
                "expose hooks for k-scaled fake-quant; integrating INT8 scan "
                "via TRT plugin is a separate deploy-time concern."
            )

        super().__init__(
            channels=channels,
            d_state=d_state,
            headdim=headdim,
            d_conv=d_conv,
            expand=expand,
            backend="torch",
        )

        unknown = set(targets) - set(self.SUPPORTED_TARGETS)
        if unknown:
            raise ValueError(
                f"Unknown k-scaled targets: {unknown}. "
                f"Supported: {self.SUPPORTED_TARGETS}"
            )
        self.targets = set(targets)
        self.num_bins = num_bins

        # Build fake-quant nodes only for active targets
        if "x_main" in self.targets:
            self.fq_x_main = KScaledFakeQuant(num_bins=num_bins)
        if "delta" in self.targets:
            self.fq_delta = KScaledFakeQuant(num_bins=num_bins)
        if "y" in self.targets:
            self.fq_y = KScaledFakeQuant(num_bins=num_bins)

    # -------------------- Calibration helpers --------------------

    def _all_fqs(self) -> List[KScaledFakeQuant]:
        out: List[KScaledFakeQuant] = []
        for name in self.SUPPORTED_TARGETS:
            attr = f"fq_{name}"
            mod = getattr(self, attr, None)
            if isinstance(mod, KScaledFakeQuant):
                out.append(mod)
        return out

    def start_calibration(self) -> None:
        for fq in self._all_fqs():
            fq.start_calibration()

    def finalize_calibration(self, percentile: float = 99.9) -> None:
        for fq in self._all_fqs():
            fq.finalize_calibration(percentile=percentile)

    def set_quant_enabled(self, enabled: bool) -> None:
        for fq in self._all_fqs():
            fq.enabled = enabled

    # -------------------- Forward --------------------

    def _forward_torch(self, x: torch.Tensor) -> torch.Tensor:
        """Mamba-2 torch forward with k-scaled fake-quant on selected paths."""
        B, L, C = x.shape

        # In projection (matches parent layout)
        z_in = self.in_proj(x)
        xz, BC, dt = torch.split(
            z_in,
            [2 * self.d_inner, 2 * self.d_state, self.nheads],
            dim=-1,
        )
        x_main, z_gate = xz.chunk(2, dim=-1)
        B_p, C_p = BC.chunk(2, dim=-1)

        # Local conv + SiLU
        x_main_t = x_main.transpose(1, 2)
        x_main_t = self.conv1d(x_main_t)[..., :L]
        x_main = F.silu(x_main_t).transpose(1, 2)

        # K-SCALED QUANT: x_main (scan input)
        if "x_main" in self.targets:
            x_main = self.fq_x_main(x_main)

        # Δ
        dt = F.softplus(dt)
        if "delta" in self.targets:
            dt = self.fq_delta(dt)

        A = -torch.exp(self.A_log)    # [nheads]
        x_head = x_main.view(B, L, self.nheads, self.headdim)

        # Sequential scan. R5 reparam: state h kept in model dtype throughout.
        h = torch.zeros(
            B, self.nheads, self.headdim, self.d_state,
            device=x.device, dtype=x.dtype,
        )
        ys: List[torch.Tensor] = []
        for t in range(L):
            dt_t = dt[:, t]
            B_t = B_p[:, t]
            C_t = C_p[:, t]
            x_t = x_head[:, t]

            dA = torch.exp(dt_t * A.unsqueeze(0)).unsqueeze(-1).unsqueeze(-1)
            dB = (dt_t.unsqueeze(-1) * B_t.unsqueeze(1)).unsqueeze(2)

            h = dA * h + dB * x_t.unsqueeze(-1)
            y_t = (h * C_t.unsqueeze(1).unsqueeze(1)).sum(dim=-1)
            y_t = y_t + self.D.unsqueeze(0).unsqueeze(-1) * x_t

            # K-SCALED QUANT: per-token y (after scan, before gate)
            if "y" in self.targets:
                y_t = self.fq_y(y_t)

            ys.append(y_t)

        y = torch.stack(ys, dim=1)
        y = y.reshape(B, L, self.d_inner)
        y = y * F.silu(z_gate)
        y = self.out_proj(y)
        return y


# ============================================================
# Model-wide calibration helpers
# ============================================================

def start_calibration(model: nn.Module) -> None:
    """Walk model, start calibration on every quant module."""
    seen = set()
    # First, KScaledMamba2Block instances drive their internal FQs.
    for m in model.modules():
        if isinstance(m, KScaledMamba2Block):
            m.start_calibration()
            for fq in m._all_fqs():
                seen.add(id(fq))
    # Then standalone KScaledFakeQuant not owned by a KScaledMamba2Block.
    for m in model.modules():
        if isinstance(m, KScaledFakeQuant) and id(m) not in seen:
            m.start_calibration()


def finalize_calibration(model: nn.Module, percentile: float = 99.9) -> None:
    """Walk model, finalize all quant modules with the same percentile."""
    seen = set()
    for m in model.modules():
        if isinstance(m, KScaledMamba2Block):
            m.finalize_calibration(percentile=percentile)
            for fq in m._all_fqs():
                seen.add(id(fq))
    for m in model.modules():
        if isinstance(m, KScaledFakeQuant) and id(m) not in seen:
            m.finalize_calibration(percentile=percentile)


def set_quant_enabled(model: nn.Module, enabled: bool) -> None:
    """Toggle every quant module on/off. Useful for FP-vs-INT8 comparison."""
    for m in model.modules():
        if isinstance(m, KScaledFakeQuant):
            m.enabled = enabled
        elif isinstance(m, KScaledMamba2Block):
            m.set_quant_enabled(enabled)


# ============================================================
# Smoke tests
# ============================================================

if __name__ == "__main__":
    torch.manual_seed(0)

    print("=== DCN-M2 OffsetClampSTE ===")
    offsets = torch.randn(2, 14, 16, 16, requires_grad=True) * 5.0    # outliers ~ ±15
    print(f"  in:  range [{offsets.min().item():.2f}, {offsets.max().item():.2f}]")
    clamped = offset_clamp_ste(offsets, kernel_size=3)
    print(f"  out: range [{clamped.min().item():.2f}, {clamped.max().item():.2f}]")
    clamped.sum().backward()
    grad_mean = offsets.grad.mean().item()
    grad_std = offsets.grad.std().item()
    print(f"  grad mean={grad_mean:.4f} std={grad_std:.4f} "
          f"(STE → exactly 1.0 for every element)")

    print("\n=== KScaledFakeQuant (num_bins=3) ===")
    fq = KScaledFakeQuant(num_bins=3)
    fq.start_calibration()
    with torch.no_grad():
        for _ in range(20):
            x = torch.randn(64, 32) * 1.5
            x[0, 0] = 25.0      # outlier
            x[3, 7] = -40.0
            _ = fq(x)
    fq.finalize_calibration(percentile=99.9)
    print(f"  thresholds: {fq.thresholds.tolist()}")
    print(f"  scales:     {[round(s, 5) for s in fq.scales.tolist()]}")

    x_test = torch.randn(8, 32) * 3
    x_test[0, 0] = 30.0
    x_q = fq(x_test)
    err = (x_q - x_test).abs().mean().item()
    rel_err = err / x_test.abs().mean().item()
    print(f"  avg abs err: {err:.4e}  (relative: {rel_err:.2%})")

    print("\n=== KScaledMamba2Block ===")
    block = KScaledMamba2Block(
        channels=128, d_state=32, headdim=32, num_bins=3,
        targets=("delta", "y"),
    )
    n_params = sum(p.numel() for p in block.parameters()) / 1e3
    print(f"  params: {n_params:.1f} K")

    # Calibration on dummy data
    block.start_calibration()
    with torch.no_grad():
        for _ in range(3):
            _ = block(torch.randn(1, 64, 128))
    block.finalize_calibration(percentile=99.9)

    # FP vs k-scaled-INT8 comparison
    x = torch.randn(2, 64, 128)
    block.set_quant_enabled(False)
    y_fp = block(x)
    block.set_quant_enabled(True)
    y_q = block(x)
    diff = (y_fp - y_q).abs().mean().item()
    rel = diff / y_fp.abs().mean().item()
    print(f"  FP output  range [{y_fp.min().item():.3f}, {y_fp.max().item():.3f}]")
    print(f"  INT8 output range [{y_q.min().item():.3f}, {y_q.max().item():.3f}]")
    print(f"  abs diff: {diff:.4e}  (relative: {rel:.2%})")
