#!/usr/bin/env python3
"""Quantize a Hugging Face directory to one ComfyUI NVFP4 safetensors file.

convert_to_quant reads one safetensors file and writes one file with packed
NVFP4 weights, ``weight_scale`` / ``weight_scale_2``, ``.comfy_quant``
tensors, and a ``_quantization_metadata`` header. A checkpoint is usually
sharded. This script runs that same conversion on every shard, then merges
the results into a single ``.safetensors`` file.

A bias is corrected only when it lives in the same shard as its weight,
which is how convert_to_quant sees it. Norms, embeddings, ``lm_head``, and
the other names on convert_to_quant's exclusion list stay in their original
dtype.
"""

from __future__ import annotations

import argparse
import json
import shutil
import struct
from pathlib import Path

from nvfp4gguf import (
    CTQ_NUM_ITER,
    DEFAULT_CONVERT_TO_QUANT,
    DEFAULT_EARLY_STOP_STALL,
    DEFAULT_SCALE_OPTIMIZATION,
    DEFAULT_SCALE_REFINEMENT_ROUNDS,
    load_learned_converter,
    quant_profile,
    resolve_optimizer,
)

INDEX_NAME = "model.safetensors.index.json"
# safetensors aligns each tensor to its dtype width, not to 8.
_DTYPE_ALIGN = {
    "BOOL": 1, "U8": 1, "I8": 1,
    "F8_E4M3": 1, "F8_E5M2": 1, "F8_E8M0": 1,
    "F16": 2, "BF16": 2, "I16": 2, "U16": 2,
    "F32": 4, "I32": 4, "U32": 4,
    "F64": 8, "I64": 8, "U64": 8,
}


def discover_shards(model_dir: Path) -> list[str]:
    """Return safetensors shard filenames in the checkpoint directory."""
    index_path = model_dir / INDEX_NAME
    if index_path.is_file():
        payload = json.loads(index_path.read_text(encoding="utf-8"))
        weight_map = payload.get("weight_map")
        if not isinstance(weight_map, dict) or not weight_map:
            raise SystemExit(f"{index_path} has no weight_map")
        names = sorted(set(weight_map.values()))
        missing = [name for name in names if not (model_dir / name).is_file()]
        if missing:
            raise SystemExit(f"Index entries are missing from {model_dir}: {', '.join(missing)}")
        return names

    single = model_dir / "model.safetensors"
    if single.is_file():
        return ["model.safetensors"]

    shards = sorted(path.name for path in model_dir.glob("*.safetensors") if path.is_file())
    if not shards:
        raise SystemExit(f"No safetensors shards in {model_dir}")
    return shards


def output_file(args, source: Path) -> Path:
    if not args.output:
        return (Path.cwd() / f"{source.name}-NVFP4.safetensors").resolve()
    output = Path(args.output)
    if output.exists() and output.is_dir():
        return (output / f"{source.name}-NVFP4.safetensors").resolve()
    if output.suffix != ".safetensors":
        output = output.with_name(output.name + ".safetensors")
    return output.resolve()


def read_header(path: Path) -> tuple[dict, int]:
    with path.open("rb") as handle:
        raw_size = handle.read(8)
        if len(raw_size) != 8:
            raise SystemExit(f"{path} is not a safetensors file")
        header_size = struct.unpack("<Q", raw_size)[0]
        header = json.loads(handle.read(header_size))
    if not isinstance(header, dict):
        raise SystemExit(f"{path} has an unreadable safetensors header")
    return header, header_size


def merge_safetensors(paths: list[Path], dest: Path) -> None:
    """Concatenate quantized shards into one safetensors file.

    Tensor bytes are copied, not loaded. ``_quantization_metadata`` layers
    from each shard are combined so ComfyUI still sees every quantized weight.
    """
    metadata: dict[str, str] = {}
    layers: dict = {}
    pieces: list[tuple[str, dict, Path, int, int]] = []
    owners: dict[str, str] = {}
    for path in paths:
        header, header_size = read_header(path)
        file_meta = dict(header.pop("__metadata__", {}) or {})
        raw_quant = file_meta.pop("_quantization_metadata", None)
        if raw_quant:
            parsed = json.loads(raw_quant)
            file_layers = parsed.get("layers", {})
            overlap = sorted(set(file_layers) & set(layers))
            if overlap:
                raise SystemExit(f"Quant metadata for {overlap[0]!r} appears in more than one shard")
            layers.update(file_layers)
        for key, value in file_meta.items():
            metadata.setdefault(key, value)
        data_base = 8 + header_size
        for key, entry in header.items():
            if key in owners:
                raise SystemExit(f"Tensor {key!r} appears in both {owners[key]} and {path.name}")
            owners[key] = path.name
            start, end = entry["data_offsets"]
            pieces.append((key, entry, path, data_base + int(start), int(end) - int(start)))

    if layers:
        metadata["_quantization_metadata"] = json.dumps(
            {"format_version": "1.0", "layers": layers},
            separators=(",", ":"),
        )

    # safetensors rejects gaps. Grouping by descending dtype width keeps every
    # tensor aligned without padding, which is how save_file lays a file out.
    ordered = sorted(
        enumerate(pieces),
        key=lambda item: (-_DTYPE_ALIGN.get(item[1][1].get("dtype", ""), 8), item[0]),
    )
    cursor = 0
    new_header: dict = {}
    if metadata:
        new_header["__metadata__"] = metadata
    layout: list[tuple[Path, int, int]] = []
    for _index, (key, entry, path, abs_start, length) in ordered:
        align = _DTYPE_ALIGN.get(entry.get("dtype", ""), 8)
        if cursor % align != 0:
            raise SystemExit(f"Cannot align {key} ({entry.get('dtype')}) at offset {cursor}")
        info = {name: value for name, value in entry.items() if name != "data_offsets"}
        info["data_offsets"] = [cursor, cursor + length]
        new_header[key] = info
        layout.append((path, abs_start, length))
        cursor += length

    header_bytes = json.dumps(new_header, separators=(",", ":")).encode("utf-8")
    header_bytes += b" " * ((8 - (len(header_bytes) % 8)) % 8)
    dest.parent.mkdir(parents=True, exist_ok=True)
    temporary = dest.with_name(dest.name + ".tmp")
    with temporary.open("wb") as out:
        out.write(struct.pack("<Q", len(header_bytes)))
        out.write(header_bytes)
        for path, abs_start, length in layout:
            with path.open("rb") as src:
                src.seek(abs_start)
                remaining = length
                while remaining:
                    chunk = src.read(min(remaining, 8 * 1024 * 1024))
                    if not chunk:
                        raise SystemExit(f"Short read from {path}")
                    out.write(chunk)
                    remaining -= len(chunk)
    temporary.replace(dest)


def quantize_shard(convert_to_nvfp4, args, optimizer: str, source: Path, destination: Path) -> None:
    if destination.exists():
        destination.unlink()
    convert_to_nvfp4(
        str(source),
        str(destination),
        exclude_layers=args.exclude_layers,
        output_dtype=args.output_dtype,
        simple=args.simple,
        num_iter=args.num_iter,
        heur=args.heur,
        calib_samples=args.calib_samples,
        seed=args.seed,
        optimizer=optimizer,
        lr=args.lr,
        lr_schedule=args.lr_schedule,
        top_p=args.top_p,
        min_k=args.min_k,
        max_k=args.max_k,
        full_matrix=args.full_matrix,
        early_stop_loss=args.early_stop_loss,
        early_stop_stall=args.early_stop_stall,
        scale_refinement_rounds=args.scale_refinement_rounds,
        scale_optimization=args.scale_optimization,
        low_memory=args.low_memory,
        use_speed=args.use_speed,
        verbose_output=args.verbose,
    )
    if not destination.is_file() or destination.stat().st_size == 0:
        raise SystemExit(f"convert_to_quant did not write {destination}")


def _apply_directory_profile(args, source: Path) -> None:
    config_path = source / "config.json"
    if not config_path.is_file():
        return
    try:
        config = json.loads(config_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return
    names: list[str] = []
    architectures = config.get("architectures")
    if isinstance(architectures, list):
        names.extend(str(item) for item in architectures)
    elif isinstance(architectures, str):
        names.append(architectures)
    model_type = config.get("model_type")
    if isinstance(model_type, str) and model_type:
        names.append(model_type)
    for name in names:
        profile = quant_profile(name)
        if profile is None:
            continue
        if not args.exclude_layers:
            args.exclude_layers = profile.skip
            print(f"{profile.arch} profile: leaving {profile.label} unquantized.")
        else:
            print(f"{profile.arch} profile: --exclude-layers replaces the default skip of {profile.label}.")
        return


def convert_directory(args) -> None:
    source = Path(args.model).resolve()
    _apply_directory_profile(args, source)
    if not source.is_dir():
        raise SystemExit(f"Not a directory: {source}")
    output = output_file(args, source)
    if output.is_relative_to(source):
        raise SystemExit("Write the safetensors file outside the input directory.")
    if output.exists() and not args.overwrite:
        raise SystemExit(f"{output} already exists. Pass --overwrite to replace it.")
    shards = discover_shards(source)

    load_learned_converter(args.convert_to_quant)
    from convert_to_quant.formats.nvfp4_conversion import convert_to_nvfp4
    from convert_to_quant.utils.logging import setup_logging

    # Without this, convert_to_quant drops the "(13/256) Processing tensor" lines
    # and only the inner 4000-step bar is visible.
    setup_logging("VERBOSE" if args.verbose else "NORMAL")

    optimizer = resolve_optimizer(args.optimizer)
    print(f"Quantizing {len(shards)} shard(s) from {source}")
    print(f"Output file: {output}")
    if len(shards) == 1:
        quantize_shard(convert_to_nvfp4, args, optimizer, source / shards[0], output)
        print(f"Wrote {output}")
        return

    work = output.with_name(output.name + ".shards")
    if work.exists():
        shutil.rmtree(work)
    work.mkdir(parents=True)
    try:
        written: list[Path] = []
        for index, name in enumerate(shards, start=1):
            destination = work / name
            print(f"\n[{index}/{len(shards)}] {name}")
            quantize_shard(convert_to_nvfp4, args, optimizer, source / name, destination)
            written.append(destination)
        print(f"\nMerging {len(written)} shards into {output.name}")
        merge_safetensors(written, output)
    finally:
        shutil.rmtree(work, ignore_errors=True)
    print(f"Wrote {output}")


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Quantize a Hugging Face directory to one ComfyUI NVFP4 safetensors file.",
    )
    parser.add_argument("model", nargs="?", help="Hugging Face model directory (config.json + safetensors)")
    parser.add_argument("-o", "--output", help="Output .safetensors file. Default: ./<model-name>-NVFP4.safetensors")
    parser.add_argument("--overwrite", action="store_true", help="Replace an existing output file")
    parser.add_argument("--convert-to-quant", type=Path, default=None,
                        help="convert_to_quant checkout, if it is not installed")
    parser.add_argument("--simple", action="store_true", help="Round onto the E2M1 grid without learned rounding")
    parser.add_argument("--optimizer", default="auto", choices=("auto", "prodigy", "adamw", "radam", "original"),
                        help="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-patience", type=int, default=5, help="Patience for learning rate decay (default: 5)")
    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 progress. At the learning-rate floor, progress means a 1%% loss drop (default: 2000)",
    )
    parser.add_argument("--calib-samples", type=int, default=3072)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--use-speed", action="store_true", help="Enable Prodigy's use_speed")
    parser.add_argument("--low-memory", action="store_true", help="Stream each shard instead of loading it all at once")
    parser.add_argument("--heur", action="store_true", help="Skip layers convert_to_quant's size heuristic rejects")
    parser.add_argument("--exclude-layers", default=None, help="Regex of tensor names to leave unquantized")
    parser.add_argument("--output-dtype", default="bfloat16", choices=("bfloat16", "float16"),
                        help="Dtype recorded for quantized layers and used for unquantized 2D weights")
    parser.add_argument("--verbose", action="store_true")
    args = parser.parse_args(argv)
    if not args.model:
        parser.error("model directory is required")
    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:
    convert_directory(parse_args(argv))


if __name__ == "__main__":
    main()