"""Weight-only int8 / int4 checkpoints for Auto's ModernBERT classifiers, in plain PyTorch.

Runs wherever PyTorch runs (CPU, Apple MPS, NVIDIA CUDA, AMD ROCm): weights stay quantized in memory and each layer
dequantizes its own weight just before the matrix multiply. Activations keep the device's float dtype, so no
quantized kernels are needed. Requires torch and transformers (5.x); nothing is compiled.

    from auto_quant import load, score
    model, tokenizer = load("ProCreations/auto-200m-2-int4")    # device and dtype chosen automatically
    p_deny = score(model, tokenizer, text)                      # one request, any length up to 65,536 tokens

By default attention is Auto's exact memory-linear attention (the same as the Auto runtime): no N x N score matrix or
mask is ever built, so a 65,536-token request fits in a few GB on CPU or Apple MPS. It scores one unpadded request
at a time through score(). Pass attention="sdpa" (or any Transformers attention implementation) to call
model(**inputs) on padded batches instead.

Format (quant_config.json, "format": "auto-quant-v1"): symmetric integers, one float16 scale per output channel
(int8) or per group of `group_size` consecutive input weights (int4). Dequantized weight = float32(q) *
float32(scale), then cast to the compute dtype. int4 values are stored as q + 8 in nibbles, two per byte, low first.
"""

import json
import os
import types

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

FORMAT = "auto-quant-v1"
RANGES = {8: (-127, 127), 4: (-8, 7)}


# ------------------------------------------------------------------------------------------------ quantized weights


def pack_int4(q):
    """int8 tensor in [-8, 7] with an even last dimension -> uint8 with two values per byte (low nibble first)."""
    u = (q.to(torch.int16) + 8).to(torch.uint8)
    return u[..., 0::2] | (u[..., 1::2] << 4)


def unpack_int4(packed):
    low = (packed & 15).to(torch.int8) - 8
    high = (packed >> 4).to(torch.int8) - 8
    return torch.stack([low, high], dim=-1).flatten(-2)


def dequantize(qweight, scales, bits, group_size, dtype):
    """Float weight from stored integers and scales (rows = output channels / embedding rows)."""
    q = unpack_int4(qweight) if bits == 4 else qweight
    if bits == 8:
        w = q.float() * scales.float()[..., None]
    else:
        cols = q.shape[-1]
        w = (q.float().reshape(*q.shape[:-1], cols // group_size, group_size) * scales.float()[..., None]).reshape(q.shape)
    return w.to(dtype)


def quantize(weight, bits, group_size, scales):
    """Integers for `weight` given its scales (rounded to float16): per row for int8, per row and group for int4."""
    lo, hi = RANGES[bits]
    w = weight.detach().float()
    s = scales.to(torch.float16).float()
    if bits == 8:
        q = torch.clamp(torch.round(w / s[:, None]), lo, hi)
    else:
        rows, cols = w.shape
        q = torch.clamp(torch.round(w.reshape(rows, cols // group_size, group_size) / s[..., None]), lo, hi).reshape(rows, cols)
    q = q.to(torch.int8)
    return (pack_int4(q) if bits == 4 else q), scales.to(torch.float16)


class QuantLinear(nn.Module):
    def __init__(self, in_features, out_features, bits, group_size, bias=False, device=None):
        super().__init__()
        self.in_features, self.out_features, self.bits, self.group_size = in_features, out_features, bits, group_size
        cols = in_features // 2 if bits == 4 else in_features
        groups = in_features // group_size if bits == 4 else None
        self.register_buffer("qweight", torch.empty(out_features, cols, dtype=torch.uint8 if bits == 4 else torch.int8, device=device))
        self.register_buffer("scales", torch.empty((out_features, groups) if groups else (out_features,), dtype=torch.float16, device=device))
        self.register_buffer("bias", torch.empty(out_features, dtype=torch.float16, device=device) if bias else None)

    def forward(self, x):
        w = dequantize(self.qweight, self.scales, self.bits, self.group_size, x.dtype)
        return F.linear(x, w, None if self.bias is None else self.bias.to(x.dtype))


class QuantEmbedding(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, bits, group_size, dtype=torch.float32, device=None):
        super().__init__()
        self.num_embeddings, self.embedding_dim, self.bits, self.group_size = num_embeddings, embedding_dim, bits, group_size
        self.dtype = dtype
        cols = embedding_dim // 2 if bits == 4 else embedding_dim
        groups = embedding_dim // group_size if bits == 4 else None
        self.register_buffer("qweight", torch.empty(num_embeddings, cols, dtype=torch.uint8 if bits == 4 else torch.int8, device=device))
        self.register_buffer("scales", torch.empty((num_embeddings, groups) if groups else (num_embeddings,), dtype=torch.float16, device=device))

    def forward(self, input_ids):
        # Only the looked-up rows are dequantized.
        return dequantize(self.qweight[input_ids], self.scales[input_ids], self.bits, self.group_size, self.dtype)


def targets(model):
    """Quantized modules: token embeddings, every encoder Linear and the head's dense layer. LayerNorm weights and
    the final 2-way classifier stay in float16."""
    names = ["model.embeddings.tok_embeddings"]
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and (name.startswith("model.layers.") or name == "head.dense"):
            names.append(name)
    return names


def swap(model, config, dtype=None, device=None):
    """Replace the target modules with (empty) quantized modules; then load with load_state_dict(strict=True)."""
    bits, group = config["bits"], config.get("group_size")
    dtype = dtype or next(model.parameters()).dtype
    for name in config["modules"]:
        parent_name, _, child = name.rpartition(".")
        parent = model.get_submodule(parent_name)
        old = getattr(parent, child)
        dev = device if device is not None else old.weight.device
        if isinstance(old, nn.Embedding):
            new = QuantEmbedding(old.num_embeddings, old.embedding_dim, bits, group, dtype=dtype, device=dev)
        else:
            new = QuantLinear(old.in_features, old.out_features, bits, group, bias=old.bias is not None, device=dev)
        setattr(parent, child, new)
    return model


# -------------------------------------------------------------- memory-linear exact attention (from the Auto runtime)

LOCAL_BLOCK, MLP_CHUNK = 1024, 8192
SCORE_BUDGET = {"cpu": 256 * 2**20}  # bytes of temporary scores per call if only the plain math kernel applies
GPU_SCORE_BUDGET = 1024 * 2**20
try:
    from torch.nn.attention import SDPBackend, sdpa_kernel

    _FUSED = {"cuda": [b for b in ("FLASH_ATTENTION", "EFFICIENT_ATTENTION", "CUDNN_ATTENTION") if hasattr(SDPBackend, b)],
              "cpu": ["FLASH_ATTENTION"]}
except ImportError:  # torch < 2.3
    sdpa_kernel, _FUSED = None, {}
_UNAVAILABLE = set()


def _sdpa(query, key, value, mask, scaling):
    return F.scaled_dot_product_attention(query, key, value, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scaling)


def _rows(device, heads, keys, element_size):
    return max(1, SCORE_BUDGET.get(device.type, GPU_SCORE_BUDGET) // max(1, heads * keys * max(4, element_size)))


def _global_attention(query, key, value, scaling):
    kind = query.device.type
    if sdpa_kernel is not None and kind in _FUSED and (kind, query.dtype) not in _UNAVAILABLE:
        try:  # one unmasked call restricted to memory-efficient fused kernels
            with sdpa_kernel([getattr(SDPBackend, b) for b in _FUSED[kind]]):
                return _sdpa(query, key, value, None, scaling)
        except RuntimeError:
            _UNAVAILABLE.add((kind, query.dtype))
    n = query.shape[-2]
    block = _rows(query.device, query.shape[1], n, query.element_size())
    if block >= n:
        return _sdpa(query, key, value, None, scaling)
    out = torch.empty_like(query)
    for s in range(0, n, block):
        out[:, :, s : s + block] = _sdpa(query[:, :, s : s + block], key, value, None, scaling)
    return out


def _local_attention(query, key, value, radius, scaling):
    """Each query i attends to keys j with |i - j| <= radius, one block of queries at a time."""
    n = query.shape[-2]
    block = max(1, min(LOCAL_BLOCK, _rows(query.device, query.shape[1], LOCAL_BLOCK + 2 * radius, query.element_size())))
    out, cached = torch.empty_like(query), None
    for start in range(0, n, block):
        end = min(n, start + block)
        low, high = max(0, start - radius), min(n, end + radius)
        interior = low == start - radius and high == end + radius and end - start == block
        if interior and cached is not None:
            mask = cached
        else:
            mask = (torch.arange(start, end, device=query.device)[:, None] - torch.arange(low, high, device=query.device)[None, :]).abs() <= radius
            if interior:
                cached = mask
        out[:, :, start:end] = _sdpa(query[:, :, start:end], key[:, :, low:high], value[:, :, low:high], mask, scaling)
    return out


def linear_attention(module, query, key, value, attention_mask=None, scaling=None, dropout=0.0, sliding_window=None, **kwargs):
    """Transformers attention interface (registered as "auto_linear"); one unpadded example, no mask."""
    if query.shape[0] != 1 or attention_mask is not None:
        raise ValueError('attention="linear" scores one unpadded request at a time: use score(), or load with attention="sdpa"')
    out = _global_attention(query, key, value, scaling) if sliding_window is None else _local_attention(query, key, value, sliding_window - 1, scaling)
    return out.transpose(1, 2).contiguous(), None


def _chunked_mlp(self, hidden):
    if hidden.shape[-2] <= MLP_CHUNK or torch.is_grad_enabled():
        return self._full_forward(hidden)
    # The GLU acts on each token independently, so chunking along the sequence is exact.
    return torch.cat([self._full_forward(hidden[..., s : s + MLP_CHUNK, :]) for s in range(0, hidden.shape[-2], MLP_CHUNK)], dim=-2)


# ------------------------------------------------------------------------------------------------------ public API


def _resolve(path_or_repo, revision):
    if os.path.isdir(path_or_repo):
        return path_or_repo
    from huggingface_hub import snapshot_download

    return snapshot_download(path_or_repo, revision=revision)


def load(path_or_repo, device="auto", dtype=None, revision=None, attention="linear", **model_kwargs):
    """Load a quantized Auto checkpoint -> (model, tokenizer), in eval mode on `device` ("auto": CUDA/ROCm, then
    Apple MPS, then CPU). dtype defaults to bfloat16 on GPUs that support it and float32 elsewhere. No full-precision
    copy of the weights is ever materialized."""
    from safetensors.torch import load_file
    from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer
    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
    from transformers.models.modernbert.modeling_modernbert import ModernBertRotaryEmbedding

    path = _resolve(path_or_repo, revision)
    with open(os.path.join(path, "quant_config.json"), encoding="utf-8") as f:
        qconfig = json.load(f)
    if qconfig.get("format") != FORMAT:
        raise ValueError(f"unsupported quantization format {qconfig.get('format')!r}")
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
    device = torch.device(device)
    if dtype is None:
        dtype = torch.bfloat16 if device.type == "cuda" and torch.cuda.is_bf16_supported() else torch.float32
    implementation = attention
    if attention == "linear":
        ALL_ATTENTION_FUNCTIONS.register("auto_linear", linear_attention)
        implementation = "auto_linear"
    config = AutoConfig.from_pretrained(path)
    with torch.device("meta"):
        model = AutoModelForSequenceClassification.from_config(config, attn_implementation=implementation, dtype=dtype, **model_kwargs)
        swap(model, qconfig, dtype=dtype)
    model.to_empty(device=device)
    model.load_state_dict(load_file(os.path.join(path, qconfig["weights"]), device=str(device)), strict=True)
    with torch.device(device):  # rotary tables are computed, not stored
        model.model.rotary_emb = ModernBertRotaryEmbedding(config=config)
    if attention == "linear":
        for layer in model.model.layers:
            layer.mlp._full_forward = layer.mlp.forward
            layer.mlp.forward = types.MethodType(_chunked_mlp, layer.mlp)
    model.auto_attention = attention
    return model.eval(), AutoTokenizer.from_pretrained(path)


@torch.inference_mode()
def logits(model, input_ids):
    """Classifier logits for one unpadded request (input_ids of shape [1, n]) on the model's device."""
    input_ids = input_ids.to(model.device)
    if getattr(model, "auto_attention", None) != "linear":
        return model(input_ids=input_ids).logits
    hidden = model.model(input_ids=input_ids, attention_mask={"full_attention": None, "sliding_attention": None})[0]
    pooled = hidden[:, 0] if model.config.classifier_pooling == "cls" else hidden.mean(dim=1)
    return model.classifier(model.drop(model.head(pooled)))


def score(model, tokenizer, text):
    """P(deny) for one serialized request (see build_input); deny if >= 0.5."""
    ids = tokenizer(text, return_tensors="pt", truncation=False)["input_ids"]
    if ids.shape[-1] > model.config.max_position_embeddings:
        raise ValueError(f"{ids.shape[-1]} tokens exceed the {model.config.max_position_embeddings}-token context")
    return logits(model, ids).float().softmax(-1)[0, 1].item()


def build_input(user_request, history, call):
    """The serialization every Auto model was trained on: proposed call, user request, agent history."""
    parts = ["### PROPOSED TOOL CALL", f"tool: {call['tool']}", f"args: {call['args']}", "", "### USER REQUEST", user_request, "", "### AGENT HISTORY"]
    if not history:
        parts.append("(no prior actions)")
    else:
        for i, h in enumerate(history):
            parts.append(f"[{i + 1}] {h['tool']}({h['args']})\n-> {h.get('result', '')}")
    return "\n".join(parts)