"""Minimal Clef runtime for one 16 GB Blackwell card.

The release's `load_release_model` puts the whole Qwen3.5 backbone, both 248k-row vocab
matrices and the vision tower on the GPU in bf16 -- 55 GB for Clef, 19 GB for Clef-Flash.
Neither fits a 16 GB card. This runtime keeps the decoder on the GPU and the rest in host RAM:

* embed_tokens -> CPU. A lookup of a few hundred rows; the result is copied to the GPU.
* lm_head      -> CPU. The joint head only gathers rows for the option tokens
                  (`output_embedding_weight[token_ids]`), never multiplies the full matrix.
* vision tower -> not loaded (text-only benchmark).

Decoder linear layers can be:
  bf16   -- as released
  nf4    -- bitsandbytes NF4 quantized at load time from the bf16 checkpoint
  nvfp4  -- a compressed-tensors NVFP4 checkpoint. W4A4 (activations quantized to FP4 per
            16-element block and multiplied on the FP4 tensor cores via torch._scaled_mm) when
            the checkpoint carries input_global_scale; otherwise W4A16 (dequantize + bf16 GEMM).
"""
from __future__ import annotations

import json
import math
import os
import re
import sys
from pathlib import Path

import torch
import torch.nn.functional as F
from safetensors import safe_open
from safetensors.torch import load_file

FP4_MAX, FP8_MAX = 6.0, 448.0
FP4_VALUES = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0])
NF4_TARGETS = re.compile(r"\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj|"
                         r"in_proj_qkv|in_proj_z|out_proj)$")
LM_PREFIX = "model.language_model."


class Checkpoint:
    """Lazy tensor access across a sharded safetensors checkpoint."""

    def __init__(self, path):
        self.path = Path(path)
        self.map = json.load(open(self.path / "model.safetensors.index.json"))["weight_map"]
        extra = self.path / "input_scales.safetensors"
        if extra.exists():
            with safe_open(extra, "pt") as h:
                for k in h.keys():
                    self.map[k] = extra.name
        self._h = {}

    def __contains__(self, k):
        return k in self.map

    def get(self, k):
        f = self.map[k]
        if f not in self._h:
            self._h[f] = safe_open(self.path / f, "pt")
        return self._h[f].get_tensor(k)


# --------------------------------------------------------------------------------- NVFP4
def to_blocked(scales: torch.Tensor) -> torch.Tensor:
    """cuBLAS block-scaled layout for 1x16 FP8 scales: 128x4 tiles, each re-tiled as 32x16."""
    rows, cols = scales.shape
    pr, pc = -rows % 128, -cols % 4
    if pr or pc:
        scales = F.pad(scales, (0, pc, 0, pr))
    r, c = scales.shape
    blocks = scales.view(r // 128, 128, c // 4, 4).permute(0, 2, 1, 3)
    return blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16).flatten()


def _quant_act(x: torch.Tensor, gscale: torch.Tensor):
    """bf16 [M, K] -> (packed fp4 [M, K/2] uint8, fp8 block scales [M, K/16]).
    Same math as compressed-tensors: block_scale = amax/6 * global_scale, stored in FP8."""
    m, k = x.shape
    xb = x.float().view(m, k // 16, 16)
    amax = xb.abs().amax(-1)
    s = (amax * (gscale / FP4_MAX)).clamp(max=FP8_MAX).to(torch.float8_e4m3fn)
    # Round |x|/eff to the e2m1 grid with ties to even, WITHOUT dividing: compare |x| against
    # the midpoints times eff. A division lets Inductor substitute x * (1/eff), which pushes the
    # (frequent -- bf16 over fp8 scales) exact ties up in magnitude: a systematic bias that took
    # Clef-Flash CLINC150 from 67 to 39. Products are correctly rounded in eager and compiled
    # code alike. Ties go to the even code: down at 0.25/1.25/2.5/5, up at 0.75/1.75/3.5.
    t = (s.float() / gscale).unsqueeze(-1)
    a = xb.abs()
    code = ((a > 0.25 * t).to(torch.uint8) + (a >= 0.75 * t) + (a > 1.25 * t) + (a >= 1.75 * t)
            + (a > 2.5 * t) + (a >= 3.5 * t) + (a > 5.0 * t))
    code = torch.where(t > 0, code, torch.zeros_like(code))
    code = code | ((xb < 0).to(torch.uint8) << 3)
    code = code.view(m, k // 2, 2)
    packed = code[..., 0] | (code[..., 1] << 4)
    return packed, s


_quant_act_compiled = None


def quant_act(x, gscale):
    global _quant_act_compiled
    if os.environ.get("CLEF_NO_COMPILE"):
        return _quant_act(x, gscale)
    if _quant_act_compiled is None:
        _quant_act_compiled = torch.compile(_quant_act, dynamic=True, fullgraph=True)
    return _quant_act_compiled(x, gscale)


def dequant_fp4(packed: torch.Tensor, scale: torch.Tensor, gscale: torch.Tensor, dtype=torch.bfloat16):
    n, half = packed.shape
    lo, hi = packed & 0x0F, packed >> 4
    codes = torch.stack((lo, hi), -1).view(n, half * 2)
    vals = FP4_VALUES.to(packed.device)[(codes & 7).long()]
    vals = torch.where((codes & 8) > 0, -vals, vals)
    w = vals.view(n, -1, 16) * (scale.float() / gscale).unsqueeze(-1)
    return w.view(n, half * 2).to(dtype)


class NVFP4Linear(torch.nn.Module):
    def __init__(self, packed, scale, wgs, igs=None, mode="w4a4"):
        super().__init__()
        self.out_features, self.in_features = packed.shape[0], packed.shape[1] * 2
        self.register_buffer("packed", packed)
        self.register_buffer("wgs", wgs.float().reshape(()))
        self.mode = mode if igs is not None else "w4a16"
        if self.mode == "w4a4":
            # only the cuBLAS-swizzled copy is kept: holding both costs 1.4 GiB on Clef 27B,
            # which is the difference between fitting 16 GB and OOM on the first request
            self.register_buffer("igs", igs.float().reshape(()))
            self.register_buffer("scale_blocked", to_blocked(scale))
            self.register_buffer("alpha", (1.0 / (self.igs * self.wgs)).reshape(()))
        else:
            self.register_buffer("scale", scale)

    def forward(self, x):
        shape = x.shape
        x2 = x.reshape(-1, self.in_features)
        if self.mode == "w4a16":
            y = F.linear(x2, dequant_fp4(self.packed, self.scale, self.wgs, x.dtype))
        else:
            xq, xs = quant_act(x2, self.igs)
            y = torch._scaled_mm(xq.view(torch.float4_e2m1fn_x2),
                                 self.packed.view(torch.float4_e2m1fn_x2).t(),
                                 to_blocked(xs), self.scale_blocked,
                                 out_dtype=torch.bfloat16)
            y = y * self.alpha.to(y.dtype)
        return y.view(*shape[:-1], self.out_features)


# --------------------------------------------------------------------------------- model
class CPUGather:
    """Stands in for the lm_head weight: the joint head only indexes rows out of it."""

    def __init__(self, w):
        self.w = w

    def __getitem__(self, idx):
        return self.w[idx.cpu()].to(idx.device, non_blocking=True)


def _set_param(module, name, t, device):
    old = getattr(module, name)
    if isinstance(old, torch.nn.Parameter):
        setattr(module, name, torch.nn.Parameter(t.to(device), requires_grad=False))
    else:
        module._buffers[name] = t.to(device)


def build_text_model(cfg):
    from transformers.models.qwen3_5.modeling_qwen3_5 import (Qwen3_5TextModel,
                                                             Qwen3_5TextRotaryEmbedding)
    cfg._attn_implementation = "sdpa"
    with torch.device("meta"):
        tm = Qwen3_5TextModel(cfg)
    tm.embed_tokens = torch.nn.Identity()  # embeddings are looked up on the CPU
    tm.rotary_emb = Qwen3_5TextRotaryEmbedding(cfg).to("cuda")
    return tm


def materialize(module, prefix, ck: Checkpoint, mode, device="cuda", dtype=torch.bfloat16):
    """Load every parameter under `module` from the checkpoint, swapping Linear layers for
    their quantized form where the mode/checkpoint says so."""
    for name, sub in list(module.named_modules()):
        full = prefix + name + ("." if name else "")
        if isinstance(sub, torch.nn.Linear):
            new = None
            if full + "weight_packed" in ck:
                igs = ck.get(full + "input_global_scale") if full + "input_global_scale" in ck else None
                new = NVFP4Linear(ck.get(full + "weight_packed").to(device),
                                  ck.get(full + "weight_scale").to(device),
                                  ck.get(full + "weight_global_scale").to(device),
                                  None if igs is None else igs.to(device),
                                  mode="w4a16" if mode == "nvfp4-w4a16" else "w4a4")
            elif mode == "nf4" and NF4_TARGETS.search("." + name):
                import bitsandbytes as bnb
                w = ck.get(full + "weight")
                new = bnb.nn.Linear4bit(sub.in_features, sub.out_features, bias=False,
                                        compute_dtype=dtype, quant_type="nf4",
                                        compress_statistics=True, device="meta")
                new.weight = bnb.nn.Params4bit(w.to(dtype), requires_grad=False, quant_type="nf4",
                                               compress_statistics=True, blocksize=64)
                new = new.to(device)
            if new is not None:
                parent_name, _, child = name.rpartition(".")
                parent = module.get_submodule(parent_name) if parent_name else module
                setattr(parent, child, new)
                continue
        for pname, p in list(sub.named_parameters(recurse=False)) + list(sub.named_buffers(recurse=False)):
            if p.device.type != "meta":
                continue
            key = full + pname
            t = ck.get(key)
            if t.is_floating_point() and pname not in ("A_log",):
                t = t.to(dtype)
            _set_param(sub, pname, t, device)


MLP_CHUNK = int(os.environ.get("CLEF_MLP_CHUNK", "4096"))


def chunk_mlp(mlp):
    """The MLP is per-token, so long prefills can run it in token chunks: same math, but the
    17408-wide intermediate of a 9k-token request no longer has to fit next to 13 GiB of weights."""
    fwd = mlp.forward

    def forward(x):
        if x.shape[-2] <= MLP_CHUNK:
            return fwd(x)
        return torch.cat([fwd(c) for c in x.split(MLP_CHUNK, dim=-2)], dim=-2)
    mlp.forward = forward


def assert_no_meta(module):
    left = [n for n, p in list(module.named_parameters()) + list(module.named_buffers())
            if p.device.type == "meta"]
    if left:
        raise RuntimeError(f"{len(left)} tensors never loaded, e.g. {left[:5]}")


class Clef:
    def __init__(self, path, mode="bf16", device="cuda"):
        from transformers import AutoConfig, AutoTokenizer
        self.path, self.mode, self.device = Path(path), mode, device
        sys.path.insert(0, str(self.path))
        import joint_schema_model as jsm
        self.jsm = jsm
        self.cfg = AutoConfig.from_pretrained(self.path).text_config
        ck = Checkpoint(self.path)
        self.tm = build_text_model(self.cfg)
        for i, layer in enumerate(self.tm.layers):
            materialize(layer, f"{LM_PREFIX}layers.{i}.", ck, mode, device)
        materialize(self.tm.norm, f"{LM_PREFIX}norm.", ck, mode, device)
        assert_no_meta(self.tm)
        for layer in self.tm.layers:
            chunk_mlp(layer.mlp)
        self.tm.eval()
        self.embed = ck.get(f"{LM_PREFIX}embed_tokens.weight").to(torch.bfloat16)
        self.lm_head = CPUGather(ck.get("lm_head.weight").to(torch.bfloat16))
        head_cfg = json.loads((self.path / "joint_head_config.json").read_text())
        self.head = jsm.JointSchemaHead(**head_cfg)
        self.head.load_state_dict(load_file(self.path / "joint_head.safetensors"), strict=True)
        self.head = self.head.to(device=device, dtype=torch.bfloat16).eval()
        self.tok = AutoTokenizer.from_pretrained(self.path)

    def encode(self, record, max_length=16384):
        return self.jsm.encode_record(self.tok, record, max_length=max_length)

    @torch.inference_mode()
    def forward(self, encoded: list):
        batch = self.jsm.collate_records(encoded, self.tok.pad_token_id, torch.device("cpu"))
        ids, mask = batch["input_ids"], batch["attention_mask"]
        emb = F.embedding(ids, self.embed).to(self.device, non_blocking=True)
        ids_d, mask_d = ids.to(self.device), mask.to(self.device)
        out = self.tm(inputs_embeds=emb, attention_mask=mask_d, use_cache=False)
        return self.head(out.last_hidden_state, ids_d, mask_d, batch["records"], self.lm_head)

    def probs(self, encoded: list):
        logits = self.forward(encoded)
        res = []
        for enc, rec_logits in zip(encoded, logits):
            res.append({q.question_id: dict(zip(q.option_ids, l.float().softmax(-1).tolist()))
                        for q, l in zip(enc.questions, rec_logits)})
        return res