"""ComfyUI single file (export_comfy.py) -> ConvRot W8A8, the Comfy-Org/MOSS-Image ``*_int8_convrot`` layout.

The 168 block linears (``layers.N.proj.{qkv,o,ffn_up_gate,ffn_down}``) are stored as ComfyUI's native
``int8_tensorwise`` with ConvRot (arXiv 2512.03673); no rotation matrix is stored, the loader rebuilds it:

    <layer>.weight         int8 [out, in]  W rotated per 256-wide input block by the regular Hadamard, quantized
    <layer>.weight_scale   f32  [out, 1]   per-output-row scale, absmax / 127 of the rotated row
    <layer>.comfy_quant    uint8 JSON      {"format": "int8_tensorwise", "convrot": true, "convrot_groupsize": 256}
    layers.N.comfy_attention.config  uint8 JSON  {"attention": "comfy_kitchen_int8"}

H is H4 kron H4 kron H4 kron H4 / 16 with H4 = [[1,1,1,-1],[1,1,-1,1],[1,-1,1,1],[-1,1,1,1]]: symmetric and
orthogonal, so x W^T = (x H)(W H^T)^T and ComfyUI rotates activations online. Everything else (text embedder,
time / modulation, norms, head) passes through unchanged; empty file metadata, as in the Comfy-Org files. The math
mirrors comfy-kitchen's int8_tensorwise (sign pattern, Kronecker order, rounding) or the file decodes to garbage.

python quantize_convrot.py kromoss_bf16.safetensors kromoss_int8_convrot.safetensors \
    [--compare moss_image_sft_int8_convrot.safetensors]
"""
from __future__ import annotations

import os

import argparse
import json
import re
import time

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

BLOCK_LINEAR = re.compile(r"^layers\.(\d+)\.proj\.(qkv|o|ffn_up_gate|ffn_down)\.weight$")
H4 = ((1, 1, 1, -1), (1, 1, -1, 1), (1, -1, 1, 1), (-1, 1, 1, 1))
CLIP_GRID = torch.linspace(0.55, 1.0, 80)


def regular_hadamard(n, device):
    h = torch.tensor(H4, dtype=torch.float32, device=device)
    while h.shape[0] < n:
        h = torch.kron(h, torch.tensor(H4, dtype=torch.float32, device=device))
    if h.shape[0] != n:
        raise ValueError(f"group size must be a power of 4, got {n}")
    return h / n ** 0.5


def rotate(w, gs):
    out_f, in_f = w.shape
    return (w.reshape(out_f, in_f // gs, gs) @ regular_hadamard(gs, w.device).T).reshape(out_f, in_f)


@torch.no_grad()
def quantize(w, gs, mseclip=False):
    """fp32 rotate, then per row scale = absmax / 127 (the Comfy-Org recipe), or with ``mseclip`` the clip ratio
    in CLIP_GRID with the lowest reconstruction MSE."""
    wf = rotate(w.float(), gs)
    absmax = wf.abs().amax(dim=1, keepdim=True).clamp(min=1e-30)
    best_mse, best_scale, best_q = torch.full_like(absmax, float("inf")), absmax / 127.0, None
    for a in CLIP_GRID.tolist() if mseclip else [1.0]:
        scale = (absmax * a / 127.0).clamp(min=1e-30)
        q = (wf / scale).round_().clamp_(-127, 127)
        mse = ((q * scale - wf) ** 2).mean(dim=1, keepdim=True)
        better = mse < best_mse
        best_mse = torch.where(better, mse, best_mse)
        best_scale = torch.where(better, scale, best_scale)
        best_q = q if best_q is None else torch.where(better, q, best_q)
    deq = rotate(best_q * best_scale, gs).double()
    cos = (deq.flatten() @ w.double().flatten() / (deq.norm() * w.double().norm())).item()
    return best_q.to(torch.int8), best_scale, cos


def blob(obj):
    return torch.tensor(list(json.dumps(obj).encode()), dtype=torch.uint8)


def convrot(path, groupsize=256, mseclip=False, int8_attention=True, device=None):
    """bf16 ComfyUI single file -> state dict of the ConvRot W8A8 file (save with metadata={})."""
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    out, start, worst = {}, time.time(), (1.0, "")
    with safe_open(str(path), "pt") as f:
        keys = list(f.keys())
        blocks = sorted({int(m.group(1)) for k in keys if (m := BLOCK_LINEAR.match(k))})
        for k in keys:
            tensor = f.get_tensor(k)
            if not BLOCK_LINEAR.match(k):
                out[k] = tensor
                continue
            q, scale, cos = quantize(tensor.to(device), groupsize, mseclip)
            if cos <= 0.99:
                raise SystemExit(f"broken quantization of {k}: cos {cos:.5f}")
            worst = min(worst, (cos, k))
            base = k[:-len(".weight")]
            out[k], out[f"{base}.weight_scale"] = q.cpu(), scale.cpu()
            out[f"{base}.comfy_quant"] = blob({"format": "int8_tensorwise", "convrot": True,
                                               "convrot_groupsize": groupsize})
    if int8_attention:
        for n in blocks:
            out[f"layers.{n}.comfy_attention.config"] = blob({"attention": "comfy_kitchen_int8"})
    print(f"quantized {sum(1 for k in out if k.endswith('.comfy_quant'))} linears in {time.time() - start:.0f}s, "
          f"worst cos {worst[0]:.6f} ({worst[1]})", flush=True)
    return out


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("input", help="bf16 ComfyUI single file (fused proj.qkv), e.g. from export_comfy.py")
    parser.add_argument("output")
    parser.add_argument("--groupsize", type=int, default=256)
    parser.add_argument("--mseclip", action="store_true", help="per-row MSE-optimal clip instead of absmax "
                        "(~1.5%% lower weight error; not how the Comfy-Org files were made)")
    parser.add_argument("--no-int8-attention", action="store_true",
                        help="omit the per-block comfy_attention.config (ComfyUI then keeps bf16 attention)")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--compare", help="reference int8 file: require the same keys/shapes/dtypes and report "
                                          "how many int8 values and scales agree")
    args = parser.parse_args()
    out = convrot(args.input, args.groupsize, args.mseclip, not args.no_int8_attention, args.device)
    if args.compare:
        with safe_open(args.compare, "pt") as f:
            if set(f.keys()) != set(out):
                raise SystemExit(f"key sets differ: {sorted(set(f.keys()) ^ set(out))[:6]}")
            q_same = q_total = s_same = s_total = exact = 0
            for k in f.keys():
                ref = f.get_tensor(k)
                if ref.shape != out[k].shape or ref.dtype != out[k].dtype:
                    raise SystemExit(f"{k}: {tuple(ref.shape)} {ref.dtype} vs {tuple(out[k].shape)} {out[k].dtype}")
                exact += torch.equal(ref, out[k])
                if ref.dtype == torch.int8:
                    q_same += (ref == out[k]).sum().item()
                    q_total += ref.numel()
                elif k.endswith(".weight_scale"):
                    s_same += (ref == out[k]).sum().item()
                    s_total += ref.numel()
        print(f"vs {args.compare}: {exact}/{len(out)} tensors bit-identical; int8 values equal "
              f"{100 * q_same / q_total:.4f}%, scales equal {100 * s_same / s_total:.4f}%")
    save_file(out, args.output, metadata={})
    print(f"{args.output}: {len(out)} tensors, {os.path.getsize(args.output) / 1e9:.2f} GB")


if __name__ == "__main__":
    main()