#!/usr/bin/env python3
"""Quantize a GGUF to llama.cpp NVFP4, or repack one that is already NVFP4.

The learned-rounding flags match hf_to_nvfp4_gguf.py. Vocabulary, tensor names,
and architecture metadata are copied from the source GGUF. --outtype nvfp4
copies an existing NVFP4 file without fitting it again. --no-quant with
f32, f16, bf16, or auto rewrites the tensors as float, or leaves them as stored.
"""

from __future__ import annotations

import argparse
import hashlib
import logging
import re
import shutil
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,
    import_gguf,
    load_learned_converter,
    must_bake_nvfp4_scale,
    mtp_layer_tensor,
    bake_packed_nvfp4,
    preserve_scale_tensor,
    quantize_and_pack,
    report_reconstruction,
    require_quant_settings,
    resolve_optimizer,
    scale_bake_notice,
)
from nvfp4gguf.bundle import (
    flatten_chat_template,
    normalize_chat_template,
    is_model_context_key,
    outtype as _outtype,
    refuse_existing_bundle,
    require_context_length,
    scalar_context_length,
    warn_missing_gguf_metadata,
    write_simple_modelfile,
)

log = logging.getLogger("gguf_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 _pythonize(value):
    if isinstance(value, np.generic):
        return value.item()
    if isinstance(value, np.ndarray):
        return value.tolist()
    if isinstance(value, list):
        return [_pythonize(item) for item in value]
    return value


def _text(value) -> str:
    value = _pythonize(value)
    if isinstance(value, bytes):
        return value.decode("utf-8", "replace")
    if isinstance(value, list):
        return "".join(_text(item) for item in value)
    return str(value)


def _logical_shape(tensor) -> tuple[int, ...]:
    return tuple(int(dim) for dim in tensor.shape[::-1])


def _to_float(tensor, gguf_module) -> np.ndarray:
    qtype = tensor.tensor_type
    if qtype == gguf_module.GGMLQuantizationType.F32:
        values = np.asarray(tensor.data, dtype=np.float32)
    elif qtype == gguf_module.GGMLQuantizationType.F16:
        values = np.asarray(tensor.data, dtype=np.float32)
    else:
        values = np.asarray(gguf_module.quants.dequantize(tensor.data, qtype), dtype=np.float32)
    # The reader maps the file read-only. Torch refuses a non-writable view.
    return np.array(values, dtype=np.float32, copy=True)


def _open_reader(gguf, path: Path):
    if not path.is_file():
        raise SystemExit(f"GGUF file does not exist: {path}")
    reader = gguf.GGUFReader(path)
    if reader.byte_order != "I":
        raise SystemExit(f"{path} is not little-endian GGUF")
    return reader


def _architecture(reader) -> str:
    field = reader.fields.get("general.architecture")
    if field is None:
        return "unknown"
    return _text(field.contents())


def _file_type_name(reader, gguf_module) -> str:
    field = reader.fields.get("general.file_type")
    if field is None:
        return ""
    try:
        return gguf_module.LlamaFileType(int(field.contents())).name
    except (TypeError, ValueError):
        return str(field.contents())


def _is_nvfp4(reader, gguf_module) -> bool:
    if _file_type_name(reader, gguf_module) == "MOSTLY_NVFP4":
        return True
    return any(tensor.tensor_type == gguf_module.GGMLQuantizationType.NVFP4 for tensor in reader.tensors)


def _copy_metadata(reader, writer, gguf_module, file_type: int | None, model_name: str | None, context_length: int | None) -> None:
    requested = require_context_length(context_length)
    source_context = None
    wrote_context = False
    for field in reader.fields.values():
        if field.name.startswith("GGUF.") or field.name == "general.architecture" or not field.types:
            continue
        value = _pythonize(field.contents())
        if value is None:
            continue
        value_type = field.types[0]
        if value_type == gguf_module.GGUFValueType.ARRAY and not value:
            continue
        if field.name == "general.file_type" and file_type is not None:
            value = int(file_type)
        if field.name == "general.name" and model_name:
            value = model_name
        if field.name == "general.alignment":
            writer.data_alignment = int(value)
        subtype = field.types[-1] if value_type == gguf_module.GGUFValueType.ARRAY else None
        if field.name == "tokenizer.chat_template" or field.name.startswith("tokenizer.chat_template."):
            value = normalize_chat_template(value)
        if is_model_context_key(field.name):
            source_context = scalar_context_length(value)
            if requested is not None:
                value = requested
                value_type = gguf_module.GGUFValueType.UINT32
                subtype = None
                wrote_context = True
        writer.add_key_value(field.name, value, value_type, sub_type=subtype)
    if file_type is not None and "general.file_type" not in reader.fields:
        writer.add_file_type(int(file_type))
    if model_name and "general.name" not in reader.fields:
        writer.add_name(model_name)
    if requested is not None and not wrote_context:
        writer.add_context_length(requested)
    if requested is not None:
        log.info("context length %d", requested)
    elif source_context is not None:
        log.info("context length %d from the source GGUF", source_context)


def _replaced_scale(name: str, quant_names: set[str], architecture: str) -> bool:
    if not (name.endswith(".scale") or name.endswith(".input_scale")):
        return False
    stem = name.rsplit(".", 1)[0]
    weight = stem + ".weight"
    if weight not in quant_names:
        return False
    return not (name.endswith(".scale") and preserve_scale_tensor(architecture, weight))


def _keep_f32(name: str, rank: int) -> bool:
    return rank < 2 or "norm" in name.lower()


def _add_float(writer, gguf_module, name: str, values: np.ndarray, outtype: str) -> None:
    values = np.ascontiguousarray(values, dtype=np.float32)
    if outtype == "f16":
        writer.add_tensor(name, values.astype(np.float16))
        return
    if outtype == "bf16":
        raw = gguf_module.quants.quantize(values, gguf_module.GGMLQuantizationType.BF16)
        writer.add_tensor(name, raw, raw_dtype=gguf_module.GGMLQuantizationType.BF16)
        return
    writer.add_tensor(name, values)


def _copy_tensor(writer, tensor) -> None:
    writer.add_tensor(tensor.name, tensor.data, raw_dtype=tensor.tensor_type)


def _load_checkpoint(directory: Path, name: str, shape: tuple[int, ...]):
    path = _checkpoint_path(directory, 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(shape):
            return None
        raw = np.array(blob["raw"], copy=True)
        scale = np.array(blob["scale"], copy=True)
        bias = np.array(blob["bias"], copy=True)
    return raw, None if scale.size == 0 else scale, None if bias.size == 0 else bias


def _save_checkpoint(directory: Path, name: str, shape: tuple[int, ...], raw, scale, bias) -> None:
    path = _checkpoint_path(directory, name)
    temporary = path.with_name(path.name + ".tmp")
    with temporary.open("wb") as handle:
        np.savez_compressed(
            handle,
            version=np.int32(QUANT_VERSION),
            shape=np.asarray(shape, dtype=np.int64),
            raw=raw,
            scale=np.zeros((0,), np.float32) if scale is None else np.asarray(scale, dtype=np.float32),
            bias=np.zeros((0,), np.float32) if bias is None else np.asarray(bias, dtype=np.float32),
        )
    temporary.replace(path)


def _assign_output(args) -> None:
    source = Path(args.gguf)
    stem = source.stem
    passthrough = args.no_quant or args.outtype == "nvfp4"
    if args.output is None:
        if passthrough and args.outtype == "auto":
            directory = Path.cwd() / stem
            gguf_name = f"{stem}.gguf"
        elif passthrough:
            label = "NVFP4" if args.outtype == "nvfp4" else args.outtype.upper()
            directory = Path.cwd() / f"{stem}-{label}"
            gguf_name = f"{directory.name}.gguf"
        else:
            directory = Path.cwd() / f"{stem}-NVFP4"
            gguf_name = f"{directory.name}.gguf"
    else:
        output = Path(args.output)
        if output.suffix == ".gguf":
            output.parent.mkdir(parents=True, exist_ok=True)
            args.output = str(output.resolve())
            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_same_file(args) -> None:
    if Path(args.output).resolve() == Path(args.gguf).resolve():
        raise SystemExit("Output GGUF is the same file as the input. Choose a different --output.")


def _target_file_type(gguf_module, outtype: str, source_name: str):
    if outtype == "nvfp4":
        return gguf_module.LlamaFileType.MOSTLY_NVFP4
    if outtype == "f32":
        return gguf_module.LlamaFileType.ALL_F32
    if outtype == "f16":
        return gguf_module.LlamaFileType.MOSTLY_F16
    if outtype == "bf16":
        return gguf_module.LlamaFileType.MOSTLY_BF16
    if source_name == "ALL_F32":
        return gguf_module.LlamaFileType.ALL_F32
    if source_name == "MOSTLY_F16":
        return gguf_module.LlamaFileType.MOSTLY_F16
    if source_name == "MOSTLY_BF16":
        return gguf_module.LlamaFileType.MOSTLY_BF16
    if source_name == "MOSTLY_NVFP4":
        return gguf_module.LlamaFileType.MOSTLY_NVFP4
    return None


def _chat_template(reader) -> str:
    parts = []
    for name, field in reader.fields.items():
        if name == "tokenizer.chat_template" or name.startswith("tokenizer.chat_template."):
            parts.append(flatten_chat_template(field.contents()))
    return "\n".join(part for part in parts if part)


def _kv_int(reader, key: str) -> int:
    field = reader.fields.get(key)
    if field is None:
        return 0
    try:
        return int(field.contents())
    except (TypeError, ValueError):
        return 0


def _mtp_counts(reader, architecture: str) -> tuple[int, int]:
    return _kv_int(reader, f"{architecture}.block_count"), _kv_int(reader, f"{architecture}.nextn_predict_layers")


def _folded_scales(reader, gguf_module, architecture: str, n_blocks: int, n_nextn: int) -> dict[str, float]:
    """Per-tensor NVFP4 scales the llama.cpp graph never multiplies."""
    by_name = {tensor.name: tensor for tensor in reader.tensors}
    nvfp4 = gguf_module.GGMLQuantizationType.NVFP4
    scales: dict[str, float] = {}
    for tensor in reader.tensors:
        if not tensor.name.endswith(".scale"):
            continue
        weight_name = tensor.name[: -len(".scale")] + ".weight"
        weight = by_name.get(weight_name)
        if weight is None or weight.tensor_type != nvfp4:
            continue
        if preserve_scale_tensor(architecture, weight_name):
            continue
        bias_name = weight_name[: -len(".weight")] + ".bias"
        fold = mtp_layer_tensor(weight_name, n_blocks, n_nextn) or must_bake_nvfp4_scale(
            architecture, weight_name, bias_name in by_name
        )
        if not fold:
            continue
        values = np.asarray(tensor.data, dtype=np.float32).reshape(-1)
        if values.size != 1:
            log.warning(
                "Leaving %s in place. Baking expects one per-tensor scale, got %d values.",
                tensor.name,
                int(values.size),
            )
            continue
        scales[tensor.name] = float(values[0])
    return scales


def export_copy(args) -> None:
    """Rewrite the GGUF without learned rounding."""
    gguf = import_gguf(args.llama_cpp)
    source = Path(args.gguf).resolve()
    output = Path(args.output).resolve()
    reader = _open_reader(gguf, source)
    if args.outtype == "nvfp4" and not _is_nvfp4(reader, gguf):
        raise SystemExit(
            f"{source} is not an NVFP4 GGUF. "
            "Use --outtype f16, bf16, f32, or auto to write float weights, "
            "or omit --outtype nvfp4 to fit NVFP4 from these weights."
        )
    architecture = _architecture(reader)
    source_type = _file_type_name(reader, gguf)
    file_type = _target_file_type(gguf, args.outtype, source_type)
    if args.outtype == "nvfp4":
        log.info("Copying existing NVFP4 tensors. Learned rounding is skipped.")
    elif args.outtype == "auto":
        log.info("Copying tensors as stored (%s). Learned rounding is skipped.", source_type or "unknown file type")
    else:
        log.info("Writing float GGUF (%s). Learned rounding is skipped.", args.outtype)
    writer = gguf.GGUFWriter(output, architecture, use_temp_file=not args.no_temp_file)
    _copy_metadata(reader, writer, gguf, None if file_type is None else int(file_type), args.model_name, args.context_length)
    n_blocks, n_nextn = _mtp_counts(reader, architecture)
    folded = _folded_scales(reader, gguf, architecture, n_blocks, n_nextn) if args.outtype in ("auto", "nvfp4") else {}
    if folded:
        log.info(
            "Folding %d scale tensors into their NVFP4 blocks. "
            "llama.cpp does not claim those sidecars.",
            len(folded),
        )
    for tensor in reader.tensors:
        if tensor.name in folded:
            continue
        scale_name = tensor.name[: -len(".weight")] + ".scale" if tensor.name.endswith(".weight") else ""
        if scale_name in folded and tensor.tensor_type == gguf.GGMLQuantizationType.NVFP4:
            baked = bake_packed_nvfp4(np.array(tensor.data, copy=True), folded[scale_name])
            writer.add_tensor(tensor.name, baked, raw_dtype=gguf.GGMLQuantizationType.NVFP4)
            continue
        if args.outtype in ("auto", "nvfp4"):
            _copy_tensor(writer, tensor)
            continue
        if _keep_f32(tensor.name, len(_logical_shape(tensor))):
            if tensor.tensor_type == gguf.GGMLQuantizationType.F32:
                _copy_tensor(writer, tensor)
            else:
                _add_float(writer, gguf, tensor.name, _to_float(tensor, gguf), "f32")
            continue
        _add_float(writer, gguf, tensor.name, _to_float(tensor, gguf), args.outtype)
    log.info("Writing %s", output)
    warn_missing_gguf_metadata(writer)
    writer.write_header_to_file()
    writer.write_kv_data_to_file()
    writer.write_tensors_to_file(progress=True)
    writer.close()
    write_simple_modelfile(output, args.overwrite, args.mmproj, _chat_template(reader))
    log.info("Wrote %s", output)


def _emit_bias(writer, gguf_module, tensor, correction: np.ndarray | None) -> bool:
    if correction is None:
        _copy_tensor(writer, tensor)
        return False
    values = _to_float(tensor, gguf_module)
    flat = values.reshape(-1)
    if flat.shape != correction.shape:
        log.warning("    bias %s shape %s does not match correction %s", tensor.name, values.shape, correction.shape)
        _copy_tensor(writer, tensor)
        return False
    values = (flat - correction).reshape(values.shape)
    log.info("    bias %s adjusted (mean %+.3e)", tensor.name, float(correction.mean()))
    if tensor.tensor_type == gguf_module.GGMLQuantizationType.F16:
        writer.add_tensor(tensor.name, np.ascontiguousarray(values, dtype=np.float32).astype(np.float16))
        return True
    if tensor.tensor_type == gguf_module.GGMLQuantizationType.BF16:
        raw = gguf_module.quants.quantize(
            np.ascontiguousarray(values, dtype=np.float32),
            gguf_module.GGMLQuantizationType.BF16,
        )
        writer.add_tensor(tensor.name, raw, raw_dtype=gguf_module.GGMLQuantizationType.BF16)
        return True
    writer.add_tensor(tensor.name, np.ascontiguousarray(values, dtype=np.float32))
    return True


def export_quant(args) -> None:
    gguf = import_gguf(args.llama_cpp)
    source = Path(args.gguf).resolve()
    output = Path(args.output).resolve()
    ensure_cuda_interpreter(args.device)
    reader = _open_reader(gguf, source)
    architecture = _architecture(reader)
    bind_exclude(args, architecture)
    notice = scale_bake_notice(architecture)
    if notice:
        log.warning(notice)
    file_type_name = _file_type_name(reader, gguf)
    if file_type_name not in {"ALL_F32", "MOSTLY_F16", "MOSTLY_BF16"}:
        log.warning(
            "Source file type is %s. NVFP4 will be fit to the dequantized weights, "
            "so detail already discarded by the source quant stays discarded.",
            file_type_name or "quantized",
        )
    names = {tensor.name for tensor in reader.tensors}
    n_blocks, n_nextn = _mtp_counts(reader, architecture)
    quant_names = [
        tensor.name
        for tensor in reader.tensors
        if eligible_weight(tensor.name, _logical_shape(tensor), args, gguf)
    ]
    quant_name_set = set(quant_names)
    log.info("%d matrices queued for NVFP4, %d tensors copied", len(quant_names), len(reader.tensors) - len(quant_names))
    if not quant_names:
        log.warning("No matrices are eligible for NVFP4. Tensors are copied and the file type stays unchanged.")

    work = Path(args.checkpoint_dir) if args.checkpoint_dir else Path(str(output) + ".nvfp4-work")
    if args.fresh and work.exists():
        shutil.rmtree(work)
    work.mkdir(parents=True, exist_ok=True)
    require_quant_settings(work, args)

    _converter_cls, unpack_uint4 = load_learned_converter(args.convert_to_quant)
    del _converter_cls
    quantizer = TensorQuantizer(args, gguf, unpack_uint4)
    writer = gguf.GGUFWriter(output, architecture, use_temp_file=not args.no_temp_file)
    _copy_metadata(
        reader, writer, gguf,
        int(gguf.LlamaFileType.MOSTLY_NVFP4) if quant_names else None,
        args.model_name,
        args.context_length,
    )

    pending: dict[str, object] = {}
    corrections: dict[str, np.ndarray] = {}
    done = 0
    baked = 0
    corrected = 0
    try:
        for tensor in reader.tensors:
            if _replaced_scale(tensor.name, quant_name_set, architecture):
                continue
            weight_name = tensor.name[: -len(".bias")] + ".weight" if tensor.name.endswith(".bias") else ""
            if weight_name in quant_name_set and not args.no_bias_correction:
                if tensor.name not in corrections:
                    pending[tensor.name] = tensor
                    continue
                if _emit_bias(writer, gguf, tensor, corrections[tensor.name]):
                    corrected += 1
                continue
            if tensor.name not in quant_name_set:
                _copy_tensor(writer, tensor)
                continue
            done += 1
            logical = _logical_shape(tensor)
            bias_name = tensor.name[: -len(".weight")] + ".bias"
            bake = must_bake_nvfp4_scale(architecture, tensor.name, bias_name in names) or mtp_layer_tensor(
                tensor.name, n_blocks, n_nextn
            )
            cached = _load_checkpoint(work, tensor.name, logical)
            if cached is not None and (cached[1] is None) != bake:
                log.info("    checkpoint scale layout does not match %s; recomputing", tensor.name)
                cached = None
            log.info(
                "NVFP4 [%d/%d] %s  %s  from %s%s",
                done, len(quant_names), tensor.name, logical, tensor.tensor_type.name,
                "  (checkpoint)" if cached else "",
            )
            if cached is None:
                weight = _to_float(tensor, gguf)
                raw, scale = quantize_and_pack(quantizer, weight, bake, tensor.name)
                report_reconstruction(quantizer, weight, raw, scale)
                bias = None
                if weight.ndim == 2 and not args.no_bias_correction and bias_name in names:
                    per_tensor = None if scale is None else float(np.reshape(scale, -1)[0])
                    runtime = quantizer.runtime_dequant(raw if raw.ndim == 2 else raw[0], per_tensor)
                    sample = weight if weight.ndim == 2 else weight[0]
                    bias = quantizer.bias_correction(sample, runtime)
                _save_checkpoint(work, tensor.name, logical, raw, scale, bias)
                del weight
            else:
                raw, scale, bias = cached
            if bias is not None:
                corrections[bias_name] = bias
                pending_tensor = pending.pop(bias_name, None)
                if pending_tensor is not None and _emit_bias(writer, gguf, pending_tensor, bias):
                    corrected += 1
            if bake and scale is None:
                baked += 1
            writer.add_tensor(tensor.name, raw, raw_dtype=gguf.GGMLQuantizationType.NVFP4)
            if scale is not None:
                writer.add_tensor(tensor.name[: -len(".weight")] + ".scale", np.ascontiguousarray(scale, dtype=np.float32))
            del raw
        for name, tensor in list(pending.items()):
            if _emit_bias(writer, gguf, tensor, corrections.get(name)):
                corrected += 1
        pending.clear()
    except KeyboardInterrupt:
        log.warning("Interrupted. Checkpoints are in %s; rerun the same command to resume.", work)
        raise

    log.info("Writing %s", output)
    warn_missing_gguf_metadata(writer)
    writer.write_header_to_file()
    writer.write_kv_data_to_file()
    writer.write_tensors_to_file(progress=True)
    writer.close()
    if not args.keep_checkpoints:
        shutil.rmtree(work, ignore_errors=True)
    write_simple_modelfile(output, args.overwrite, args.mmproj, _chat_template(reader))
    log.info(
        "Wrote %s  (%d NVFP4 tensors, %d with scale folded into the blocks, %d biases adjusted)",
        output, done, baked, corrected,
    )


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Quantize a GGUF to NVFP4, or repack one that is already NVFP4, with the same options as hf_to_nvfp4_gguf.py.",
    )
    parser.add_argument("gguf", nargs="?", help="Input GGUF")
    parser.add_argument("-o", "--output", help="Output directory. Default: ./<name>-NVFP4, or ./<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 a source value, or sets one when the source has no single context length.")
    parser.add_argument("--no-quant", action="store_true",
                        help="Skip learned rounding. Writes a float GGUF, or copies an existing NVFP4 GGUF with --outtype nvfp4")
    parser.add_argument("--outtype", default="auto", type=_outtype, choices=("auto", "f32", "f16", "bf16", "nvfp4"),
                        help="Type used when learned rounding is skipped. nvfp4 copies a GGUF that is already NVFP4. auto leaves tensors as stored (default: auto)")
    parser.add_argument("--llama-cpp", type=Path, default=DEFAULT_LLAMA_CPP,
                        help="llama.cpp checkout used for the GGUF reader and 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 their source type")
    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="Accepted for the same command line as hf_to_nvfp4_gguf. The GGUF reader maps the file either way.")
    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.gguf:
        parser.error("GGUF file is required")
    _assign_output(args)
    _refuse_same_file(args)
    if not args.modelfile_only:
        refuse_existing_bundle(Path(args.output), args.overwrite, args.mmproj)
    else:
        refuse_existing_bundle(Path(args.output), args.overwrite, args.mmproj, skip_gguf=True)
    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:
        from hf_to_nvfp4_gguf import self_test
        self_test(args)
        return
    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(message)s")
    if args.modelfile_only:
        output = Path(args.output)
        output.parent.mkdir(parents=True, exist_ok=True)
        log.info("Skipping quantization and GGUF creation; writing Modelfile only")
        reader = _open_reader(import_gguf(args.llama_cpp), Path(args.gguf))
        write_simple_modelfile(output, args.overwrite, args.mmproj, _chat_template(reader))
        return
    if args.no_quant or args.outtype == "nvfp4":
        export_copy(args)
        return
    args.optimizer = resolve_optimizer(args.optimizer)
    export_quant(args)


if __name__ == "__main__":
    main()