"""Native post-training int8 runtime for the packed LSCodec / FocalCodec model repos.

Packed format (what `export_lscodec_int8.py` / `export_quantized.py` write): every weight as
int8 levels under its original state-dict key plus an fp32 scale under `_scale/<key>`.

`apply_int8(model, packed)` swaps each Linear / Conv1d / ConvTranspose1d / Embedding for a twin
that keeps the weight int8-resident and computes int8 x int8 -> int32 GEMMs (`torch._int_mm`)
with per-tensor dynamic activation quantization (W8A8 dynamic). Biases, norms and other
activation-side constants are materialized as fp32 at load (KB-scale); no weight matrix is ever
dequantized to run it. fp4/fp8/fp16 have no int8 GEMM path — fake-quantize those instead.
"""

from __future__ import annotations

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

SCALE_PREFIX = "_scale/"
_ROUND = 32  # cublasLt int8 GEMM (via torch._int_mm) wants %32 dims on CUDA; %16 is not enough (K=16 fails)


def _up(n: int) -> int:
    return (n + _ROUND - 1) // _ROUND * _ROUND


def _act_quant(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Per-row (last-dim) symmetric absmax activation quantization -> (int8 levels, scales).

    One scale per GEMM row (token / output position): a single per-tensor absmax lets a few
    outlier channels set the step for everything, which LSCodec's GroupNorms amplify badly.
    """
    scale = (x.abs().amax(dim=-1, keepdim=True) / 127.0).clamp_min(1e-12)
    return (x / scale).round().clamp(-128, 127).to(torch.int8), scale


def _int_mm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """(M, K) int8 x (K, N) int8 -> (M, N) int32. M is padded: CUDA rejects M <= 16."""
    m = a.shape[0]
    mp = max(_ROUND, _up(m))
    if mp != m:
        padded = torch.zeros((mp, a.shape[1]), dtype=a.dtype, device=a.device)
        padded[:m] = a
        a = padded
    try:
        return torch._int_mm(a, b)[:m]
    except (RuntimeError, NotImplementedError):  # backend refuses the shape
        if a.is_cuda:
            # exact in float32 (|int8 * int8 * K| << 2^24); int32 matmul has no CUDA kernel
            return (a.to(torch.float32) @ b.to(torch.float32))[:m]
        return (a.to(torch.int32) @ b.to(torch.int32))[:m]


def _k_pad(a: torch.Tensor, k: int) -> torch.Tensor:
    """Zero-pad the K dim; the weight's K is padded the same way, so the product is exact."""
    if a.shape[-1] == k:
        return a
    padded = torch.zeros((*a.shape[:-1], k), dtype=a.dtype, device=a.device)
    padded[..., : a.shape[-1]] = a
    return padded


def _quant_channel(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """fp32 -> int8 + per-output-channel absmax scale (same grid as core.int_levels)."""
    scale = w.abs().amax(dim=tuple(range(1, w.dim())), keepdim=True) if w.dim() > 1 else w.abs().amax()
    scale = (scale / 128.0).clamp_min(1e-12)
    return (w / scale).round().clamp(-128, 127).to(torch.int8), scale.float()


def _quant_tensor(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """fp32 -> int8 + one per-tensor scale (used where per-channel scales don't fit the GEMM)."""
    scale = (w.abs().amax() / 128.0).clamp_min(1e-12)
    return (w / scale).round().clamp(-128, 127).to(torch.int8), scale.float()


def _dequant(packed: dict, key: str, device: torch.device) -> torch.Tensor:
    value = packed[key].to(device)
    scale = packed.get(SCALE_PREFIX + key)
    return value if scale is None else value.float() * scale.to(device)


def _int1(v) -> int:
    if isinstance(v, int):
        return v
    assert len(v) == 1, f"unsupported tuple attribute {v}"
    return v[0]


def _pair(v) -> tuple[int, int]:
    if isinstance(v, int):
        return (v, v)
    if len(v) == 1:
        return (v[0], v[0])
    assert len(v) == 2, f"unsupported padding {v}"
    return (v[0], v[1])


class QuantLinear(nn.Module):
    def __init__(self, mod: nn.Linear, w: torch.Tensor, scale: torch.Tensor, bias: torch.Tensor | None):
        super().__init__()
        n, k = w.shape
        wt = torch.zeros((_up(n), _up(k)), dtype=torch.int8, device=w.device)
        wt[:n, :k] = w
        self.register_buffer("wT", wt.t().contiguous())  # (Kp, Np), the _int_mm operand
        sc = torch.zeros((_up(n),), dtype=torch.float32, device=w.device)
        sc[:n] = scale.reshape(-1)
        self.register_buffer("w_scale", sc)
        self.register_buffer("bias", bias)
        self.k, self.n = k, n

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        assert x.shape[-1] == self.k, f"expected in_features {self.k}, got {tuple(x.shape)}"
        q, sx = _act_quant(x)
        a = _k_pad(q.reshape(-1, self.k), self.wT.shape[0])
        y = _int_mm(a, self.wT)[:, : self.n].float() * (sx.reshape(-1, 1) * self.w_scale[: self.n])
        y = y.view(*x.shape[:-1], self.n)
        return y if self.bias is None else y + self.bias


class QuantConv1d(nn.Module):
    def __init__(self, mod: nn.Conv1d, w: torch.Tensor, scale: torch.Tensor, bias: torch.Tensor | None):
        super().__init__()
        self.register_buffer("bias", bias)  # before the early returns below
        self.k, self.s, self.d = _int1(mod.kernel_size), _int1(mod.stride), _int1(mod.dilation)
        if isinstance(mod.padding, str):  # FocalCodec's "same" convs: stride 1, left-heavy pad
            assert mod.padding in {"same", "valid"} and self.s == 1, f"unsupported padding {mod.padding!r}"
            left = 0 if mod.padding == "valid" else (self.k - 1) // 2
            self.pl, self.pr = (left, (self.k - 1) - left) if mod.padding == "same" else (0, 0)
        else:
            self.pl, self.pr = _pair(mod.padding)
        assert mod.groups == 1 or mod.groups == mod.in_channels or mod.in_channels % mod.groups == 0
        self.kind = "dense"
        if mod.groups > 1:
            self.kind = "depthwise" if mod.groups == mod.in_channels and mod.out_channels == mod.groups else "grouped"
        self.groups = mod.groups
        self.ci = mod.in_channels
        w = w.reshape(w.shape[0], -1)  # (co, (ci / groups) * k); accepts a 3-D conv weight
        n, ck = w.shape
        self.n, self.ck = n, ck
        sc = scale.reshape(-1)
        if self.kind == "depthwise":
            self.register_buffer("w_d", w.reshape(n, -1))  # (C, k), elementwise int MAC below
            self.register_buffer("w_scale_d", sc)
            return
        if self.kind == "grouped":
            co_g, ck_g = n // mod.groups, ck  # ck is already (ci / groups) * k, group-major
            wts, scales = [], []
            for gi in range(mod.groups):
                wt = torch.zeros((_up(co_g), _up(ck_g)), dtype=torch.int8, device=w.device)
                wt[:co_g, :ck_g] = w[gi * co_g : (gi + 1) * co_g]
                wts.append(wt.t().contiguous())
                s = torch.zeros((_up(co_g),), dtype=torch.float32, device=w.device)
                s[:co_g] = sc[gi * co_g : (gi + 1) * co_g]
                scales.append(s)
            self.wT, self.w_scale_g = wts, scales
            self.co_g, self.ck_g = co_g, ck_g
            return
        wt = torch.zeros((_up(n), _up(ck)), dtype=torch.int8, device=w.device)
        wt[:n, :ck] = w
        self.register_buffer("wT", wt.t().contiguous())
        s = torch.zeros((_up(n),), dtype=torch.float32, device=w.device)
        s[:n] = sc
        self.register_buffer("w_scale", s)

    def _cols(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:
        """fp32 input -> quantized im2col windows (B, Lout, ck) + per-position scales."""
        x = F.pad(x, (self.pl, self.pr)) if (self.pl or self.pr) else x
        cols = F.unfold(x.unsqueeze(2), (1, self.k), dilation=(1, self.d), stride=(1, self.s))
        q, sx = _act_quant(cols.permute(0, 2, 1))  # rows = output positions
        return q, sx, cols.shape[-1]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        squeeze = x.dim() == 2
        if squeeze:
            x = x.unsqueeze(0)
        assert x.shape[1] == self.ci, f"expected {self.ci} input channels, got {x.shape[1]}"
        b = x.shape[0]
        q, sx, lout = self._cols(x)
        sx_rows = sx.reshape(-1, 1)
        if self.kind == "depthwise":
            c = self.w_d.shape[0]
            acc = q.permute(0, 2, 1).view(b, c, self.k, lout).to(torch.int32) * self.w_d.to(torch.int32).view(c, self.k, 1)
            y = acc.sum(2).float() * sx.transpose(1, 2) * self.w_scale_d.view(1, -1, 1)  # (B, C, Lout)
            y = y.permute(0, 2, 1)
        elif self.kind == "grouped":
            qg = q.view(b, lout, self.groups, self.ck_g)
            outs = []
            for gi in range(self.groups):
                a = _k_pad(qg[:, :, gi].reshape(-1, self.ck_g), self.wT[gi].shape[0])
                yi = _int_mm(a, self.wT[gi])[:, : self.co_g].float() * (sx_rows * self.w_scale_g[gi][: self.co_g])
                outs.append(yi.view(b, lout, self.co_g))
            y = torch.cat(outs, dim=-1)
        else:
            a = _k_pad(q.reshape(-1, self.ck), self.wT.shape[0])
            y = _int_mm(a, self.wT)[:, : self.n].float() * (sx_rows * self.w_scale[: self.n])
            y = y.view(b, lout, self.n)
        if self.bias is not None:
            y = y + self.bias
        y = y.transpose(1, 2)
        return y.squeeze(0) if squeeze else y


class QuantConvTranspose1d(nn.Module):
    def __init__(self, mod: nn.ConvTranspose1d, w_eff: torch.Tensor, bias: torch.Tensor | None):
        super().__init__()
        assert _int1(mod.groups) == 1, "grouped ConvTranspose1d not supported"
        self.k, self.s, self.d = _int1(mod.kernel_size), _int1(mod.stride), _int1(mod.dilation)
        self.pl, self.pr = _pair(mod.padding)
        self.op = _int1(mod.output_padding)
        ci, co = w_eff.shape[0], w_eff.shape[1]
        self.ci, self.n = ci, co
        w, sc = _quant_tensor(w_eff)  # per-K scales in the pack don't fit a GEMM: per-tensor here
        wts = []
        for ki in range(self.k):
            wt = torch.zeros((_up(ci), _up(co)), dtype=torch.int8, device=w.device)
            wt[:ci, :co] = w[:, :, ki]
            wts.append(wt)
        self.wT, self.scale = wts, sc
        self.register_buffer("bias", bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        assert x.dim() == 3 and x.shape[1] == self.ci, f"expected (B, {self.ci}, L), got {tuple(x.shape)}"
        b, _, l = x.shape
        lout = (l - 1) * self.s - (self.pl + self.pr) + self.d * (self.k - 1) + 1 + self.op
        q, sx = _act_quant(x.permute(0, 2, 1))  # rows = (batch, time) positions
        a = _k_pad(q.reshape(-1, self.ci), self.wT[0].shape[0])
        acc = torch.zeros((b, lout, self.n), dtype=torch.float32, device=x.device)
        rows = torch.arange(l, device=x.device) * self.s
        sx_rows = sx.reshape(-1, 1)
        for ki, wt in enumerate(self.wT):
            y = _int_mm(a, wt)[:, : self.n].float() * (sx_rows * self.scale)
            t = rows + ki * self.d - self.pl
            keep = (t >= 0) & (t < lout)
            if keep.any():
                acc[:, t[keep]] += y.view(b, l, self.n)[:, keep]
        y = acc if self.bias is None else acc + self.bias
        return y.transpose(1, 2)


class QuantEmbedding(nn.Module):
    def __init__(self, w: torch.Tensor, scale: torch.Tensor, bias: torch.Tensor | None = None):
        super().__init__()
        # named `weight` so legacy hooks that touch `.weight.device` keep working; stays int8
        self.register_buffer("weight", w)
        self.register_buffer("w_scale", scale.reshape(-1))

    def forward(self, idx: torch.Tensor) -> torch.Tensor:
        return self.weight[idx].float() * self.w_scale[idx].unsqueeze(-1)


_TARGETS = (nn.Linear, nn.Conv1d, nn.ConvTranspose1d, nn.Embedding)


def _effective_weight(mod, name: str, packed: dict, device: torch.device, consumed: set) -> torch.Tensor:
    """Materialize a weight-normed layer's effective weight from the packed g/v via its own pre-hook."""
    g_key, v_key = f"{name}.weight_g", f"{name}.weight_v"
    g = _dequant(packed, g_key, device)
    v = _dequant(packed, v_key, device)
    consumed.update({g_key, v_key})
    hooks = tuple(mod._forward_pre_hooks.values())
    assert hooks, f"{name}: weight-normed but no forward pre-hook to recompute `weight`"
    with torch.no_grad():
        mod.weight_g.copy_(g)
        mod.weight_v.copy_(v)
        for hook in hooks:
            hook(mod, ())
    return mod.weight.detach().clone()


def _build_twin(mod, name: str, packed: dict, device: torch.device, consumed: set) -> nn.Module:
    bias = None
    bias_key = f"{name}.bias"
    if bias_key in packed:
        consumed.add(bias_key)
        bias = _dequant(packed, bias_key, device)
    if f"{name}.weight_g" in packed:
        w, scale, effective = _effective_weight(mod, name, packed, device, consumed), None, True
    elif f"{name}.weight" in packed:
        consumed.add(f"{name}.weight")
        raw = packed[f"{name}.weight"].to(device)
        scale = packed.get(SCALE_PREFIX + f"{name}.weight")
        scale = scale.to(device) if scale is not None else None
        w, effective = raw, raw.is_floating_point()  # int8 packed -> direct; fp32 -> requantize effective
    else:
        raise KeyError(f"{name}: no weight in packed state")
    if isinstance(mod, nn.ConvTranspose1d):
        w_eff = w if effective else w.float() * scale
        return QuantConvTranspose1d(mod, w_eff, bias)
    if effective:  # weight-normed (or caller passed fp32): quantize the effective weight now
        w, scale = _quant_channel(w)
    if isinstance(mod, nn.Linear):
        return QuantLinear(mod, w, scale, bias)
    if isinstance(mod, nn.Embedding):
        return QuantEmbedding(w, scale)
    return QuantConv1d(mod, w, scale, bias)


def apply_int8(model: nn.Module, packed: dict) -> nn.Module:
    """Swap every Linear/Conv1d/ConvTranspose1d/Embedding for its int8 twin, in place.

    `packed` is the export format: int8 levels + `_scale/` fp32 scales (fp32 tensors are
    accepted too). Everything else in `packed` (biases of unswapped modules, norms, buffers)
    is materialized as fp32 and loaded; weights of swapped modules stay int8.
    """
    device = next(model.parameters()).device
    consumed: set = set()
    twins: dict[str, nn.Module] = {}
    for name, mod in model.named_modules():
        if name and isinstance(mod, _TARGETS):
            twin = _build_twin(mod, name, packed, device, consumed)
            twin.train(mod.training)
            twins[name] = twin
    for name, twin in twins.items():
        parent_name, _, attr = name.rpartition(".")
        parent = model.get_submodule(parent_name) if parent_name else model
        parent._modules[attr] = twin
    fp_state = {}
    for key, value in packed.items():
        if key.startswith(SCALE_PREFIX) or key in consumed:
            continue
        scale = packed.get(SCALE_PREFIX + key)
        # packed int8/fp8 -> dequantize (int8 without a scale would be a bug, not a buffer);
        # unscaled non-float buffers (int64 masks etc.) pass through with their own dtype
        fp_state[key] = value.float() * scale if scale is not None else value
    _, unexpected = model.load_state_dict(fp_state, strict=False)
    assert not unexpected, f"packed keys with no home in the model: {unexpected}"
    return model


def load_packed(repo_id: str, filename: str) -> dict:
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file

    return load_file(hf_hub_download(repo_id, filename, repo_type="model"))