"""qpack3: int4 group-wise packing for the EmbeddingGemma 2 text and vision towers.

Scheme: per group of 16 input columns, asymmetric min-max 4-bit codes (two codes per byte, low nibble
first). The per-group scale and minimum are stored as 8-bit affine integers per row with fp16 row
bounds. The audio tower is kept: its block linears are packed the same way. The vision position table is cut to the first 2,520 rows per axis,
which is exact for the default 280 soft tokens (at most 2,520 patches). Do not raise max_soft_tokens.

  build_package(model, out_dir, side_dir)   quantize a loaded EmbeddingGemma2Model in place and save it
  load_packed3(path, dtype)                 rebuild the module tree on empty weights and attach packed tensors
  load_st(path, dtype, device)              SentenceTransformer around the packed model

Requires transformers 5.19.0 and sentence-transformers >= 6.1.
"""
import contextlib
import json
import os
import re
import shutil

import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file, save_file

BITS, G, POS_ROWS = 4, 16, 2520
LEAVE = {"model.safetensors", "quant_config.json", "qpack3.py"}


def _enc_meta(t):
    """(rows, groups) fp16 -> uint8 codes plus per-row fp16 lo and hi."""
    x = t.float()
    lo = x.amin(1, keepdim=True).half()
    hi = x.amax(1, keepdim=True).half()
    span = (hi.float() - lo.float()).clamp(min=1e-6)
    q = ((x - lo.float()) / span * 255).round().clamp(0, 255).to(torch.uint8)
    return q, lo, hi


def _dec_meta(q, lo, hi, dtype):
    span = (hi.float() - lo.float()).clamp(min=1e-6)
    return (q.float() / 255 * span + lo.float()).half().to(dtype)


CLIP_GRID = (1.0, 0.97, 0.94, 0.91, 0.88, 0.85, 0.82, 0.79, 0.76)


def _clip_range(Wg, wmin, wmax, qmax):
    """Per group, pick the shrink of the min-max range that minimizes squared rounding error."""
    center, half = (wmin + wmax) / 2, (wmax - wmin) / 2
    best_err = best_lo = best_hi = None
    for f in CLIP_GRID:
        lo, hi = center - half * f, center + half * f
        sc = ((hi - lo) / qmax).half().float().clamp(min=1e-4)
        z = lo.half().float()
        q = ((Wg - z.unsqueeze(-1)) / sc.unsqueeze(-1)).round().clamp(0, qmax)
        err = ((q * sc.unsqueeze(-1) + z.unsqueeze(-1) - Wg) ** 2).sum(-1)
        if best_err is None:
            best_err, best_lo, best_hi = err, lo, hi
        else:
            better = err < best_err
            best_err = torch.where(better, err, best_err)
            best_lo, best_hi = torch.where(better, lo, best_lo), torch.where(better, hi, best_hi)
    return best_lo, best_hi


def quantize(W, g=G, bits=4, chunk=32768, clip=False):
    """W (rows, cols) -> dict of packed uint8 codes and 8-bit metadata, chunked over rows."""
    rows, cols = W.shape
    assert cols % g == 0 and cols % 2 == 0, (rows, cols)
    qmax = 2 ** bits - 1
    parts = {k: [] for k in ("packed", "sq", "slo", "shi", "zq", "zlo", "zhi")}
    for a in range(0, rows, chunk):
        Wg = W[a:a + chunk].float().reshape(-1, cols // g, g)
        wmin, wmax = Wg.amin(-1), Wg.amax(-1)
        if clip:
            wmin, wmax = _clip_range(Wg, wmin, wmax, qmax)
        scale = ((wmax - wmin) / qmax).half().clamp(min=1e-4)
        sq, slo, shi = _enc_meta(scale)
        zq, zlo, zhi = _enc_meta(wmin.half())
        s = _dec_meta(sq, slo, shi, torch.float32).unsqueeze(-1)
        z = _dec_meta(zq, zlo, zhi, torch.float32).unsqueeze(-1)
        q = ((Wg - z) / s).round().clamp(0, qmax).to(torch.uint8).reshape(-1, cols)
        parts["packed"].append((q[:, 0::2] | (q[:, 1::2] << 4)).contiguous() if bits == 4 else q.contiguous())
        for k, v in zip(("sq", "slo", "shi", "zq", "zlo", "zhi"), (sq, slo, shi, zq, zlo, zhi)):
            parts[k].append(v)
    return {k: torch.cat(v) for k, v in parts.items()}


def _empty(m, rows, cols, g, bits=4):
    mk = lambda *s, dt: torch.empty(*s, dtype=dt, device="meta")
    m.bits = bits
    m.register_buffer("packed", mk(rows, cols // 2 if bits == 4 else cols, dt=torch.uint8))
    for k in ("sq", "zq"):
        m.register_buffer(k, mk(rows, cols // g, dt=torch.uint8))
    for k in ("slo", "shi", "zlo", "zhi"):
        m.register_buffer(k, mk(rows, 1, dt=torch.float16))


def _deq(m, idx, dtype):
    pick = (lambda t: t) if idx is None else (lambda t: t[idx])
    s = _dec_meta(pick(m.sq), pick(m.slo), pick(m.shi), dtype).unsqueeze(-1)
    z = _dec_meta(pick(m.zq), pick(m.zlo), pick(m.zhi), dtype).unsqueeze(-1)
    p = pick(m.packed)
    q = torch.stack([p & 15, p >> 4], -1).reshape(p.shape[0], -1) if m.bits == 4 else p
    r = q.shape[0]
    return (q.to(dtype).reshape(r, -1, m.g) * s + z).reshape(r, -1)


class QLinear3(nn.Module):
    def __init__(self, in_features, out_features, bias, g=G, bits=4):
        super().__init__()
        self.in_features, self.out_features, self.g = in_features, out_features, g
        _empty(self, out_features, in_features, g, bits)
        if bias:
            self.register_buffer("bias", torch.empty(out_features, dtype=torch.float32, device="meta"))
        else:
            self.bias = None

    def forward(self, x):
        W = _deq(self, None, x.dtype)
        return F.linear(x, W, None if self.bias is None else self.bias.to(x.dtype))


class QEmbedding3(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, padding_idx, embed_scale, g=G, bits=4):
        super().__init__()
        self.num_embeddings, self.embedding_dim, self.g = num_embeddings, embedding_dim, g
        self.padding_idx, self.embed_scale = padding_idx, float(embed_scale)
        self.compute_dtype = torch.float32
        _empty(self, num_embeddings, embedding_dim, g, bits)

    def forward(self, ids):
        flat = ids.reshape(-1)
        W = _deq(self, flat, self.compute_dtype)
        return W.reshape(*ids.shape, self.embedding_dim) * self.embed_scale


def bits_for(path, int8_layers, audio_bits=4):
    """int8 for the layers listed per tower, audio_bits for audio, int4 elsewhere."""
    if path.startswith("audio_tower."):
        return audio_bits
    tower = path.split(".")[0]
    m = re.search(r"layers\.(\d+)\.", path)
    if m:
        return 8 if int(m.group(1)) in int8_layers.get(tower, ()) else 4
    return 8 if "other" in int8_layers.get(tower, ()) else 4


def want_linear(path, _mod):
    if path.startswith("language_model."):
        return True
    if path.startswith("vision_tower."):
        return bool(re.search(r"layers\.\d+\.", path)) and _mod.in_features % G == 0
    if path.startswith("audio_tower."):
        return _mod.in_features % G == 0 and _mod.in_features * _mod.out_features >= 65536
    return False


def want_embedding(path, mod):
    return path.startswith("language_model.") and mod.num_embeddings >= 100_000


def swap(root, make_lin, make_emb, prefix=""):
    for name, child in list(root.named_children()):
        path = f"{prefix}{name}"
        if isinstance(child, nn.Embedding) and want_embedding(path, child):
            setattr(root, name, make_emb(path, child))
        elif isinstance(child, nn.Linear) and want_linear(path, child):
            setattr(root, name, make_lin(path, child))
        else:
            swap(child, make_lin, make_emb, path + ".")


def build_package(model, out_dir, side_dir, audio_bits=4, int8_layers=None):
    int8_layers = int8_layers or {}
    """Quantize `model` (EmbeddingGemma2Model) in place, then write the package to out_dir."""
    scales = {}

    def make_lin(_p, lin):
        b = bits_for(_p, int8_layers, audio_bits)
        m = QLinear3(lin.in_features, lin.out_features, lin.bias is not None, bits=b)
        for k, v in quantize(lin.weight.data, bits=b).items():
            setattr(m, k, v)
        if lin.bias is not None:
            m.bias = lin.bias.data.float()
        return m

    def make_emb(path, emb):
        scales[path] = float(getattr(emb, "embed_scale", 1.0))
        m = QEmbedding3(emb.num_embeddings, emb.embedding_dim, emb.padding_idx, scales[path])
        for k, v in quantize(emb.weight.data).items():
            setattr(m, k, v)
        return m

    swap(model, make_lin, make_emb)
    pe = model.vision_tower.patch_embedder
    pe.position_embedding_table = nn.Parameter(pe.position_embedding_table.data[:, :POS_ROWS].clone(),
                                               requires_grad=False)
    os.makedirs(out_dir, exist_ok=True)
    sd = {k: v.detach().cpu().contiguous() for k, v in model.state_dict().items()}
    save_file(sd, os.path.join(out_dir, "model.safetensors"), metadata={"format": "pt"})
    json.dump({"format": 3, "bits": BITS, "audio_bits": audio_bits, "int8_layers": int8_layers, "group": G, "position_rows": POS_ROWS, "embed_scale": scales,
               "meta": "uint8 affine per row with fp16 lo and hi", "audio": f"tower block linears int{audio_bits}; convolutions, norms and embed_audio bf16",
               "text": "linears and token table int4", "vision": "block linears int4",
               "kept_bf16": "patch input projection, position table, embed_vision, norms"},
              open(os.path.join(out_dir, "quant_config.json"), "w"), indent=2)
    for root, _dirs, files in os.walk(side_dir):
        rel = os.path.relpath(root, side_dir)
        if rel.startswith(".") and rel != ".":
            continue
        for f in files:
            if f.endswith((".safetensors", ".bin", ".gguf", ".onnx", ".md", ".gitattributes")) or f in LEAVE:
                continue
            dst = os.path.join(out_dir, rel if rel != "." else "")
            os.makedirs(dst, exist_ok=True)
            shutil.copy(os.path.join(root, f), dst)
    shutil.copy(os.path.abspath(__file__), os.path.join(out_dir, "qpack3.py"))
    return sd


_LAZY_TOWERS = ("audio_tower.", "embed_audio.")
_LAZY = (".packed", ".sq", ".zq", ".slo", ".shi", ".zlo", ".zhi")


def _read_lazy(wfile, dtype):
    """Memory-map the audio tower so text and image workloads never make it resident; read everything
    else normally (memory-mapping the hot towers slowed the forward pass in a test)."""
    import struct
    import warnings

    import numpy as np
    from safetensors import safe_open

    with open(wfile, "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        hdr = json.loads(f.read(n))
    base = 8 + n
    dts = {"U8": np.uint8, "F16": np.float16}
    sd = {}
    with safe_open(wfile, framework="pt") as f:
        for k in f.keys():
            h = hdr[k]
            if k.startswith(_LAZY_TOWERS) and k.endswith(_LAZY) and h["dtype"] in dts:
                s0, _ = h["data_offsets"]
                arr = np.memmap(wfile, dtype=dts[h["dtype"]], mode="r", offset=base + s0, shape=tuple(h["shape"]))
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    sd[k] = torch.from_numpy(arr)
            else:
                v = f.get_tensor(k)
                sd[k] = v.to(dtype) if (dtype != torch.bfloat16 and v.dtype == torch.bfloat16) else v
    return sd


def load_packed3(path, dtype=torch.bfloat16):
    from accelerate import init_empty_weights
    from transformers import AutoConfig
    from transformers.models.embedding_gemma2.modeling_embedding_gemma2 import EmbeddingGemma2Model

    qc = json.load(open(os.path.join(path, "quant_config.json")))
    cfg = AutoConfig.from_pretrained(path)
    with init_empty_weights(include_buffers=False):
        model = EmbeddingGemma2Model(cfg)
    swap(model,
         lambda _p, l: QLinear3(l.in_features, l.out_features, l.bias is not None,
                                bits=bits_for(_p, {k: set(v) for k, v in qc.get("int8_layers", {}).items()}, qc.get("audio_bits", 4))),
         lambda p, e: QEmbedding3(e.num_embeddings, e.embedding_dim, e.padding_idx, qc["embed_scale"][p]))
    pe = model.vision_tower.patch_embedder
    pe.position_embedding_table = nn.Parameter(
        torch.empty(2, qc["position_rows"], pe.position_embedding_table.shape[-1], device="meta"),
        requires_grad=False)
    sd = _read_lazy(os.path.join(path, "model.safetensors"), dtype)
    missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
    assert not unexpected, unexpected[:6]
    left = [n for n, t in list(model.named_parameters()) + list(model.named_buffers()) if t.is_meta]
    assert not left, f"still on meta: {left[:6]} missing={missing[:6]}"
    for m in model.modules():
        if isinstance(m, QEmbedding3):
            m.compute_dtype = dtype
    return model.eval()


@contextlib.contextmanager
def _return_model(model):
    """Make from_pretrained hand back an already built model, so SentenceTransformer reads no weights."""
    import transformers.modeling_utils as mu
    from transformers.models.auto import auto_factory as af
    targets = [mu.PreTrainedModel, af._BaseAutoModelClass]
    saved = [(t, vars(t)["from_pretrained"]) for t in targets]
    try:
        for t in targets:
            t.from_pretrained = classmethod(lambda cls, *a, **k: model)
        yield
    finally:
        for t, orig in saved:
            t.from_pretrained = orig


def load_st(path, dtype=torch.bfloat16, device="cpu"):
    from sentence_transformers import SentenceTransformer
    model = load_packed3(path, dtype)
    with _return_model(model):
        st = SentenceTransformer(path, device=device)
    return st