#!/usr/bin/env python3
"""Quantize a native ComfyUI single-file OmniRef checkpoint to native FP8 or NVFP4.

The input is the native ComfyUI layout produced by ``convert_omniref_comfy.py``:
Q/K/V are already fused into ``attn.qkv_proj`` and the PDD output heads use the
native first-head-plus-offset representation. The quantizer therefore sees one
fused QKV tensor and applies one Comfy weight scale to it, exactly as the
pinned ComfyUI runtime expects.

Only the per-block attention/MLP projections are quantized. The token refiner,
AdaLN curves and projections, patch/condition/output projections, norms, RoPE,
and FP32 PDD heads stay dense. This matches the native ComfyUI MiniMax-H3
quantized-layer policy used by the released V2/Trim Comfy checkpoints.

Writes a JSON receipt with source/output hashes, the exact format markers,
per-linear reconstruction metrics, and runtime versions.
"""
from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import importlib.metadata
import json
from pathlib import Path
import re
import sys
import time

import torch

# Native Comfy MiniMax-H3 block linears. qkv_proj is the fused Q/K/V module.
QUANT_LINEAR = re.compile(
    r"^blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj|to_gate_compress)|mlp\.(?:fc1|fc2))\.weight$"
)

FORMATS = {
    "fp8": {
        "layout": "TensorCoreFP8Layout",
        "tag": {"format": "float8_e4m3fn"},
        "documented_name": "FP8 E4M3, per-tensor weight scale, dynamic per-tensor activation scale",
        "activation": "dynamic-per-tensor",
    },
    "nvfp4": {
        "layout": "TensorCoreNVFP4Layout",
        "tag": {"format": "nvfp4"},
        "documented_name": "NVFP4 E2M1, per-tensor global scale plus FP8 block scales (group 16), dynamic per-tensor activation scale",
        "activation": "dynamic-per-tensor",
    },
}


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while True:
            block = handle.read(64 * 1024 * 1024)
            if not block:
                break
            digest.update(block)
    return digest.hexdigest()


def tensor_digest(tensor: torch.Tensor) -> str:
    return hashlib.sha256(tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()).hexdigest()


def safetensors_header(path: Path) -> dict:
    from safetensors import safe_open

    with safe_open(str(path), framework="pt", device="cpu") as handle:
        keys = list(handle.keys())
        metadata = handle.metadata() or {}
    return {"tensor_count": len(keys), "metadata": metadata, "keys": keys}


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--src", required=True, type=Path, help="Native ComfyUI BF16 single-file checkpoint")
    parser.add_argument("--dst", required=True, type=Path, help="Output quantized single-file checkpoint")
    parser.add_argument("--format", required=True, choices=sorted(FORMATS))
    parser.add_argument("--receipt", required=True, type=Path)
    parser.add_argument("--comfyui-commit", default="b1a247c9cf9dbfda4b2c0a0c1d50950667eeaa0c")
    parser.add_argument("--expect-tensors", type=int, default=None)
    args = parser.parse_args()

    if args.src == args.dst:
        raise SystemExit("source and destination must differ")
    if args.dst.exists():
        raise SystemExit(f"refusing to overwrite existing output: {args.dst}")
    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required for comfy-kitchen quantization")

    import comfy_kitchen
    import comfy_kitchen.tensor as ckt
    from safetensors import safe_open
    from safetensors.torch import load_file, save_file

    started = time.time()
    spec = FORMATS[args.format]
    layout_cls = getattr(ckt, spec["layout"])
    tag_bytes = json.dumps(spec["tag"], separators=(",", ":")).encode()

    print(f"[quantize] format={args.format} src={args.src} dst={args.dst}", flush=True)
    src_header = safetensors_header(args.src)
    if args.expect_tensors is not None and src_header["tensor_count"] != args.expect_tensors:
        raise SystemExit(f"expected {args.expect_tensors} tensors, found {src_header['tensor_count']}")
    print(f"[quantize] source tensors={src_header['tensor_count']} hashing source", flush=True)
    src_sha256 = sha256_file(args.src)

    # The native Comfy BF16 export is ~44 GB. Keeping the pass-through tensors
    # plus the quantized output in CPU RAM stays well under a GB200 node's RAM.
    source = load_file(str(args.src), device="cpu")
    metadata = dict(src_header["metadata"])
    metadata["quantization"] = json.dumps(
        {
            "format": spec["tag"]["format"],
            "layout": spec["layout"],
            "weights": "per-tensor global scale" if args.format == "fp8" else "per-tensor global scale plus FP8 block scales, group 16",
            "activations": spec["activation"],
            "quantized_layers": "all 50 DiT blocks: fused qkv_proj, out_proj, to_gate_compress, mlp.fc1, mlp.fc2",
            "dense_layers": "token_refiner, AdaLN, norms, embeddings, patch/condition projections, FP32 PDD heads",
            "qkv_treatment": "fused Q/K/V quantized as one tensor under one weight scale",
            "comfyui_commit": args.comfyui_commit,
        },
        separators=(",", ":"),
    )
    metadata["quantization_source_sha256"] = src_sha256
    metadata["quantization_recipe"] = "omniref-comfy-ptq-v1"

    output: dict[str, torch.Tensor] = {}
    linears = []
    for index, (name, tensor) in enumerate(source.items(), 1):
        if not QUANT_LINEAR.match(name):
            output[name] = tensor
            continue

        weight = tensor.to(device="cuda", dtype=torch.bfloat16).contiguous()
        qdata, params = layout_cls.quantize(weight)
        serialized = layout_cls.state_dict_tensors(qdata, params)
        for suffix, value in serialized.items():
            output[name + suffix] = value.detach().cpu().contiguous()
        output[name[: -len(".weight")] + ".comfy_quant"] = torch.tensor(list(tag_bytes), dtype=torch.uint8)

        dequantized = layout_cls.dequantize(qdata, params).float()
        reference = weight.float()
        diff = dequantized - reference
        rel_l2 = float(diff.norm() / reference.norm().clamp(min=1e-30))
        cosine = float(torch.nn.functional.cosine_similarity(dequantized.flatten(), reference.flatten(), dim=0))
        entry = {
            "name": name,
            "logical_shape": list(tensor.shape),
            "weight_storage_shape": list(serialized[""].shape),
            "weight_storage_dtype": str(serialized[""].dtype).replace("torch.", ""),
            "reconstruction_relative_l2": rel_l2,
            "reconstruction_cosine": cosine,
            "reconstruction_max_abs": float(diff.abs().max()),
            "source_amax": float(reference.abs().max()),
        }
        if args.format == "fp8":
            entry["weight_scale"] = float(params.scale.reshape(-1)[0])
        else:
            entry["weight_scale_2"] = float(params.scale.reshape(-1)[0])
            entry["weight_scale_shape"] = list(serialized["_scale"].shape)
            entry["weight_scale_dtype"] = str(serialized["_scale"].dtype).replace("torch.", "")
            entry["block_scale_min"] = float(serialized["_scale"].to(torch.float32).min())
            entry["block_scale_max"] = float(serialized["_scale"].to(torch.float32).max())
        linears.append(entry)
        print(
            f"[quantize] {index:4d}/{len(source)} {name} relL2={rel_l2:.6f} cos={cosine:.8f}",
            flush=True,
        )
        del weight, qdata, params, dequantized, reference, diff
        torch.cuda.empty_cache()

    expected_linears = 50 * 5
    if len(linears) != expected_linears:
        raise SystemExit(f"expected {expected_linears} quantized linears, got {len(linears)}")

    print(f"[quantize] writing {len(output)} tensors to {args.dst}", flush=True)
    args.dst.parent.mkdir(parents=True, exist_ok=True)
    save_file(output, str(args.dst), metadata=metadata)
    out_sha256 = sha256_file(args.dst)
    out_header = safetensors_header(args.dst)

    # Pass-through tensors must survive byte-for-byte. Compare every dense tensor
    # from the written file against the source file, not against the in-memory copy.
    pass_through = [name for name in source if not QUANT_LINEAR.match(name)]
    identical = 0
    with safe_open(str(args.src), framework="pt", device="cpu") as src_reader, safe_open(
        str(args.dst), framework="pt", device="cpu"
    ) as dst_reader:
        for name in pass_through:
            left = src_reader.get_tensor(name)
            right = dst_reader.get_tensor(name)
            if left.dtype != right.dtype or left.shape != right.shape or not torch.equal(left, right):
                raise SystemExit(f"pass-through tensor changed: {name}")
            identical += 1

    receipt = {
        "schema": "omniref-comfy-quant-receipt-v1",
        "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
        "format": args.format,
        "format_description": spec["documented_name"],
        "source": {
            "path": str(args.src),
            "bytes": args.src.stat().st_size,
            "sha256": src_sha256,
            "tensor_count": src_header["tensor_count"],
            "metadata": src_header["metadata"],
        },
        "output": {
            "path": str(args.dst),
            "bytes": args.dst.stat().st_size,
            "sha256": out_sha256,
            "tensor_count": out_header["tensor_count"],
        },
        "runtime": {
            "torch": torch.__version__,
            "cuda": torch.version.cuda,
            "gpu": torch.cuda.get_device_name(0),
            "comfy_kitchen": importlib.metadata.version("comfy-kitchen"),
            "comfyui_commit": args.comfyui_commit,
        },
        "policy": {
            "quantized_linears": "all 50 DiT blocks: fused qkv_proj, out_proj, to_gate_compress, mlp.fc1, mlp.fc2",
            "dense_linears": "token refiner, AdaLN, norms, embeddings, patch/condition projections, FP32 PDD heads",
            "qkv": "fused by the native Comfy conversion, then quantized as one tensor with one weight scale",
            "activation_quantization": spec["activation"],
            "native_comfy_layout": spec["layout"],
        },
        "quantized_linear_count": len(linears),
        "quantized_linears": linears,
        "pass_through_tensor_count": identical,
        "pass_through_byte_identical": True,
        "elapsed_seconds": round(time.time() - started, 1),
    }
    args.receipt.parent.mkdir(parents=True, exist_ok=True)
    args.receipt.write_text(json.dumps(receipt, indent=2) + "\n")
    print(json.dumps({k: receipt[k] for k in ("format", "output", "quantized_linear_count", "elapsed_seconds")}, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())