# SPDX-License-Identifier: Apache-2.0
"""Serialize the MiniMax-H3 Qwen3-VL text encoder as NVFP4.

H3 conditions on hidden state 50 of a 64-layer Qwen3-VL and never predicts a
token, so the conditioner builds 50 decoder layers, no ``lm_head``, and reads
nothing above (#1711). Every linear in those layers is quantized with
``flashinfer.nvfp4_quantize`` exactly as ``convert_model_to_nvfp4`` would at
runtime and stored as ``weight_packed`` / ``weight_scale`` /
``weight_global_scale``; the byte layout is documented in
``fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py``. Everything else
(token embedding, norms, the vision tower) is copied unchanged. ``config.json``
gains the ``quantization_config`` block that makes the loader select the
serialized NVFP4 path and a ``num_hidden_layers_override`` equal to the kept
layer count, so the loader builds exactly the layers the file holds.

For the FastH3 conditioner: 50 layers x 487.6M values = 24.4B values, 12.2 GB
packed plus 1.5 GB of scales, next to about 3 GB of unquantized tensors,
against 48.8 GB of bf16 language layers (63 GB on disk with the 14 dropped
layers and ``lm_head``).

Needs a Blackwell GPU with FlashInfer: the quantizer is a CUDA kernel, and the
error report runs the same ``mm_fp4`` path the loader executes.

Usage::

    python scripts/checkpoint_conversion/convert_minimax_h3_text_encoder_nvfp4.py \
        --src /path/to/FastH3/text_encoder \
        --dst /path/to/FastH3-nvfp4/text_encoder --keep-bf16 mlp.down_proj

    # use it with any FastH3 model directory
    ComponentConfig(text_encoder_weights="/path/to/FastH3-nvfp4/text_encoder")

``--keep-bf16 mlp.down_proj`` leaves one projection kind in bf16 in every
layer and records it in ``modules_to_not_convert`` so the loader builds those
linears unquantized. The SwiGLU product feeding ``down_proj`` carries the
widest activation outliers in the stack; measured on one GB10 it is the
projection whose 4-bit form moves the layer-50 output most.

Every quantized linear is probed: random bf16 rows go through the same
``mm_fp4`` path the loader runs and the relative error against the bf16
product must stay under ``--max-probe-error``, otherwise the conversion stops
and removes what it wrote. The genuine W4A4 noise on random inputs is about
0.13; a wrong scale layout reads as 1.0. ``--report-only`` runs the probe and
writes nothing. The kept layer count is the one the conditioner builds, 50
for H3, and is also written into ``config.json`` for a directory used as the
``text_encoder`` component; when the checkpoint is supplied through
``text_encoder_weights`` the model structure comes from the model directory,
which builds the same 50 layers.
"""

from __future__ import annotations

import argparse
import json
import re
import shutil
import time
from collections import Counter
from importlib import metadata as importlib_metadata
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import save_file

from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConfig
from fastvideo.layers.quantization.nvfp4_config import _nvfp4_quantize, _require_flashinfer
from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import (
    LANGUAGE_PROJECTIONS,
    NVFP4_TENSOR_SUFFIXES,
    _nvfp4_linear,
    _quantize_activation_nvfp4,
    nvfp4_packed_weight_shape,
    nvfp4_scale_shape,
    nvfp4_weight_global_scale,
    serialized_nvfp4_quantization_config,
    validate_nvfp4_geometry,
)
from fastvideo.models.loader.weight_utils import SAFE_WEIGHTS_INDEX_NAME, resolve_safetensors_files

LANGUAGE_LINEAR = re.compile(r"^model\.language_model\.layers\.(?P<layer>\d+)\.(?P<proj>" +
                             "|".join(re.escape(name) for name in LANGUAGE_PROJECTIONS) + r")\.weight$")
LANGUAGE_LAYER = re.compile(r"^model\.language_model\.layers\.(?P<layer>\d+)\.")

# Keys the converter drops on purpose: id -> (pattern, reason).
SKIPPED = {
    "lm_head": ("lm_head.weight", "the conditioner never predicts tokens; H3 reads a hidden state"),
    "upper_layers": ("model.language_model.layers.{N >= kept}.*",
                     "layers above the one H3 reads are never built (#1711)"),
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--src", type=Path, required=True, help="text_encoder directory of the bf16 checkpoint")
    parser.add_argument("--dst", type=Path, help="text_encoder directory to write (required unless --report-only)")
    parser.add_argument("--keep-bf16", action="append", default=[], metavar="PROJ",
                        help="projection kind to leave in bf16 in every layer, e.g. mlp.down_proj; "
                        "repeat or comma-separate for several")
    parser.add_argument("--device", default="cuda", help="CUDA device that runs the quantizer")
    parser.add_argument("--shard-size-gb", type=float, default=4.0, help="safetensors shard size")
    parser.add_argument("--probe-rows", type=int, default=512,
                        help="rows of random activations per linear for the error probe, 0 disables it")
    parser.add_argument("--max-probe-error", type=float, default=0.5,
                        help="stop when a linear's FP4 GEMM relative error exceeds this; W4A4 noise is about 0.13")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--report-only", action="store_true", help="quantize and report errors, write nothing")
    args = parser.parse_args()
    if not args.report_only and args.dst is None:
        parser.error("--dst is required unless --report-only is given")
    if args.probe_rows < 0:
        parser.error("--probe-rows must be 0 or positive")
    if args.shard_size_gb <= 0:
        parser.error("--shard-size-gb must be positive")
    if args.max_probe_error <= 0:
        parser.error("--max-probe-error must be positive")
    keep = [name.strip() for entry in args.keep_bf16 for name in entry.split(",") if name.strip()]
    unknown = sorted(set(keep) - set(LANGUAGE_PROJECTIONS))
    if unknown:
        parser.error(f"--keep-bf16 got unknown projection(s) {unknown}; choose from {list(LANGUAGE_PROJECTIONS)}")
    args.keep_bf16 = tuple(dict.fromkeys(keep))
    if set(args.keep_bf16) == set(LANGUAGE_PROJECTIONS):
        parser.error("--keep-bf16 names every projection kind, so nothing would be quantized; "
                     "use the bf16 checkpoint as is")
    return args


def layer_plan(source_config: dict) -> tuple[int, int, int]:
    """Return (kept, total, tap): the language layers the conditioner builds from
    this config.json, how many the source has, and the hidden state H3 reads,
    derived the way the model derives them."""
    model_config = MiniMaxH3Qwen3VLConfig()
    model_config.update_model_arch(source_config)
    arch = model_config.arch_config
    total = int(arch.num_hidden_layers)
    tap = int(arch.output_hidden_state_index)
    override = getattr(arch, "num_hidden_layers_override", None)
    kept = min(total, int(override)) if override else total
    return kept, total, tap


def scan_source(shards: list[str], kept: int) -> tuple[Counter, int, dict[str, set[int]]]:
    """Count the language linears the conversion will touch before any work starts.

    Returns per-projection counts below ``kept``, the highest language layer index
    seen, and the layer indexes present per projection for diagnostics.
    """
    counts: Counter = Counter()
    layers_seen: dict[str, set[int]] = {proj: set() for proj in LANGUAGE_PROJECTIONS}
    highest_layer = -1
    for shard in shards:
        with safe_open(shard, framework="pt", device="cpu") as handle:
            for key in handle.keys():
                layer_match = LANGUAGE_LAYER.match(key)
                if layer_match is not None:
                    highest_layer = max(highest_layer, int(layer_match["layer"]))
                linear_match = LANGUAGE_LINEAR.match(key)
                if linear_match is not None and int(linear_match["layer"]) < kept:
                    counts[linear_match["proj"]] += 1
                    layers_seen[linear_match["proj"]].add(int(linear_match["layer"]))
    return counts, highest_layer, layers_seen


class ShardWriter:
    """Accumulate tensors and flush them as fixed-size safetensors shards plus an index.

    Shards are renamed to the ``model-00001-of-0000N.safetensors`` convention
    once the count is known, so the output looks like any Hub checkpoint.
    """

    def __init__(self, dst: Path | None, shard_bytes: int) -> None:
        self.dst = dst
        self.shard_bytes = shard_bytes
        self.pending: dict[str, torch.Tensor] = {}
        self.pending_bytes = 0
        self.weight_map: dict[str, str] = {}
        self.total_bytes = 0
        self.shard_index = 0
        self.written: list[Path] = []

    def add(self, name: str, tensor: torch.Tensor) -> None:
        nbytes = tensor.numel() * tensor.element_size()
        if self.pending and self.pending_bytes + nbytes > self.shard_bytes:
            self.flush()
        self.pending[name] = tensor
        self.pending_bytes += nbytes
        self.total_bytes += nbytes

    def flush(self) -> None:
        if not self.pending:
            return
        self.shard_index += 1
        filename = f"model-{self.shard_index:05d}.partial.safetensors"
        if self.dst is not None:
            save_file(self.pending, str(self.dst / filename), metadata={"format": "pt"})
            self.written.append(self.dst / filename)
        for name in self.pending:
            self.weight_map[name] = filename
        self.pending = {}
        self.pending_bytes = 0

    def finish(self) -> None:
        self.flush()
        if self.dst is None:
            return
        final_names = {}
        for index in range(1, self.shard_index + 1):
            partial = f"model-{index:05d}.partial.safetensors"
            final = f"model-{index:05d}-of-{self.shard_index:05d}.safetensors"
            (self.dst / partial).rename(self.dst / final)
            final_names[partial] = final
        self.written = [self.dst / final for final in final_names.values()]
        weight_map = {name: final_names[partial] for name, partial in self.weight_map.items()}
        index_payload = {"metadata": {"total_size": self.total_bytes}, "weight_map": weight_map}
        (self.dst / SAFE_WEIGHTS_INDEX_NAME).write_text(json.dumps(index_payload, indent=2, sort_keys=True),
                                                        encoding="utf-8")

    def abort(self) -> None:
        """Remove every shard this writer produced so a failed conversion leaves no half checkpoint."""
        for path in self.written:
            path.unlink(missing_ok=True)
        self.written = []


def probe_relative_error(
    weight: torch.Tensor,
    packed: torch.Tensor,
    scale: torch.Tensor,
    global_scale: torch.Tensor,
    rows: int,
    generator: torch.Generator,
) -> float:
    """||fp4(x) @ fp4(W).T - x @ W.T|| / ||x @ W.T|| on random bf16 rows, through the loader's own GEMM path.

    The reference is the bf16 tensor-core product with fp32 accumulation; its
    rounding sits two orders of magnitude below the FP4 error being measured.
    """
    x = torch.randn(rows, weight.shape[1], generator=generator, device=weight.device,
                    dtype=torch.float32).to(torch.bfloat16)
    reference = (x @ weight.t()).float()
    x_fp4, x_scale = _quantize_activation_nvfp4(x, torch.ones((), dtype=torch.float32, device=weight.device))
    output = _nvfp4_linear(x_fp4, x_scale, packed, scale, (1.0 / global_scale).reshape(()))
    return ((output.float() - reference).norm() / reference.norm().clamp_min(1e-12)).item()


def quantize_language_linear(
    weight: torch.Tensor,
    device: torch.device,
    sf_layout: object,
    probe_rows: int,
    generator: torch.Generator | None,
    max_probe_error: float,
    key: str = "",
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float | None]:
    if weight.dtype not in (torch.bfloat16, torch.float16, torch.float32):
        raise ValueError(f"Expected a bf16, fp16 or fp32 source weight, got {weight.dtype}; "
                         "the source text_encoder must be unquantized")
    output_size, input_size = weight.shape
    validate_nvfp4_geometry(output_size, input_size)
    weight_device = weight.to(device=device, dtype=torch.bfloat16)
    global_scale = nvfp4_weight_global_scale(weight_device)
    packed, scale = _nvfp4_quantize(weight_device, global_scale, sfLayout=sf_layout.layout_128x4, do_shuffle=False)
    if packed.dtype != torch.uint8:
        packed = packed.view(torch.uint8)
    if scale.dtype != torch.uint8:
        scale = scale.view(torch.uint8)
    expected_packed = nvfp4_packed_weight_shape(output_size, input_size)
    expected_scale = nvfp4_scale_shape(output_size, input_size)
    if tuple(packed.shape) != expected_packed or tuple(scale.shape) != expected_scale:
        raise RuntimeError("flashinfer.nvfp4_quantize returned an unexpected layout for a "
                           f"{output_size}x{input_size} weight: packed {tuple(packed.shape)} "
                           f"(expected {expected_packed}), scale {tuple(scale.shape)} (expected {expected_scale}). "
                           "The loader allocates exactly these shapes; check the FlashInfer version.")
    error = None
    if probe_rows and generator is not None:
        error = probe_relative_error(weight_device, packed, scale, global_scale, probe_rows, generator)
        if not error <= max_probe_error:
            raise SystemExit(f"FP4 GEMM relative error {error:.3e} on {key or 'a language linear'} exceeds "
                             f"--max-probe-error {max_probe_error}; the packed layout or scales do not reproduce "
                             "the bf16 product, nothing was written")
    return packed.cpu().contiguous(), scale.cpu().contiguous(), global_scale.reshape(1).cpu(), error


def flashinfer_version() -> str:
    try:
        return importlib_metadata.version("flashinfer-python")
    except importlib_metadata.PackageNotFoundError:
        import flashinfer
        return str(getattr(flashinfer, "__version__", "unknown"))


def main() -> None:
    args = parse_args()
    src: Path = args.src
    config_path = src / "config.json"
    if not config_path.is_file():
        raise SystemExit(f"{src} has no config.json; point --src at the text_encoder directory")
    source_config = json.loads(config_path.read_text(encoding="utf-8"))
    if source_config.get("quantization_config"):
        raise SystemExit(f"{src} already carries a quantization_config "
                         f"({source_config['quantization_config'].get('quant_method')!r}); this converter "
                         "quantizes the bf16 text_encoder, not a serialized one")
    kept, total, tap = layer_plan(source_config)

    shards = resolve_safetensors_files(str(src))
    counts, highest_layer, layers_seen = scan_source(shards, kept)
    if not counts:
        raise SystemExit(f"No language linear in {src} matches model.language_model.layers.N.<proj>.weight; "
                         "this converter expects the Qwen3-VL checkpoint layout")
    if highest_layer + 1 < kept:
        raise SystemExit(f"{src} holds {highest_layer + 1} language layers but the conditioner builds {kept}")
    short = {}
    for proj in LANGUAGE_PROJECTIONS:
        missing = sorted(set(range(kept)) - layers_seen[proj])
        if missing:
            short[proj] = f"{len(missing)} missing, e.g. model.language_model.layers.{missing[0]}.{proj}.weight"
    if short:
        raise SystemExit(f"Expected {kept} of every language projection below layer {kept}: {short}")

    device = torch.device(args.device)
    if device.type != "cuda":
        raise SystemExit("The NVFP4 quantizer is a CUDA kernel; pass --device cuda")
    sf_layout, _, _ = _require_flashinfer()

    dst: Path | None = None
    if not args.report_only:
        dst = args.dst
        dst.mkdir(parents=True, exist_ok=True)
        existing = sorted(path.name for path in dst.glob("*.safetensors")) + \
            [name for name in ("config.json", SAFE_WEIGHTS_INDEX_NAME) if (dst / name).exists()]
        if existing:
            raise SystemExit(f"{dst} already holds {existing[:4]}{' and more' if len(existing) > 4 else ''}; "
                             "refusing to overwrite or mix outputs")
    writer = ShardWriter(dst, int(args.shard_size_gb * (1 << 30)))
    generator = torch.Generator(device=device).manual_seed(args.seed) if args.probe_rows else None

    quantized = 0
    kept_bf16 = 0
    kept_bf16_bytes = 0
    copied = 0
    skipped: Counter = Counter()
    source_language_bytes = 0
    written_language_bytes = 0
    errors: dict[str, list[float]] = {}
    started = time.perf_counter()
    print(f"keeping {kept} of {total} language layers, H3 reads hidden state {tap}; "
          f"{sum(counts.values())} language linears to visit", flush=True)
    packed_name, scale_name, global_scale_name = NVFP4_TENSOR_SUFFIXES

    try:
        for shard in shards:
            with safe_open(shard, framework="pt", device="cpu") as handle:
                for key in sorted(handle.keys()):
                    layer_match = LANGUAGE_LAYER.match(key)
                    if key == "lm_head.weight":
                        skipped["lm_head"] += 1
                        continue
                    if layer_match is not None and int(layer_match["layer"]) >= kept:
                        skipped["upper_layers"] += 1
                        continue
                    tensor = handle.get_tensor(key)
                    linear_match = LANGUAGE_LINEAR.match(key)
                    if linear_match is None:
                        writer.add(key, tensor.contiguous())
                        copied += 1
                        continue
                    if linear_match["proj"] in args.keep_bf16:
                        # Kept linears are stored as bf16 whatever the source dtype, matching the
                        # precision the conditioner runs them in.
                        writer.add(key, tensor.to(torch.bfloat16).contiguous())
                        kept_bf16 += 1
                        kept_bf16_bytes += tensor.numel() * 2
                        continue
                    packed, scale, global_scale, error = quantize_language_linear(
                        tensor, device, sf_layout, args.probe_rows, generator, args.max_probe_error, key)
                    prefix = key[:-len(".weight")]
                    writer.add(f"{prefix}.{packed_name}", packed)
                    writer.add(f"{prefix}.{scale_name}", scale)
                    writer.add(f"{prefix}.{global_scale_name}", global_scale)
                    source_language_bytes += tensor.numel() * tensor.element_size()
                    written_language_bytes += packed.numel() + scale.numel() + 4
                    quantized += 1
                    if error is not None:
                        errors.setdefault(linear_match["proj"], []).append(error)
                    if quantized % 35 == 0:
                        print(f"  quantized {quantized} language linears, {time.perf_counter() - started:.0f}s",
                              flush=True)
        writer.finish()
    except BaseException:
        writer.abort()
        raise

    if dst is not None:
        config = dict(source_config)
        config["num_hidden_layers_override"] = kept
        config["quantization_config"] = serialized_nvfp4_quantization_config(
            keep_bf16=args.keep_bf16,
            producer={
                "converter": Path(__file__).name,
                "flashinfer": flashinfer_version(),
                "kept_language_layers": kept,
            })
        (dst / "config.json").write_text(json.dumps(config, indent=2), encoding="utf-8")
        for extra in src.iterdir():
            if extra.is_file() and extra.name not in {SAFE_WEIGHTS_INDEX_NAME, "config.json"} \
                    and not extra.name.endswith(".safetensors"):
                shutil.copy2(extra, dst / extra.name)

    print(f"language linears quantized: {quantized}")
    if args.keep_bf16:
        print(f"language linears kept bf16: {kept_bf16} ({', '.join(args.keep_bf16)}; "
              f"{kept_bf16_bytes / 1e9:.2f} GB)")
    print(f"tensors copied unchanged:   {copied}")
    for skip_id, count in sorted(skipped.items()):
        pattern, reason = SKIPPED[skip_id]
        print(f"skipped {count:>4} keys: {pattern}: {reason}")
    print(f"language linear bytes: {source_language_bytes / 1e9:.2f} GB bf16 -> "
          f"{written_language_bytes / 1e9:.2f} GB NVFP4 (packed values + E4M3 scales)")
    print(f"artifact total: {writer.total_bytes / 1e9:.2f} GB in {writer.shard_index} shard(s)"
          f"{'' if dst is not None else ' (not written, --report-only)'}")
    if errors:
        print(f"FP4 GEMM relative error vs bf16, {args.probe_rows} random rows per linear, "
              f"every linear under --max-probe-error {args.max_probe_error}:")
        print(f"  {'projection':<20}{'linears':>8}{'max':>12}{'mean':>12}")
        for proj in LANGUAGE_PROJECTIONS:
            values = errors.get(proj)
            if values:
                print(f"  {proj:<20}{len(values):>8}{max(values):>12.3e}{sum(values) / len(values):>12.3e}")
    elif quantized:
        print("FP4 GEMM error probe skipped because --probe-rows is 0; nothing verified the written layout")
    print(f"done in {time.perf_counter() - started:.0f}s")


if __name__ == "__main__":
    main()