"""Load Whittle-Qwen-3.8 checkpoints (35B, 45B) with stock transformers (>= 5.18, Qwen4Exp).

Whittle keeps the Qwen4Exp architecture and its sharded n-gram table format, with two differences from the stock code:
  * geometry: 8 hash heads of exactly 4,880,000 rows each (the geometry its GGUFs carry for llama.cpp), where the stock
    code sizes each head as a distinct prime above `ngram_vocab_size_base`;
  * key prefix: the table pieces are saved as `model.ngram_embedding.shard_{i}.weight` and the hash buffers as
    `model.ple_embedding.*`, where the stock code uses `model.layers.{L}.ple.ple_embedding.*`.
This loader sizes the PLE n-gram table from the checkpoint's own hash buffers, maps the two prefixes onto the stock path
and lets transformers do the rest (it concatenates the pieces on CPU and keeps the ~20 GB lookup table off the GPU).
"""
import json
import mmap
import os
import re
import struct

import torch
from huggingface_hub import hf_hub_download
from safetensors import safe_open
from transformers import AutoConfig, AutoModelForCausalLM
from transformers.models.qwen4_exp import modeling_qwen4_exp as _qwen4exp


def _local(repo: str, filename: str) -> str:
    return os.path.join(repo, filename) if os.path.isdir(repo) else hf_hub_download(repo, filename)


FULL_INDEX = "whittle_full_index.json"   # optional: full index kept beside a load index that omits the table pieces


def _weight_map(repo: str) -> dict:
    if os.path.isdir(repo) and os.path.isfile(os.path.join(repo, FULL_INDEX)):
        return json.load(open(os.path.join(repo, FULL_INDEX)))["weight_map"]
    return json.load(open(_local(repo, "model.safetensors.index.json")))["weight_map"]


def _hash_buffers(repo: str) -> tuple[dict, int]:
    """Whittle's hash buffers, and the number of table rows its checkpoint pieces hold."""
    weight_map = _weight_map(repo)
    out = {}
    for name in ("ngram_heads_vocab_sizes", "ngram_heads_offsets", "layer_multipliers"):
        key = f"model.ple_embedding.{name}"
        with safe_open(_local(repo, weight_map[key]), "pt") as f:
            out[name] = f.get_tensor(key).tolist()
    rows = 0
    for key in (k for k in weight_map if k.startswith("model.ngram_embedding.shard_")):
        with safe_open(_local(repo, weight_map[key]), "pt") as f:
            rows += f.get_slice(key).get_shape()[0]
    return out, rows


class MmapNGramTable(torch.nn.Module):
    """The n-gram table read in place from the checkpoint's row pieces (memory-mapped): only looked-up rows are read,
    so the ~20 GB table costs almost no RAM - the same idea as llama.cpp's mmap. Behaves like nn.Embedding for lookups."""

    def __init__(self, pieces, dim: int, dtype):
        super().__init__()
        self._pieces, self._parts, starts, row = list(pieces), [], [], 0
        for path, offset, rows, stored in pieces:
            with open(path, "rb") as fh:
                buf = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ)
            self._parts.append(torch.frombuffer(buf, dtype=stored, count=rows * dim, offset=offset).view(rows, dim))
            starts.append(row); row += rows
        self.num_embeddings, self.embedding_dim, self._starts = row, dim, starts
        self.register_buffer("_bounds", torch.tensor(starts[1:], dtype=torch.long), persistent=False)
        self.register_buffer("weight", torch.empty(0, dim, dtype=dtype), persistent=False)   # device/dtype marker only

    def to_ram(self, chunk: int = 64 << 20) -> torch.nn.Embedding:
        """Copy the whole table into one in-memory nn.Embedding with large sequential file reads (much faster than
        faulting the memory map in page by page, especially on a network mount)."""
        stored = self._pieces[0][3]
        assert all(p[3] == stored for p in self._pieces), "pieces with mixed dtypes"
        table = torch.empty((self.num_embeddings, self.embedding_dim), dtype=stored)
        raw = memoryview(table.view(torch.uint8).numpy().reshape(-1))
        pos = 0
        for path, offset, rows, _ in self._pieces:
            nbytes = rows * self.embedding_dim * table.element_size()
            with open(path, "rb", buffering=0) as fh:
                fh.seek(offset)
                done = 0
                while done < nbytes:
                    n = fh.readinto(raw[pos + done: pos + min(nbytes, done + chunk)])
                    if not n:
                        raise IOError(f"short read in {path} at {offset + done}")
                    done += n
            pos += nbytes
        return torch.nn.Embedding.from_pretrained(table.to(self.weight.dtype), freeze=True)

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        flat = ids.reshape(-1).cpu()
        which = torch.bucketize(flat, self._bounds, right=True)
        out = torch.empty(flat.numel(), self.embedding_dim, dtype=self.weight.dtype)
        for i, part in enumerate(self._parts):
            sel = which == i
            if sel.any():
                out[sel] = part[flat[sel] - self._starts[i]].to(out.dtype)
        return out.view(*ids.shape, self.embedding_dim)


def _table_pieces(repo: str):
    """(path, absolute byte offset, rows, stored dtype) of each n-gram table piece, in row order."""
    weight_map = _weight_map(repo)
    keys = sorted((k for k in weight_map if re.fullmatch(r"model\.ngram_embedding\.shard_\d+\.weight", k)),
                  key=lambda k: int(k.split(".")[2].split("_")[1]))
    out = []
    for key in keys:
        path = _local(repo, weight_map[key])
        with open(path, "rb") as fh:
            n = struct.unpack("<Q", fh.read(8))[0]
            header = json.loads(fh.read(n))
        start, _ = header[key]["data_offsets"]
        stored = {"BF16": torch.bfloat16, "F16": torch.float16, "F32": torch.float32}[header[key]["dtype"]]
        out.append((path, 8 + n + start, header[key]["shape"][0], stored))
    return out


def load_whittle(repo: str, dtype=torch.bfloat16, device_map="cuda", table: str = "load", **kwargs):
    """table=\"load\" puts the n-gram table in CPU memory (~20 GB, plus a transient copy while loading);
    table=\"mmap\" reads it in place from the checkpoint files (needs the shards on a local or mounted path)."""
    assert table in ("load", "mmap"), table
    cfg = AutoConfig.from_pretrained(repo)
    prefix = f"model.layers.{cfg.ple_layer_ids[0] - 1}.ple.ple_embedding"
    hb, rows = _hash_buffers(repo)
    sizes, offsets = hb["ngram_heads_vocab_sizes"], hb["ngram_heads_offsets"]
    total = offsets[-1] + sizes[-1]
    assert rows >= total, f"table pieces hold {rows:,} rows, the hash addresses {total:,}"

    original_init = _qwen4exp.Qwen4ExpTextNGramEmbedding.__init__

    def whittle_init(self, config, embedding_dim, layer_idx, ple_layer_index=0):
        original_init(self, config, embedding_dim, layer_idx, ple_layer_index)
        assert len(sizes) == self.ngram_heads, (len(sizes), self.ngram_heads)
        self.head_vocab_sizes, self.head_offsets, self.total_vocab_size = list(sizes), list(offsets), total
        self.ngram_heads_vocab_sizes = torch.nn.Buffer(torch.tensor(sizes, dtype=torch.long))
        self.ngram_heads_offsets = torch.nn.Buffer(torch.tensor(offsets, dtype=torch.long))
        dim = self.ngram_embedding.embedding_dim
        # load: exactly the checkpoint's rows; mmap: a one-row placeholder, replaced by the memory-mapped table after loading
        self.ngram_embedding = torch.nn.Embedding(rows if table == "load" else 1, dim)

    _qwen4exp.Qwen4ExpTextNGramEmbedding.__init__ = whittle_init
    try:
        model = AutoModelForCausalLM.from_pretrained(
            repo, dtype=dtype, device_map=device_map,
            key_mapping={r"^model\.ple_embedding\.": f"{prefix}.",
                         # mmap: rename the pieces so transformers neither concatenates nor loads them
                         r"^model\.ngram_embedding\.": f"{prefix}.ngram_embedding." if table == "load" else "whittle_unloaded_table."},
            **kwargs,
        )
    finally:
        _qwen4exp.Qwen4ExpTextNGramEmbedding.__init__ = original_init

    emb = model.get_submodule(prefix)
    if table == "mmap":
        emb.ngram_embedding = MmapNGramTable(_table_pieces(repo), emb.ngram_embedding.embedding_dim, dtype)
    for name, want in hb.items():                     # the forward pass hashes with these buffers: they must be Whittle's
        got = getattr(emb, name).tolist()
        assert got == want, f"{name}: {got[:3]}... != checkpoint {want[:3]}..."
    assert getattr(emb.ngram_embedding, "num_embeddings", emb.ngram_embedding.weight.shape[0]) >= total
    return model