"""Building blocks for the 8-bit SmolVLM WebGPU recipe (smolvlm-500m-instruct-q8.py).

The one 8-bit linear that public ExecuTorch lowers for WebGPU and the public WebGPU runtime runs is
et_vk.linear_q8ta_q8csw: int8 activations with one static per-tensor scale, times int8 weights with
one scale per output channel. It is emitted for linear(dq(q(x, s, z)), dq_per_channel(W8, ws, 0)),
which Q8Linear below writes out directly with quantized_decomposed ops (no PT2E prepare/convert:
the scales come from torchao MinMaxObservers run over calibration data on the fp32 model).

Two things make it accurate enough for this model (see VLM_FOLLOWUP.md for the measurements):

- Two passes. y = W8.dq(q(x, s1)) + W8.dq(q(r, s2)), r = x - dq(q(x, s1)): the second pass quantizes
  the first pass's rounding error with its own static scale (s2 ~ s1/255) and reuses the same int8
  weight, so the activation error drops from s1/2 to about s1/510. The residual is written with
  plain ops (round/clamp/mul): the lowering would fold an explicit dq(q(x)) back to x for a
  non-quantized consumer, which makes the residual zero.
- A built-in attention sink. The text decoder's position 0 is always <|im_start|>, whose hidden
  state carries massive activations (|x| ~ 3000 at the input of layer 3's down projection, vs ~4 at
  every other position). Its K/V are the same in every conversation, so they are computed once in
  fp32 and written into the cache by the program itself whenever a call starts at position 0
  (SinkKVCache); the calibration excludes that row, so it never sets a scale.
"""

import math

import torch
import torch.nn.functional as F

qd = torch.ops.quantized_decomposed
QMIN, QMAX = -128, 127


# ----------------------------------------------------------------------------- quantization


def weight_int8(w):
    """VulkanQuantizer's weight spec: int8 [-128, 127], per-output-channel symmetric, scale from
    torchao PerChannelMinMaxObserver (amax / 127.5). Returns (int8 [N, K], fp32 scales [N])."""
    from torchao.quantization.pt2e import PerChannelMinMaxObserver

    obs = PerChannelMinMaxObserver(
        ch_axis=0,
        dtype=torch.int8,
        quant_min=QMIN,
        quant_max=QMAX,
        qscheme=torch.per_channel_symmetric,
    )
    obs(w)
    scale, zp = obs.calculate_qparams()
    assert int(zp.abs().max()) == 0
    scale = scale.float()
    q = torch.clamp(torch.round(w / scale[:, None]), QMIN, QMAX).to(torch.int8)
    return q, scale


def act_observer():
    """Static per-tensor affine int8 activation observer (min/max of everything it saw)."""
    from torchao.quantization.pt2e import MinMaxObserver

    return MinMaxObserver(
        dtype=torch.int8,
        quant_min=QMIN,
        quant_max=QMAX,
        qscheme=torch.per_tensor_affine,
    )


def qparams(obs):
    s, z = obs.calculate_qparams()
    return float(s), int(z)


def fake_quant(x, s, z):
    """dq(q(x)) exactly as quantized_decomposed computes it (round(x * (1/s)) + z, clamp)."""
    return (torch.clamp(torch.round(x * (1.0 / s)) + z, QMIN, QMAX) - z) * s


class Q8Linear(torch.nn.Module):
    """et_vk.linear_q8ta_q8csw, one or two passes (see the module docstring)."""

    def __init__(self, lin, act, second=None):
        super().__init__()
        q, scale = weight_int8(lin.weight.detach().float())
        self.register_buffer("weight_int8", q)
        self.register_buffer("weight_scale", scale)
        self.register_buffer("weight_zero", torch.zeros(q.shape[0], dtype=torch.int64))
        self.bias = (
            None
            if lin.bias is None
            else torch.nn.Parameter(lin.bias.detach().float(), requires_grad=False)
        )
        self.act = (float(act[0]), int(act[1]))
        self.second = None if second is None else (float(second[0]), int(second[1]))

    def forward(self, x):
        s1, z1 = self.act
        w = qd.dequantize_per_channel(
            self.weight_int8,
            self.weight_scale,
            self.weight_zero,
            0,
            QMIN,
            QMAX,
            torch.int8,
        )
        x1 = qd.dequantize_per_tensor(
            qd.quantize_per_tensor(x, s1, z1, QMIN, QMAX, torch.int8),
            s1,
            z1,
            QMIN,
            QMAX,
            torch.int8,
        )
        y = F.linear(x1, w, self.bias)
        if self.second is not None:
            s2, z2 = self.second
            r = x - fake_quant(x, s1, z1)
            r1 = qd.dequantize_per_tensor(
                qd.quantize_per_tensor(r, s2, z2, QMIN, QMAX, torch.int8),
                s2,
                z2,
                QMIN,
                QMAX,
                torch.int8,
            )
            y = y + F.linear(r1, w)
        return y


def linear_names(module):
    return [n for n, m in module.named_modules() if isinstance(m, torch.nn.Linear)]


def set_module(root, name, new):
    parent, _, leaf = name.rpartition(".")
    setattr(root.get_submodule(parent) if parent else root, leaf, new)


class Recorder:
    """Forward pre-hooks feeding one observer per linear. `skip_first_row` (set by the feeder while
    a call starts at position 0) drops row 0, the attention-sink position, from what is observed."""

    def __init__(self, module, names, transform=None):
        self.obs = {n: act_observer() for n in names}
        self.skip_first_row = False
        self.hooks = []
        for n in names:

            def hook(m, args, n=n):
                x = args[0].detach().float()
                if self.skip_first_row and x.dim() == 3:
                    x = x[:, 1:]
                x = x.reshape(-1, x.shape[-1])
                if x.numel():
                    self.obs[n](transform(n, x) if transform else x)

            self.hooks.append(module.get_submodule(n).register_forward_pre_hook(hook))

    def close(self):
        for h in self.hooks:
            h.remove()
        return {n: qparams(o) for n, o in self.obs.items()}


def calibrate(module, names, feed, two_pass):
    """Static activation qparams for every linear in `names`: {name: (first, second or None)}.
    feed(recorder) runs the fp32 module over the calibration data (twice for two passes: the
    second run observes each linear's first-pass rounding error)."""
    rec = Recorder(module, names)
    with torch.no_grad():
        feed(rec)
    first = rec.close()
    second = {}
    if two_pass:
        rec = Recorder(
            module, names, transform=lambda n, x: x - fake_quant(x, *first[n])
        )
        with torch.no_grad():
            feed(rec)
        second = rec.close()
    return {n: (first[n], second.get(n)) for n in names}


def quantize(module, act):
    """Replace every linear named in `act` by Q8Linear (in place)."""
    for n, (first, second) in act.items():
        set_module(module, n, Q8Linear(module.get_submodule(n), first, second))
    return module


_FAKE_FIXED = []


def fix_q8ta_fake():
    """executorch 1.5.1's reference kernel for et_vk.linear_q8ta_q8csw computes linear(x_int8, w_fp32),
    so under FakeTensorMode (the Vulkan FusePatternsPass re-trace) its output comes out int8. Every
    consumer then inherits int8 metadata: llama::custom_sdpa's meta check fails on the query, and
    the serialized delegate would declare fp32 activations as int8 (the WebGPU kernel rejects a
    non-fp32 output). A fake kernel with the real output dtype (fp32, [..., N]) fixes the metadata;
    the serialized ops and constants are otherwise unchanged."""
    if _FAKE_FIXED:
        return
    import executorch.backends.vulkan.custom_ops_lib  # noqa: F401 (defines et_vk ops)

    @torch.library.register_fake("et_vk::linear_q8ta_q8csw")
    def _fake(x, input_scale, input_zero_point, weights, weight_sums, weight_scales, bias=None):
        return x.new_empty((*x.shape[:-1], weights.shape[0]), dtype=torch.float32)

    _FAKE_FIXED.append(True)


# ----------------------------------------------------------------------------- attention sink


class SinkKVCache(torch.nn.Module):
    """executorch's CustomKVCache (llama::update_cache into a [1, C, H, D] cache) with the K/V of
    position 0 built in: when a call starts at position 0, its row 0 is replaced by the constants
    before the cache update. sink = relu(1 - input_pos) is 1 at position 0 and 0 everywhere else,
    so later calls are untouched (x + 0 * (c - x) = x)."""

    def __init__(self, base, k0, v0, rows):
        super().__init__()
        self.max_context_length = base.max_context_length
        self.max_batch_size, self.n_heads, self.head_dim = (
            base.max_batch_size,
            base.n_heads,
            base.head_dim,
        )
        self.register_buffer("k_cache", base.k_cache)
        self.register_buffer("v_cache", base.v_cache)
        # [B, S, H, D], the cache layout (update() receives [B, H, S, D] and transposes back first,
        # so the arithmetic below sees one dim order: the portable kernels require that).
        self.register_buffer("k0", k0.detach().float().contiguous())
        self.register_buffer("v0", v0.detach().float().contiguous())
        first_row = torch.zeros(1, rows, 1, 1)
        first_row[0, 0, 0, 0] = 1.0
        self.register_buffer("first_row", first_row)

    def update(self, input_pos, k_val, v_val, indices=None):
        assert indices is None
        k_val = k_val.transpose(1, 2)
        v_val = v_val.transpose(1, 2)
        sink = torch.relu(1.0 - input_pos[:1].to(torch.float32)).reshape(1, 1, 1, 1) * self.first_row
        k_val = k_val + sink * (self.k0 - k_val)
        v_val = v_val + sink * (self.v0 - v_val)
        start_pos = input_pos[0].item()
        _ = torch.ops.llama.update_cache(k_val, self.k_cache, start_pos)
        _ = torch.ops.llama.update_cache(v_val, self.v_cache, start_pos)
        return self.k_cache.transpose(1, 2), self.v_cache.transpose(1, 2)


def sink_kv(transformer, row0):
    """Per-layer K/V of position 0 for the fp32 transformer, given its input row [hidden]."""
    for name, buf in transformer.named_buffers():
        if name.endswith("k_cache") or name.endswith("v_cache"):
            buf.zero_()
    with torch.no_grad():
        transformer(None, {"input_pos": torch.tensor([0])}, row0[None, None])
    out = []
    for layer in transformer.layers:
        kv = layer.attention.kv_cache
        out.append((kv.k_cache[:, :1].clone(), kv.v_cache[:, :1].clone()))
        kv.k_cache.zero_()
        kv.v_cache.zero_()
    return out


def install_sink(transformer, kvs, rows):
    for layer, (k0, v0) in zip(transformer.layers, kvs):
        layer.attention.kv_cache = SinkKVCache(layer.attention.kv_cache, k0, v0, rows)
    return transformer


class RefSDPA(torch.nn.Module):
    """Causal attention over the KV cache in plain PyTorch, used only while collecting calibration
    statistics. llama::custom_sdpa's CPU kernel divides its work by the threads it gets, so its float
    sums (and with them a few activation ranges, in the last bits) vary from run to run; this keeps
    the calibration, and the exported files, byte-reproducible. Same inputs/outputs as SDPACustom."""

    def __init__(self, dim, n_rep):
        super().__init__()
        self.dim, self.n_rep = dim, n_rep

    def forward(self, input_pos, q, k, v, bsz, seqlen, mask):
        # q [B, Hq, S, D]; k, v [B, Hkv, L, D] (the whole cache)
        end = int(input_pos[0]) + seqlen
        k = k[:, :, :end].repeat_interleave(self.n_rep, dim=1)
        v = v[:, :, :end].repeat_interleave(self.n_rep, dim=1)
        scores = (q @ k.transpose(-1, -2)) / math.sqrt(q.shape[-1])
        rows = torch.arange(end - seqlen, end)[:, None]
        scores = scores.masked_fill(torch.arange(end)[None, :] > rows, float("-inf"))
        out = torch.softmax(scores, dim=-1) @ v
        return out.transpose(1, 2).reshape(bsz, seqlen, self.dim)


class reference_attention:
    """Context manager: every layer's SDPA module replaced by RefSDPA, restored on exit."""

    def __init__(self, transformer):
        self.t = transformer

    def __enter__(self):
        self.saved = []
        for layer in self.t.layers:
            a = layer.attention
            self.saved.append(a.SDPA)
            a.SDPA = RefSDPA(a.SDPA.dim, a.n_heads // a.n_kv_heads)
        return self.t

    def __exit__(self, *exc):
        for layer, sdpa in zip(self.t.layers, self.saved):
            layer.attention.SDPA = sdpa


# ----------------------------------------------------------------------------- windows


def window_starts(start, end, rows):
    """First positions of the `rows`-row windows that feed positions [start, end): each window ends
    at most at `end`, the last exactly there (the same plan as windowStarts() in vlm.js)."""
    if start >= end:
        return []
    assert end >= rows, (start, end, rows)
    out, stop = [], min(start + rows, end)
    while True:
        out.append(stop - rows)
        if stop == end:
            return out
        stop = min(stop + rows, end)


# ----------------------------------------------------------------------------- inspection


def q8_ops_vs_eager(pte_bytes, module):
    """Every et_vk.linear_q8ta_q8csw in the delegates: its int8 weight and weight scales must be
    bit-identical to one Q8Linear's buffers, and its activation scale / zero point one of that
    linear's static qparams (scale to 1e-7 relative: the flatc JSON dump prints about 8 significant
    digits). Returns counts."""
    import hashlib

    import _smolvlm as S
    import numpy as np
    from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
        VulkanDelegateHeader,
    )
    from executorch.exir._serialize._program import deserialize_pte_binary

    eager = {}
    for m in module.modules():
        if isinstance(m, Q8Linear):
            key = hashlib.sha256(m.weight_int8.numpy().tobytes()).hexdigest()
            eager[key] = (
                m.weight_scale.numpy(),
                {m.act} | ({m.second} if m.second else set()),
            )
    pte = deserialize_pte_binary(pte_bytes)
    prog = pte.program if hasattr(pte, "program") else pte
    nd = getattr(pte, "named_data", None)
    seg_base = int.from_bytes(pte_bytes[24:32], "little")
    total = weights_ok = scales_ok = act_ok = 0
    for plan in prog.execution_plan:
        for d in plan.delegates:
            proc = d.processed
            if proc.location.name == "INLINE":
                blob = bytes(prog.backend_delegate_data[proc.index].data)
            else:
                seg = prog.segments[proc.index]
                blob = pte_bytes[
                    seg_base + seg.offset : seg_base + seg.offset + seg.size
                ]
            h = VulkanDelegateHeader.from_bytes(
                blob[: VulkanDelegateHeader.EXPECTED_LENGTH]
            )
            raw = blob[h.bytes_offset : h.bytes_offset + h.bytes_size]
            g = S._vk_graph(blob)

            def const(vid):
                t = g.values[vid].value
                c = g.constants[t.constant_id]
                if c.named_key:
                    return bytes(
                        nd.buffers[nd.pte_data[c.named_key].buffer_index]
                    ), list(t.dims)
                return raw[c.offset : c.offset + c.length], list(t.dims)

            for op in g.chain:
                if op.name != "et_vk.linear_q8ta_q8csw.default":
                    continue
                total += 1
                wb, (n, k) = const(op.args[3])
                hit = eager.get(hashlib.sha256(wb[: n * k]).hexdigest())
                if hit is None:
                    continue
                weights_ok += 1
                sb, _ = const(op.args[5])
                scales = np.frombuffer(sb, dtype=np.float32)[:n]
                scales_ok += int(np.array_equal(scales, hit[0]))
                s = g.values[op.args[1]].value.double_val
                z = g.values[op.args[2]].value.int_val
                act_ok += int(
                    any(abs(s - a) <= 1e-7 * a and z == b for a, b in hit[1])
                )
    return {
        "linear_q8ta_q8csw_ops": total,
        "weight_int8_bit_identical": weights_ok,
        "weight_scales_bit_identical": scales_ok,
        "activation_qparams_match": act_ok,
        "eager_q8_linears": len(eager),
    }


# ----------------------------------------------------------------------------- GPTQ (optional)


def gptq_int8(w, H, damp=0.01, block=128):
    """GPTQ (Frantar et al. 2022) for the same int8 grid as weight_int8 (per-output-channel symmetric,
    scale = amax / 127.5 of the original weight): columns are rounded in order and each rounding
    error is pushed onto the not-yet-rounded columns through the inverse Hessian H = sum x x^T of the
    layer's calibration inputs. Returns (int8 [N, K], fp32 scales [N]); float64 throughout."""
    _, scale = weight_int8(w)
    s = scale.double()[:, None]
    w = w.detach().double().clone()
    H = H.detach().double().clone()
    K = w.shape[1]
    dead = torch.diag(H) == 0
    H[dead, dead] = 1
    w[:, dead] = 0
    H += damp * torch.mean(torch.diag(H)) * torch.eye(K, dtype=H.dtype)
    hinv = torch.linalg.cholesky(torch.cholesky_inverse(torch.linalg.cholesky(H)), upper=True)
    q = torch.zeros_like(w)
    for i1 in range(0, K, block):
        i2 = min(i1 + block, K)
        w1 = w[:, i1:i2].clone()
        err1 = torch.zeros_like(w1)
        h1 = hinv[i1:i2, i1:i2]
        for i in range(i2 - i1):
            col = w1[:, i : i + 1]
            qc = torch.clamp(torch.round(col / s), QMIN, QMAX)
            q[:, i1 + i : i1 + i + 1] = qc
            e = (col - qc * s) / h1[i, i]
            w1[:, i:] -= e * h1[i : i + 1, i:]
            err1[:, i : i + 1] = e
        w[:, i2:] -= err1 @ hinv[i1:i2, i2:]
    return q.to(torch.int8), scale


def hessians(module, names, feed):
    """sum x x^T (float64) of every named linear's input over feed(recorder-like ctx)."""
    hs, hooks = {}, []

    class Ctx:
        skip_first_row = False

    ctx = Ctx()
    for n in names:

        def hook(m, args, n=n):
            x = args[0].detach().double()
            if ctx.skip_first_row and x.dim() == 3:
                x = x[:, 1:]
            x = x.reshape(-1, x.shape[-1])
            hs[n] = hs[n] + x.T @ x if n in hs else x.T @ x
            rows[n] = rows.get(n, 0) + x.shape[0]

        hooks.append(module.get_submodule(n).register_forward_pre_hook(hook))
    rows = {}
    with torch.no_grad():
        feed(ctx)
    for h in hooks:
        h.remove()
    hessians.rows = rows
    return hs


# ----------------------------------------------------------------------------- model builders


def quantize_with(module, act, weights=None):
    """Like quantize(), with optional precomputed (int8, scale) weights per linear (GPTQ)."""
    for n, (first, second) in act.items():
        lin = module.get_submodule(n)
        q = Q8Linear(lin, first, second)
        if weights and n in weights:
            q.weight_int8.copy_(weights[n][0])
            q.weight_scale.copy_(weights[n][1])
        set_module(module, n, q)
    return module


def build_vision(vision_fp32, images, *, two_pass, gptq):
    """Quantize a (deep-copied) fp32 VisionEncoder: every linear becomes Q8Linear, scales from the
    calibration images (pixel tensors [1,3,S,S])."""
    import copy

    v = copy.deepcopy(vision_fp32).eval()
    names = linear_names(v)

    def feed(rec):
        for px in images:
            v(px)

    act = calibrate(v, names, feed, two_pass)
    weights = None
    if gptq:
        # GPTQ only where the calibration rows outnumber the input features: the connector
        # (12288 inputs, 64 rows per image) keeps round-to-nearest.
        hs = hessians(v, names, feed)
        weights = {n: gptq_int8(v.get_submodule(n).weight, hs[n]) for n in names if hessians.rows[n] >= hs[n].shape[0]}
        del hs
    return quantize_with(v, act, weights), act


def build_text(transformer, calib_rows, *, rows, two_pass, gptq, sink_row=None):
    """Quantize the fp32 executorch llama Transformer in place.

    calib_rows: list of [n, hidden] embedding sequences (prompt + continuation), fed in the same
    `rows`-row windows the browser uses; row 0 (the sink) is excluded from every statistic.
    sink_row: the input row of position 0 (the <|im_start|> embedding as the host feeds it), or None
    for no built-in sink. Returns (transformer, activation qparams, sink K/V)."""
    import _smolvlm as S

    names = linear_names(transformer)
    sink = sink_row is not None
    with reference_attention(transformer):
        kvs = sink_kv(transformer, sink_row) if sink else None

    def feed(rec):
        for seq in calib_rows:
            S.reset_kv(transformer)
            for start in window_starts(0, seq.shape[0], rows):
                rec.skip_first_row = sink and start == 0
                transformer(None, {"input_pos": torch.tensor([start])}, seq[None, start : start + rows])
            rec.skip_first_row = False
        S.reset_kv(transformer)

    with reference_attention(transformer):
        act = calibrate(transformer, names, feed, two_pass)
        hs = hessians(transformer, names, feed) if gptq else None
    weights = None
    if gptq:
        weights = {n: gptq_int8(transformer.get_submodule(n).weight, hs.pop(n)) for n in names}
    quantize_with(transformer, act, weights)
    if sink:
        install_sink(transformer, kvs, rows)
    return transformer, act, kvs