#!/usr/bin/env python3
"""Quantize a Hugging Face model to llama.cpp NVFP4 GGUF.

Quantization lives in nvfp4gguf. This script loads the model through llama.cpp
so the vocabulary, tensor names, and architecture metadata match a normal GGUF.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import logging
import re
import shutil
import sys
from pathlib import Path

import numpy as np

from nvfp4gguf import (
    CTQ_NUM_ITER,
    DEFAULT_CONVERT_TO_QUANT,
    DEFAULT_LLAMA_CPP,
    DEFAULT_EARLY_STOP_FLOOR_STALL,
    DEFAULT_TARGET_LOSS,
    DEFAULT_EARLY_STOP_STALL,
    DEFAULT_SCALE_OPTIMIZATION,
    DEFAULT_SCALE_REFINEMENT_ROUNDS,
    QUANT_VERSION,
    TensorQuantizer,
    bind_exclude,
    eligible_weight,
    ensure_cuda_interpreter,
    load_gguf_modules,
    load_learned_converter,
    must_bake_nvfp4_scale,
    mtp_layer_tensor,
    bake_packed_nvfp4,
    preserve_scale_tensor,
    quantize_and_pack,
    report_reconstruction,
    require_quant_settings,
    scale_bake_notice,
    resolve_optimizer,
)
from nvfp4gguf.bundle import (
    apply_hf_context_length,
    chat_template_has_tools,
    flatten_chat_template,
    normalize_chat_template,
    outtype as _outtype,
    place_mmprojs as _place_mmprojs,
    quote_modelfile_value as _quote_modelfile_value,
    refuse_existing_bundle,
    tool_modelfile_lines,
    warn_missing_gguf_metadata,
)

log = logging.getLogger("hf_to_nvfp4_gguf")


def _checkpoint_path(directory: Path, name: str) -> Path:
    digest = hashlib.sha1(name.encode("utf-8")).hexdigest()[:16]
    safe = re.sub(r"[^A-Za-z0-9._-]+", "_", name)[:80]
    return directory / f"{safe}.{digest}.npz"


def _eligible(name: str, tensor, raw_dtype, args, gguf_module) -> bool:
    if raw_dtype != gguf_module.GGMLQuantizationType.BF16:
        return False
    shape = tuple(int(dim) for dim in tensor.shape)
    return eligible_weight(name, shape, args, gguf_module)


def _count_quant_targets(model, args, gguf_module) -> int:
    """Count matrices the writer will quantize, using shapes from the file header."""
    total = 0
    for hf_name, lazy in model.get_tensors():
        try:
            shape = tuple(int(dim) for dim in lazy.shape)
        except Exception:
            continue
        if len(shape) < 2 or not str(hf_name).endswith(".weight"):
            continue
        try:
            gguf_name = model.map_tensor_name(hf_name)
        except Exception:
            gguf_name = str(hf_name)
        raw_dtype = gguf_module.GGMLQuantizationType.BF16
        if gguf_name.endswith("_norm.weight"):
            raw_dtype = gguf_module.GGMLQuantizationType.F32
        if raw_dtype == gguf_module.GGMLQuantizationType.BF16 and eligible_weight(gguf_name, shape, args, gguf_module):
            total += 1
    return total


def _materialize(array) -> np.ndarray:
    gguf_lazy = sys.modules.get("gguf")
    lazy_cls = getattr(gguf_lazy, "LazyNumpyTensor", None) if gguf_lazy else None
    if lazy_cls is not None and isinstance(array, lazy_cls):
        array = lazy_cls.to_eager(array)
    array = np.asarray(array)
    if array.dtype != np.float32:
        array = array.astype(np.float32, copy=False)
    return np.ascontiguousarray(array)


def _as_float_tensor(array, gguf_module):
    lazy_cls = getattr(gguf_module, "LazyNumpyTensor", None)
    chunk_cls = getattr(gguf_module, "LazyChunkedTensor", None)
    if chunk_cls is not None and isinstance(array, chunk_cls):
        return array
    if lazy_cls is not None and isinstance(array, lazy_cls):
        return array
    return np.asarray(array)


class Nvfp4Writer:
    def __init__(self, model, quantizer: TensorQuantizer, args, bias_names: set[str]):
        self.model = model
        self.quantizer = quantizer
        self.args = args
        self.gguf = quantizer.gguf
        self.bias_names = set(bias_names)
        self.bias_corrections: dict[str, np.ndarray] = {}
        self.pending_bias: dict[str, tuple] = {}
        self.quantized = 0
        self.baked = 0
        self.corrected = 0
        self.passthrough = 0
        self.planned = _count_quant_targets(model, args, self.gguf)
        self.quant_index = 0
        if self.planned:
            log.info("%d weight matrices queued for NVFP4", self.planned)
        names = getattr(self.gguf, "MODEL_ARCH_NAMES", {})
        self.architecture = names.get(getattr(model, "model_arch", None), "") or ""
        notice = scale_bake_notice(self.architecture)
        if notice:
            log.warning(notice)
        self.ckpt_dir = Path(args.checkpoint_dir) if args.checkpoint_dir else Path(str(args.output) + ".nvfp4-work")
        self.ckpt_dir.mkdir(parents=True, exist_ok=True)
        self._orig_quantize = self.gguf.quants.quantize
        self._orig_add = model.gguf_writer.add_tensor
        self._orig_prepare = model.prepare_tensors
        self._orig_metadata = model.prepare_metadata

    def install(self) -> None:
        gguf_module = self.gguf

        def quantize(data, qtype):
            # Keep the source float for matrices we may repack. Everything else
            # still goes through llama.cpp's quantizer.
            if qtype == gguf_module.GGMLQuantizationType.BF16:
                return data
            return self._orig_quantize(data, qtype)

        self.gguf.quants.quantize = quantize
        self.model.gguf_writer.add_tensor = self._add_tensor
        self.model.prepare_tensors = self._prepare_tensors
        self.model.prepare_metadata = self._prepare_metadata

    def restore(self) -> None:
        self.gguf.quants.quantize = self._orig_quantize
        self.model.gguf_writer.add_tensor = self._orig_add
        self.model.prepare_tensors = self._orig_prepare
        self.model.prepare_metadata = self._orig_metadata

    def _prepare_tensors(self) -> None:
        self._orig_prepare()
        self._flush_pending()

    def _prepare_metadata(self, vocab_only: bool = False) -> None:
        if self.quantized:
            self.model.ftype = self.gguf.LlamaFileType.MOSTLY_NVFP4
        self._orig_metadata(vocab_only)

    def _write(self, name, tensor, raw_dtype=None) -> None:
        self._orig_add(name, tensor, raw_dtype=raw_dtype)

    def _add_tensor(self, name, tensor, raw_shape=None, raw_dtype=None, tensor_endianess=None) -> None:
        del raw_shape, tensor_endianess
        if name.endswith(".bias"):
            self._handle_bias(name, tensor, raw_dtype)
            return
        if _eligible(name, tensor, raw_dtype, self.args, self.gguf):
            self._quantize_and_write(name, tensor)
            return
        if raw_dtype == self.gguf.GGMLQuantizationType.BF16:
            tensor = self._orig_quantize(tensor, raw_dtype)
        self._write(name, tensor, raw_dtype=raw_dtype)
        self.passthrough += 1

    def _handle_bias(self, name: str, tensor, raw_dtype) -> None:
        correction = self.bias_corrections.get(name)
        if correction is None and not self.args.no_bias_correction:
            self.pending_bias[name] = (tensor, raw_dtype)
            return
        self._emit_bias(name, tensor, raw_dtype, correction)

    def _emit_bias(self, name: str, tensor, raw_dtype, correction: np.ndarray | None) -> None:
        if correction is not None:
            values = _materialize(tensor).astype(np.float32, copy=True)
            flat = values.reshape(-1)
            if flat.shape == correction.shape:
                flat -= correction
                values = flat.reshape(values.shape)
                self.corrected += 1
                log.info("    bias %s adjusted (mean %+.3e)", name, float(correction.mean()))
            else:
                log.warning("    bias %s shape %s does not match correction %s", name, values.shape, correction.shape)
            tensor = values.astype(np.float32, copy=False)
            if raw_dtype == self.gguf.GGMLQuantizationType.BF16:
                tensor = self._orig_quantize(tensor, raw_dtype)
            elif raw_dtype not in (None, self.gguf.GGMLQuantizationType.F32):
                tensor = self._orig_quantize(tensor, raw_dtype)
        elif raw_dtype == self.gguf.GGMLQuantizationType.BF16:
            tensor = self._orig_quantize(tensor, raw_dtype)
        self._write(name, tensor, raw_dtype=raw_dtype)
        self.passthrough += 1

    def _flush_pending(self) -> None:
        for name, (tensor, raw_dtype) in list(self.pending_bias.items()):
            self._emit_bias(name, tensor, raw_dtype, self.bias_corrections.get(name))
        self.pending_bias.clear()

    def _load_checkpoint(self, name: str, logical_shape: tuple[int, ...]):
        path = _checkpoint_path(self.ckpt_dir, name)
        if not path.is_file():
            return None
        with np.load(path) as blob:
            if int(blob["version"]) != QUANT_VERSION:
                return None
            if tuple(int(v) for v in blob["shape"]) != tuple(logical_shape):
                return None
            raw = np.array(blob["raw"], copy=True)
            scale = np.array(blob["scale"], copy=True)
            bias = np.array(blob["bias"], copy=True)
        scale_out = None if scale.size == 0 else scale
        bias_out = None if bias.size == 0 else bias
        return raw, scale_out, bias_out

    def _save_checkpoint(self, name: str, logical_shape, raw, scale, bias) -> None:
        path = _checkpoint_path(self.ckpt_dir, name)
        tmp = path.with_name(path.name + ".tmp")
        with tmp.open("wb") as handle:
            np.savez_compressed(
                handle,
                version=np.int32(QUANT_VERSION),
                shape=np.array(logical_shape, dtype=np.int64),
                raw=raw,
                scale=np.zeros((0,), np.float32) if scale is None else scale.astype(np.float32, copy=False),
                bias=np.zeros((0,), np.float32) if bias is None else bias.astype(np.float32, copy=False),
            )
        tmp.replace(path)

    def _quantize_matrix_group(self, name: str, weight: np.ndarray) -> tuple[np.ndarray, np.ndarray | None, np.ndarray | None]:
        baked = self._must_bake(name)
        raw, scale = quantize_and_pack(self.quantizer, weight, baked, name)
        bias = None
        bias_name = name.replace(".weight", ".bias")
        if weight.ndim == 2 and not self.args.no_bias_correction and bias_name in self.bias_names:
            per_tensor = None if scale is None else float(np.reshape(scale, -1)[0])
            runtime = self.quantizer.runtime_dequant(raw, per_tensor)
            bias = self.quantizer.bias_correction(weight, runtime)
        report_reconstruction(self.quantizer, weight, raw, scale)
        return raw, scale, bias

    def _must_bake(self, name: str) -> bool:
        bias_name = name.replace(".weight", ".bias")
        has_bias = bias_name in self.bias_names or bias_name in self.pending_bias
        if must_bake_nvfp4_scale(self.architecture, name, has_bias):
            return True
        return mtp_layer_tensor(name, *_mtp_depth(self.model))

    def _progress(self, name: str, shape: tuple[int, ...], checkpoint: bool) -> None:
        self.quant_index += 1
        if self.planned:
            slot = f"[{self.quant_index}/{self.planned}]"
        else:
            slot = f"[{self.quant_index}]"
        note = "  (checkpoint)" if checkpoint else ""
        log.info("NVFP4 %s %s  %s%s", slot, name, shape, note)

    def _quantize_and_write(self, name: str, tensor) -> None:
        weight = _materialize(tensor)
        baked = self._must_bake(name)
        cached = self._load_checkpoint(name, weight.shape)
        if cached is not None and (cached[1] is None) != baked:
            log.info("    checkpoint scale layout does not match %s; recomputing", name)
            cached = None
        self._progress(name, tuple(weight.shape), cached is not None)
        if cached is None:
            raw, scale, bias = self._quantize_matrix_group(name, weight)
            self._save_checkpoint(name, weight.shape, raw, scale, bias)
        else:
            raw, scale, bias = cached
        if bias is not None:
            self.bias_corrections[name.replace(".weight", ".bias")] = bias
            pending = self.pending_bias.pop(name.replace(".weight", ".bias"), None)
            if pending is not None:
                self._emit_bias(name.replace(".weight", ".bias"), pending[0], pending[1], bias)
        if self._must_bake(name) and scale is None:
            self.baked += 1
        self._write(name, raw, raw_dtype=self.gguf.GGMLQuantizationType.NVFP4)
        if scale is not None:
            scale_name = name[: -len(".weight")] + ".scale"
            self._write(scale_name, np.ascontiguousarray(scale, dtype=np.float32))
        self.quantized += 1
        del weight, raw


def _accept_list_extra_special_tokens() -> None:
    """Let transformers 4.x load a Gemma 4 tokenizer.

    Gemma 4 stores extra_special_tokens as a list of token strings. The 4.x
    loader calls .keys() on that field, which is what aborts set_vocab after
    the weights have already been packed.
    """
    try:
        from transformers.tokenization_utils_base import PreTrainedTokenizerBase
    except ImportError:
        return
    original = PreTrainedTokenizerBase._set_model_specific_special_tokens
    if getattr(original, "_accepts_token_list", False):
        return

    def _set_model_specific_special_tokens(self, special_tokens):
        if isinstance(special_tokens, list):
            log.info(
                "tokenizer extra_special_tokens is a list; treating each entry as its own special token."
            )
            special_tokens = {str(token): str(token) for token in special_tokens}
        return original(self, special_tokens)

    _set_model_specific_special_tokens._accepts_token_list = True
    PreTrainedTokenizerBase._set_model_specific_special_tokens = _set_model_specific_special_tokens


def _accept_missing_tokenizer_class() -> None:
    """Load tokenizer.json when tokenizer_class is newer than this transformers.

    LFM3 names its tokenizer ``TokenizersBackend``. That class arrived in
    transformers 5, and a 4.x install otherwise stops before the vocab is written.
    """
    try:
        from transformers import AutoTokenizer, PreTrainedTokenizerFast
    except ImportError:
        return
    current = AutoTokenizer.from_pretrained
    if getattr(current, "_accepts_missing_class", False):
        return
    original = current.__func__

    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        try:
            return original(cls, pretrained_model_name_or_path, *args, **kwargs)
        except ValueError as exc:
            if "does not exist or is not currently imported" not in str(exc):
                raise
            log.info(
                "Tokenizer class is not in this transformers build; loading tokenizer.json directly."
            )
            return PreTrainedTokenizerFast.from_pretrained(
                pretrained_model_name_or_path, *args, **kwargs
            )

    from_pretrained._accepts_missing_class = True
    wrapped = classmethod(from_pretrained)
    wrapped._accepts_missing_class = True
    AutoTokenizer.from_pretrained = wrapped


def _install_tokenizer_fallback(model) -> None:
    """Use tokenizer.json when sentencepiece is not installed and no tokenizer.model is present."""
    _accept_list_extra_special_tokens()
    _accept_missing_tokenizer_class()
    try:
        import sentencepiece  # noqa: F401
        return
    except ImportError:
        pass
    if (model.dir_model / "tokenizer.model").is_file():
        original = model.set_vocab

        def set_vocab_needs_sentencepiece() -> None:
            try:
                original()
            except ModuleNotFoundError as exc:
                if getattr(exc, "name", None) == "sentencepiece":
                    raise ModuleNotFoundError(
                        "This model ships tokenizer.model, which needs the sentencepiece package. "
                        "pip install sentencepiece"
                    ) from exc
                raise

        model.set_vocab = set_vocab_needs_sentencepiece
        return
    if not (model.dir_model / "tokenizer.json").is_file():
        return

    def set_vocab() -> None:
        config_path = model.dir_model / "tokenizer_config.json"
        if config_path.is_file():
            with config_path.open("r", encoding="utf-8") as handle:
                prefix = json.load(handle).get("add_prefix_space")
            if prefix is not None:
                model.gguf_writer.add_add_space_prefix(prefix)
        log.info("sentencepiece is not installed; loading tokenizer.json")
        model._set_vocab_gpt2()

    model.set_vocab = set_vocab


def _collect_bias_names(model) -> set[str]:
    names: set[str] = set()
    for hf_name in model.model_tensors:
        if not hf_name.endswith(".bias"):
            continue
        try:
            names.add(model.map_tensor_name(hf_name))
        except Exception:
            continue
    return names


def _float_file_type(gguf, outtype: str):
    return {
        "f32": gguf.LlamaFileType.ALL_F32,
        "f16": gguf.LlamaFileType.MOSTLY_F16,
        "bf16": gguf.LlamaFileType.MOSTLY_BF16,
        "auto": gguf.LlamaFileType.GUESSED,
    }[outtype]


_PACKED_WEIGHT_DTYPES = {"U8", "I8"}
_PACKED_SCALE_DTYPES = {"U8", "I8", "F8_E4M3", "F8_E5M2"}
_FP8_DTYPES = {"F8_E4M3", "F8_E5M2"}


def _numel(shape: tuple[int, ...]) -> int:
    count = 1
    for dim in shape:
        count *= int(dim)
    return count


def _is_packed_nvfp4_layout(
    weight_dtype: str,
    weight_shape: tuple[int, ...],
    scale_dtype: str,
    scale_shape: tuple[int, ...],
) -> bool:
    """Match llama.cpp's packed-4-bit check using safetensors header dtypes.

    Packed NVFP4 weights are U8 or I8, two 4-bit values per byte, with one
    scale byte per 16 values (8 stored weight bytes per scale). A weight
    labeled F8_E4M3 is an 8-bit float, not packed 4-bit data.
    """
    if weight_dtype not in _PACKED_WEIGHT_DTYPES or scale_dtype not in _PACKED_SCALE_DTYPES:
        return False
    if len(weight_shape) < 2 or len(scale_shape) < 2:
        return False
    rows = int(weight_shape[0])
    n_blocks = int(scale_shape[1])
    if rows <= 0 or n_blocks <= 0 or n_blocks % 4 != 0:
        return False
    return _numel(weight_shape) == rows * n_blocks * 8 and _numel(scale_shape) == rows * n_blocks


def _scale_name_for_weight(name: str) -> str | None:
    if name.endswith(".weight_packed"):
        return name.removesuffix("_packed") + "_scale"
    if name.endswith(".weight"):
        return name[: -len(".weight")] + ".weight_scale"
    return None


def _sample_names(names: list[str], limit: int = 4) -> str:
    if len(names) <= limit:
        return ", ".join(names)
    return ", ".join(names[:limit]) + f", ... ({len(names) - limit} more)"


class _TensorLayout:
    def __init__(self) -> None:
        self.packed: list[str] = []
        self.fp8: list[str] = []
        self.unrecognized: list[str] = []

    @property
    def found(self) -> bool:
        return bool(self.packed or self.fp8 or self.unrecognized)


def _checkpoint_tensor_layout(model_dir: Path) -> _TensorLayout:
    """Classify safetensors by dtype and shape without reading the bytes."""
    layout = _TensorLayout()
    files = sorted(model_dir.glob("*.safetensors"))
    if not files:
        return layout
    try:
        from safetensors import safe_open
    except ImportError:
        log.warning("safetensors is not installed, so packed 4-bit tensors were not prechecked.")
        return layout
    meta: dict[str, tuple[str, tuple[int, ...]]] = {}
    for path in files:
        with safe_open(path, framework="pt", device="cpu") as handle:
            for key in handle.keys():
                header = handle.get_slice(key)
                meta[key] = (str(header.get_dtype()), tuple(int(dim) for dim in header.get_shape()))
    for name, (dtype, shape) in meta.items():
        scale_name = _scale_name_for_weight(name)
        if scale_name is None:
            continue
        scale = meta.get(scale_name)
        if scale is not None and _is_packed_nvfp4_layout(dtype, shape, scale[0], scale[1]):
            layout.packed.append(name)
            continue
        if dtype in _FP8_DTYPES:
            layout.fp8.append(name)
        elif dtype in {"U8", "I8"} and scale is not None:
            layout.unrecognized.append(name)
    layout.packed.sort()
    layout.fp8.sort()
    layout.unrecognized.sort()
    return layout


def _log_tensor_layout(layout: _TensorLayout) -> None:
    if not layout.found:
        return
    log.info(
        "Safetensors layout: %d packed 4-bit, %d FP8 weights that are not 4-bit, %d other byte weights that are not packed 4-bit.",
        len(layout.packed),
        len(layout.fp8),
        len(layout.unrecognized),
    )
    if layout.packed:
        log.info("Packed 4-bit, repacked as NVFP4: %s", _sample_names(layout.packed))
    if layout.fp8:
        log.info("FP8, stored as float: %s", _sample_names(layout.fp8))
    if layout.unrecognized:
        log.warning(
            "These tensors are U8 or I8 but do not match packed 4-bit NVFP4, so they will not be repacked: %s",
            _sample_names(layout.unrecognized),
        )


def _checkpoint_is_nvfp4(model_dir: Path, hparams: dict) -> bool:
    """Match llama.cpp's NVFP4 detection in ModelBase.prepare_tensors."""
    quantization_config = hparams.get("quantization_config") or {}
    quant_algo = quantization_config.get("quant_algo")
    quant_method = quantization_config.get("quant_method")
    quant_format = quantization_config.get("format")
    quant_groups = quantization_config.get("config_groups") or {}
    quant_layers = quantization_config.get("quantized_layers") or {}
    quant_config_file = model_dir / "hf_quant_config.json"
    if (not quant_algo or not quant_layers) and quant_config_file.is_file():
        with quant_config_file.open("r", encoding="utf-8") as handle:
            hf_quant_config = json.load(handle)
        quant_config = hf_quant_config.get("quantization") or {}
        producer = hf_quant_config.get("producer") or {}
        if quant_method is None:
            quant_method = (producer.get("name") or "").lower()
        quant_algo = quant_config.get("quant_algo", quant_algo)
        quant_method = quant_config.get("quant_method", quant_method)
        quant_format = quant_config.get("format", quant_format)
        quant_groups = quant_config.get("config_groups", quant_groups) or {}
        quant_layers = quant_config.get("quantized_layers", quant_layers) or {}
    if quant_algo in ("NVFP4", "W4A16_NVFP4"):
        return True
    if quant_method == "compressed-tensors" and (
        quant_format == "nvfp4-pack-quantized"
        or (
            quant_format == "mixed-precision"
            and any(
                isinstance(group, dict) and group.get("format") == "nvfp4-pack-quantized"
                for group in quant_groups.values()
            )
        )
    ):
        return True
    return any(
        isinstance(entry, dict) and str(entry.get("quant_algo", "")).endswith("NVFP4")
        for entry in quant_layers.values()
    )


def _token_text(value) -> str | None:
    if isinstance(value, str):
        return value or None
    if isinstance(value, dict):
        content = value.get("content")
        return content if isinstance(content, str) and content else None
    return None


def _hf_stops(model_dir: Path) -> list[str]:
    """Stop strings Ollama should use for this tokenizer.

    Llama 3's EOS list includes <|eom_id|>, which ends a tool-call turn inside
    the chat template. Stopping on it clips the call. The working Llama 3.1
    Modelfile stops on the header tokens and <|eot_id|> instead.
    """
    config_path = model_dir / "config.json"
    config = json.loads(config_path.read_text(encoding="utf-8")) if config_path.is_file() else {}
    tok_cfg: dict = {}
    tokenizer_config = model_dir / "tokenizer_config.json"
    if tokenizer_config.is_file():
        tok_cfg = json.loads(tokenizer_config.read_text(encoding="utf-8"))
    decoder = tok_cfg.get("added_tokens_decoder") or {}
    known = set()
    by_id = {}
    for key, value in decoder.items():
        try:
            token = _token_text(value)
        except (TypeError, ValueError):
            continue
        if token:
            known.add(token)
        try:
            by_id[int(key)] = token
        except (TypeError, ValueError):
            continue
    eos_text = _token_text(tok_cfg.get("eos_token"))
    if eos_text:
        known.add(eos_text)
    if "<|eot_id|>" in known or "<|start_header_id|>" in known:
        return [token for token in ("<|start_header_id|>", "<|end_header_id|>", "<|eot_id|>") if token in known]
    stops: list[str] = []

    def add(token: str | None) -> None:
        if token and token not in stops and token != "<|eom_id|>":
            stops.append(token)

    add(eos_text)
    eos_ids = config.get("eos_token_id")
    if isinstance(eos_ids, int):
        eos_ids = [eos_ids]
    for token_id in eos_ids or []:
        try:
            add(by_id.get(int(token_id)))
        except (TypeError, ValueError):
            continue
    if str(config.get("model_type") or "").lower().startswith("gemma"):
        if "<end_of_turn>" in known:
            stops = ["<end_of_turn>"] + [token for token in stops if token != "<end_of_turn>"]
    return stops


def _hf_sampler(model_dir: Path) -> list[tuple[str, object]]:
    path = model_dir / "generation_config.json"
    if not path.is_file():
        return []
    config = json.loads(path.read_text(encoding="utf-8"))
    chosen = []
    for key in ("temperature", "top_p", "top_k", "min_p", "repeat_penalty"):
        value = config.get(key)
        if isinstance(value, bool) or value is None:
            continue
        if isinstance(value, (int, float)):
            chosen.append((key, value))
    return chosen


def _hf_chat_template(model_dir: Path) -> str:
    """The chat template llama.cpp will store, when the checkpoint has one."""
    tok_cfg: dict = {}
    tokenizer_config = model_dir / "tokenizer_config.json"
    if tokenizer_config.is_file():
        tok_cfg = json.loads(tokenizer_config.read_text(encoding="utf-8"))
    if "chat_template" in tok_cfg and tok_cfg.get("chat_template") is not None:
        return flatten_chat_template(tok_cfg.get("chat_template"))
    jinja = model_dir / "chat_template.jinja"
    if jinja.is_file():
        return jinja.read_text(encoding="utf-8")
    return ""


def _write_hf_modelfile(
    model_dir: Path,
    gguf_path: Path,
    overwrite: bool,
    mmprojs: list[Path] | None = None,
) -> Path:
    """Write a Modelfile next to the GGUF."""
    modelfile = gguf_path.parent / "Modelfile"
    if modelfile.exists() and not overwrite:
        raise SystemExit(f"{modelfile} already exists. Pass --overwrite to replace it.")
    projectors = _place_mmprojs(gguf_path, mmprojs, overwrite)
    stops = _hf_stops(model_dir)
    lines = [f"FROM ./{gguf_path.name}"]
    for projector in projectors:
        lines.append(f"FROM ./{projector.name}")
    lines.append("")
    lines.append("# The chat template is the one stored in the GGUF.")
    template = _hf_chat_template(model_dir)
    lines.extend(tool_modelfile_lines(template))
    if chat_template_has_tools(template):
        log.info("Chat template supports tools.")
    else:
        log.info("Chat template has no tool round-trip. Tool requests stay rejected.")
    for stop in stops:
        lines.append(f"PARAMETER stop {_quote_modelfile_value(stop)}")
    for key, value in _hf_sampler(model_dir):
        rendered = str(value).lower() if isinstance(value, bool) else str(value)
        lines.append(f"PARAMETER {key} {rendered}")
    lines.append("# Context length comes from OLLAMA_CONTEXT_LENGTH.")
    lines.append("")
    modelfile.write_text("\n".join(lines), encoding="utf-8")
    log.info("Wrote %s", modelfile)
    return modelfile


def _mtp_depth(model) -> tuple[int, int]:
    n_blocks = int(getattr(model, "block_count", 0) or 0)
    hparams = getattr(model, "hparams", {}) or {}
    n_nextn = int(hparams.get("num_nextn_predict_layers") or 0)
    hidden = hparams.get("num_hidden_layers")
    if n_nextn <= 0 and hidden is not None and n_blocks > int(hidden):
        n_nextn = n_blocks - int(hidden)
    return n_blocks, n_nextn


def _fold_unclaimed_scale(name: str, architecture: str, n_blocks: int, n_nextn: int) -> bool:
    """True when this sidecar is a per-tensor NVFP4 scale the graph never multiplies."""
    if not name.endswith(".scale"):
        return False
    weight = name[: -len(".scale")] + ".weight"
    if preserve_scale_tensor(architecture, weight):
        return False
    if mtp_layer_tensor(weight, n_blocks, n_nextn):
        return True
    return must_bake_nvfp4_scale(architecture, weight)


def _fold_mtp_scales(model) -> None:
    """Bake NVFP4 scales the graph never multiplies, and do not write the sidecars."""
    writer = model.gguf_writer
    architecture = str(getattr(writer, "arch", "") or "")
    n_blocks, n_nextn = _mtp_depth(model)
    original = writer.add_tensor
    pending: dict[str, tuple] = {}

    def add_tensor(name, tensor, raw_shape=None, raw_dtype=None, tensor_endianess=None):
        text = str(name)
        if _fold_unclaimed_scale(text, architecture, n_blocks, n_nextn):
            weight_name = text[: -len(".scale")] + ".weight"
            held = pending.pop(weight_name, None)
            if held is None:
                return
            scale = float(np.asarray(tensor).reshape(-1)[0])
            original(weight_name, bake_packed_nvfp4(held[0], scale), raw_dtype=held[1])
            return
        scale_name = text[: -len(".weight")] + ".scale" if text.endswith(".weight") else ""
        if getattr(raw_dtype, "name", "") == "NVFP4" and _fold_unclaimed_scale(scale_name, architecture, n_blocks, n_nextn):
            pending[text] = (np.array(tensor, copy=True), raw_dtype)
            return
        original(name, tensor, raw_shape=raw_shape, raw_dtype=raw_dtype, tensor_endianess=tensor_endianess)

    def write_header_to_file(path=None):
        for held_name, (raw, dtype) in pending.items():
            original(held_name, raw, raw_dtype=dtype)
        pending.clear()
        return original_header(path)

    writer.add_tensor = add_tensor
    original_header = writer.write_header_to_file
    writer.write_header_to_file = write_header_to_file


def _normalize_chat_templates(model) -> None:
    writer = model.gguf_writer
    original = writer.add_chat_template

    def add_chat_template(template, *args, **kwargs):
        return original(normalize_chat_template(template), *args, **kwargs)

    writer.add_chat_template = add_chat_template


def _warn_when_header_is_written(model) -> None:
    writer = model.gguf_writer
    original = writer.write_header_to_file

    def write_header_to_file(path=None):
        warn_missing_gguf_metadata(writer)
        return original(path)

    writer.write_header_to_file = write_header_to_file


def _assign_output(args) -> None:
    """Point args.output at the GGUF file inside the bundle directory."""
    model_dir = Path(args.model)
    if args.output is None:
        if args.no_quant and args.outtype == "auto":
            directory = Path.cwd() / model_dir.name
            gguf_name = f"{model_dir.name}-{{FTYPE}}.gguf"
        elif args.no_quant:
            label = args.outtype.upper()
            directory = Path.cwd() / f"{model_dir.name}-{label}"
            gguf_name = f"{directory.name}.gguf"
        else:
            directory = Path.cwd() / f"{model_dir.name}-NVFP4"
            gguf_name = f"{directory.name}.gguf"
    else:
        output = Path(args.output)
        if output.suffix == ".gguf":
            gguf_path = output.resolve()
            gguf_path.parent.mkdir(parents=True, exist_ok=True)
            args.output = str(gguf_path)
            return
        directory = output
        gguf_name = f"{directory.name}.gguf"
    directory = directory.resolve()
    directory.mkdir(parents=True, exist_ok=True)
    args.output = str(directory / gguf_name)


def _refuse_existing_output(args) -> None:
    refuse_existing_bundle(
        Path(args.output),
        args.overwrite,
        getattr(args, "mmproj", None),
        skip_gguf=bool(args.modelfile_only),
    )


def export_float(args) -> None:
    """Convert the Hugging Face model to GGUF without learned rounding.

    Packed 4-bit tensors are repacked as NVFP4. U8 is the packed weight and
    F8_E4M3 is usually its block scale; a weight whose own dtype is F8_E4M3
    is an 8-bit float and is stored as float. The shape has to be 8 weight
    bytes per scale before a U8 tensor is treated as 4-bit. ``--outtype nvfp4``
    asks for that repack. MOSTLY_NVFP4 is selected only after at least one
    tensor is actually repacked.
    """
    gguf, _model_base, model_type_enum, get_model_architecture, get_model_class = load_gguf_modules(args.llama_cpp)
    model_dir = Path(args.model).resolve()
    if not model_dir.is_dir():
        raise FileNotFoundError(f"model directory does not exist: {model_dir}")
    passthrough = args.outtype == "nvfp4"

    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
    log.info("Loading %s", model_dir)
    hparams = _model_base.load_hparams(model_dir, False)
    layout = _checkpoint_tensor_layout(model_dir)
    _log_tensor_layout(layout)
    config_nvfp4 = _checkpoint_is_nvfp4(model_dir, hparams)
    if passthrough and not config_nvfp4 and not layout.packed:
        if layout.fp8:
            log.info(
                "No packed 4-bit tensors were found. %d FP8 weights will be stored as float.",
                len(layout.fp8),
            )
        else:
            raise SystemExit(
                f"{model_dir} has no packed 4-bit NVFP4 tensors. "
                "Use --outtype f16, bf16, f32, or auto to write float weights, "
                "or omit --outtype nvfp4 to fit NVFP4 from float weights."
            )
    architecture = get_model_architecture(hparams, model_type_enum.TEXT)
    log.info("Architecture: %s", architecture)
    model_cls = get_model_class(architecture)
    file_type = gguf.LlamaFileType.MOSTLY_BF16 if args.outtype == "nvfp4" else _float_file_type(gguf, args.outtype)
    model = model_cls(
        model_dir,
        file_type,
        Path(args.output).resolve(),
        use_temp_file=not args.no_temp_file,
        eager=args.eager,
        model_name=args.model_name,
    )
    quant_config = model.hparams.get("quantization_config")
    if layout.fp8 and not layout.packed and not config_nvfp4 and not (isinstance(quant_config, dict) and quant_config.get("quant_method")):
        merged = dict(quant_config) if isinstance(quant_config, dict) else {}
        merged["quant_method"] = "modelopt"
        model.hparams["quantization_config"] = merged
        log.info("FP8 weights have no quant method in the config. Their scales will be applied and the weights stored as float.")
    _install_tokenizer_fallback(model)
    apply_hf_context_length(model.hparams, args.context_length)
    _warn_when_header_is_written(model)
    _normalize_chat_templates(model)
    _fold_mtp_scales(model)
    if layout.packed or config_nvfp4:
        log.info("Repacking tensors that are packed 4-bit. Other weights are stored as float. Learned rounding is skipped.")
    else:
        stored = "bf16" if args.outtype == "nvfp4" else args.outtype
        log.info("Writing float GGUF (%s). NVFP4 quantization is skipped.", stored)
    model.write()
    _write_hf_modelfile(model_dir, Path(model.fname_out), args.overwrite, args.mmproj)
    log.info("Wrote %s", model.fname_out)


def export_nvfp4(args) -> None:
    if args.modelfile_only:
        logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
        model_dir = Path(args.model).resolve()
        if not model_dir.is_dir():
            raise FileNotFoundError(f"model directory does not exist: {model_dir}")
        log.info("Skipping quantization and GGUF creation; writing Modelfile only")
        _write_hf_modelfile(model_dir, Path(args.output), args.overwrite, args.mmproj)
        return
    _refuse_existing_output(args)
    if args.no_quant or args.outtype == "nvfp4":
        export_float(args)
        return
    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
    ensure_cuda_interpreter(args.device)
    gguf, _model_base, model_type_enum, get_model_architecture, get_model_class = load_gguf_modules(args.llama_cpp)
    model_dir = Path(args.model).resolve()
    if not model_dir.is_dir():
        raise FileNotFoundError(f"model directory does not exist: {model_dir}")

    args.optimizer = resolve_optimizer(args.optimizer)
    args.output = str(Path(args.output).resolve())
    log.info("Loading %s", model_dir)
    hparams = _model_base.load_hparams(model_dir, False)
    architecture = get_model_architecture(hparams, model_type_enum.TEXT)
    log.info("Architecture: %s", architecture)
    bind_exclude(args, architecture)
    settings_path_parent = Path(args.checkpoint_dir) if args.checkpoint_dir else Path(args.output + ".nvfp4-work")
    if args.fresh and settings_path_parent.exists():
        shutil.rmtree(settings_path_parent)
    settings_path_parent.mkdir(parents=True, exist_ok=True)
    require_quant_settings(settings_path_parent, args)
    model_cls = get_model_class(architecture)
    output = Path(args.output).resolve()
    model = model_cls(
        model_dir,
        gguf.LlamaFileType.MOSTLY_BF16,
        output,
        use_temp_file=not args.no_temp_file,
        eager=args.eager,
        model_name=args.model_name,
    )
    quantizer = TensorQuantizer(args, gguf, load_learned_converter(args.convert_to_quant)[1])
    _install_tokenizer_fallback(model)
    apply_hf_context_length(model.hparams, args.context_length)
    _warn_when_header_is_written(model)
    _normalize_chat_templates(model)
    writer = Nvfp4Writer(model, quantizer, args, _collect_bias_names(model))
    writer.install()
    try:
        log.info(
            "llama.cpp prints BF16 for matrices before they are repacked. "
            "Selected matrices are stored as NVFP4."
        )
        model.write()
    except KeyboardInterrupt:
        log.warning("Interrupted. Checkpoints are in %s; rerun the same command to resume.", writer.ckpt_dir)
        raise
    finally:
        writer.restore()

    if not args.keep_checkpoints and writer.quantized:
        shutil.rmtree(writer.ckpt_dir, ignore_errors=True)
    _write_hf_modelfile(model_dir, Path(model.fname_out), args.overwrite, args.mmproj)
    log.info(
        "Wrote %s  (%d NVFP4 tensors, %d with scale folded into the blocks, %d biases adjusted, %d other tensors)",
        model.fname_out,
        writer.quantized,
        writer.baked,
        writer.corrected,
        writer.passthrough,
    )


def self_test(args) -> None:
    """Check that packed blocks dequantize to the same values llama.cpp will load."""
    logging.basicConfig(level=logging.INFO, format="%(message)s")
    from nvfp4gguf import pack_nvfp4, ue4m3_bits
    from nvfp4gguf.quant import check_scale_policy
    check_scale_policy()
    log.info("scale policy ok")
    gguf, *_rest = load_gguf_modules(args.llama_cpp)
    _converter_cls, unpack_uint4 = load_learned_converter(args.convert_to_quant)

    codes = np.zeros((1, 64), dtype=np.uint8)
    codes[0, :16] = 7  # +6
    codes[0, 16:32] = 15  # -6
    scale_float = np.array([[1.0, 0.5, 256.0, 448.0]], dtype=np.float32)
    bits = ue4m3_bits(scale_float)
    raw = pack_nvfp4(codes, bits)
    decoded = gguf.quants.dequantize(raw, gguf.GGMLQuantizationType.NVFP4)
    expected = np.zeros((1, 64), dtype=np.float32)
    expected[0, :16] = 6.0 * scale_float[0, 0]
    expected[0, 16:32] = -6.0 * scale_float[0, 1]
    if not np.allclose(decoded, expected, rtol=1e-5, atol=1e-5):
        raise AssertionError(f"manual NVFP4 roundtrip failed\n{decoded}\n{expected}")
    log.info("manual block roundtrip ok")

    def check(quantizer, label: str, shapes: tuple[tuple[int, int], ...]) -> None:
        generator = np.random.default_rng(0)
        for shape in shapes:
            weight = generator.normal(size=shape).astype(np.float32) * 0.05
            codes, block_scale, per_tensor, dequant = quantizer.quantize_matrix(weight)
            runtime = quantizer.runtime_dequant(pack_nvfp4(codes, ue4m3_bits(block_scale)), per_tensor)
            delta = float(np.max(np.abs(runtime - dequant)))
            if delta > 1e-4:
                raise AssertionError(f"{label} {shape} ggml dequant disagrees with the quantizer by {delta}")
            log.info("%s %s roundtrip ok (max abs %.3e)", label, shape, delta)
            baked_codes, baked_bits = quantizer.bake_global_scale(codes, block_scale, per_tensor)
            baked = quantizer.runtime_dequant(pack_nvfp4(baked_codes, baked_bits), None)
            rel = float(np.max(np.abs(baked - dequant))) / (float(np.max(np.abs(dequant))) + 1e-8)
            if rel > 0.15:
                raise AssertionError(f"{label} {shape} baked scale moved the reconstruction by {rel}")
            log.info("%s %s baked-scale relative error %.3e", label, shape, rel)

    def configured(**overrides):
        values = vars(args).copy()
        values.update(overrides)
        return argparse.Namespace(**values)

    check(TensorQuantizer(configured(simple=True, num_iter=1, device="cpu", optimizer="adamw"), gguf, unpack_uint4),
          "simple", ((64, 128), (40, 64)))
    check(TensorQuantizer(configured(
        simple=False, num_iter=4, device="cpu", optimizer="original", scale_optimization="fixed",
    ), gguf, unpack_uint4), "learned", ((64, 128),))
    log.info("self-test passed")


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Quantize a Hugging Face model to NVFP4 GGUF with convert_to_quant's learned E2M1 rounding.",
    )
    parser.add_argument("model", nargs="?", help="Hugging Face model directory (config.json + safetensors)")
    parser.add_argument("-o", "--output", help="Output directory. Default: ./<model-name>-NVFP4, or ./<model-name> with --no-quant and auto. A path ending in .gguf is the GGUF file itself.")
    parser.add_argument("--modelfile-only", action="store_true", help="Write the Modelfile and skip quantization and GGUF creation")
    parser.add_argument("--overwrite", action="store_true", help="Replace an existing GGUF, Modelfile, or copied projector")
    parser.add_argument("--mmproj", action="append", type=Path, default=None,
                        help="Vision projector GGUF. Copied into the output directory under its own name and added as another FROM line. Repeat for more than one projector.")
    parser.add_argument("--context-length", type=int, default=None,
                        help="Max context stored in the GGUF. Replaces the checkpoint value, or sets one when the checkpoint has no single context length.")
    parser.add_argument("--no-quant", action="store_true",
                        help="Skip learned rounding. U8 tensors are repacked as NVFP4 only when their shape is packed 4-bit. FP8 weights are stored as float.")
    parser.add_argument("--outtype", default="auto", type=_outtype, choices=("auto", "f32", "f16", "bf16", "nvfp4"),
                        help="Float type for tensors that are not packed 4-bit. nvfp4 repacks a checkpoint that already has packed 4-bit weights. auto matches the source float weights (default: auto)")
    parser.add_argument("--llama-cpp", type=Path, default=DEFAULT_LLAMA_CPP,
                        help="llama.cpp checkout used for architecture conversion and the GGUF writer (default: ~/llama.cpp)")
    parser.add_argument("--convert-to-quant", type=Path, default=None,
                        help="convert_to_quant checkout, if it is not installed in the environment")
    parser.add_argument("--model-name", default=None, help="Override the model name stored in the GGUF")
    parser.add_argument("--device", default="auto", help="auto, cpu, cuda, or cuda:N")
    parser.add_argument("--simple", action="store_true", help="Round each value onto the E2M1 grid without learned rounding")
    parser.add_argument("--optimizer", default="auto", choices=("auto", "prodigy", "adamw", "radam", "original"),
                        help="Learned-rounding optimizer. auto uses prodigy when it is installed, otherwise AdamW")
    parser.add_argument("--num-iter", type=int, default=CTQ_NUM_ITER, help="Learned-rounding iterations per matrix (convert_to_quant default: 4000)")
    parser.add_argument("--lr", type=float, default=1.0, help="Initial learning rate")
    parser.add_argument("--lr-min", type=float, default=1e-8, help="Lowest learning rate. The flat-loss check starts here (default: 1e-8).")
    parser.add_argument("--lr-schedule", default="plateau", choices=("plateau", "exponential", "adaptive"))
    parser.add_argument(
        "--scale-optimization", default=DEFAULT_SCALE_OPTIMIZATION, choices=("fixed", "iterative", "joint"),
        help="fixed keeps the initial scale. iterative refreshes it on the original optimizer. "
             "joint trains the block scales with the codes (default: joint)",
    )
    parser.add_argument(
        "--scale-refinement-rounds", type=int, default=DEFAULT_SCALE_REFINEMENT_ROUNDS,
        help="Scale refreshes when --scale-optimization is iterative (default: 2). Ignored for fixed and joint.",
    )
    parser.add_argument("--full-matrix", action="store_true", help="Use a full SVD instead of a low-rank SVD")
    parser.add_argument("--top-p", type=float, default=0.2)
    parser.add_argument("--min-k", type=int, default=128)
    parser.add_argument("--max-k", type=int, default=1280)
    parser.add_argument("--early-stop-loss", type=float, default=5e-9)
    parser.add_argument(
        "--early-stop-stall", type=int, default=DEFAULT_EARLY_STOP_STALL,
        help="Stop after this many steps without a new best, at any learning rate (default: 2000)",
    )
    parser.add_argument(
        "--target-loss", type=float, default=DEFAULT_TARGET_LOSS,
        help="Stop when the best loss reaches this value, at any learning rate (default: 1e-5). 0 disables it.",
    )
    parser.add_argument(
        "--early-stop-floor-stall", type=int, default=DEFAULT_EARLY_STOP_FLOOR_STALL,
        help="If the loss never reaches --target-loss, stop after this many steps without a new best once it has otherwise flattened. 0 disables it (default: 64)",
    )
    parser.add_argument("--calib-samples", type=int, default=3072, help="Random inputs used to correct a sibling bias")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--no-bias-correction", action="store_true")
    parser.add_argument("--quantize-output", action="store_true", help="Also quantize the output (lm_head) matrix")
    parser.add_argument("--exclude", default=None, help="Regex of GGUF tensor names to leave in BF16")
    parser.add_argument("--checkpoint-dir", type=Path, default=None, help="Per-tensor resume directory")
    parser.add_argument("--fresh", action="store_true", help="Delete the checkpoint directory before starting")
    parser.add_argument("--keep-checkpoints", action="store_true")
    parser.add_argument("--eager", action="store_true", help="Load every tensor up front (uses more RAM)")
    parser.add_argument("--no-temp-file", action="store_true", help="Hold the finished GGUF tensors in RAM instead of a temp file")
    parser.add_argument("--verbose", action="store_true")
    parser.add_argument("--self-test", action="store_true", help="Check NVFP4 packing against llama.cpp's dequantizer and exit")
    args = parser.parse_args(argv)
    if args.self_test:
        if args.device == "auto":
            args.device = "cpu"
        return args
    if not args.model:
        parser.error("model directory is required")
    _assign_output(args)
    _refuse_existing_output(args)
    args.exclude_re = re.compile(args.exclude) if args.exclude else None
    if args.convert_to_quant is None and DEFAULT_CONVERT_TO_QUANT.is_dir():
        args.convert_to_quant = DEFAULT_CONVERT_TO_QUANT
    return args


def main(argv: list[str] | None = None) -> None:
    args = parse_args(argv)
    if args.self_test:
        self_test(args)
        return
    export_nvfp4(args)


if __name__ == "__main__":
    main()