"""Quantize MOSS-Image components into ComfyUI-native ConvRot formats (comfy-kitchen layouts).

dit int4 : convrot_w4a4 (group 256, kernel group 64), header _quantization_metadata  (= toxicdog/MOSS-Image-SFT-INT4-ConvRot-ComfyUI)
dit int8 : int8_tensorwise + convrot (per-row scale), per-layer .comfy_quant tensors (= Abiray/MOSS-Image-SFT-INT8-convrot)
te  int8 : same as Abiray text encoder (embed_tokens + LM linears except in_proj_a/b + vision block linears)
"""
import argparse, json, os, re, sys, time
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from comfy_kitchen.tensor import TensorWiseINT8Layout, TensorCoreConvRotW4A4Layout

DIT_RE = re.compile(r"^layers\.\d+\.proj\.(q|k|v|o|ffn_up_gate|ffn_down)\.weight$")
TE_RE = re.compile(r"^(language_model\.embed_tokens|language_model\.layers\.\d+\.(linear_attn\.(in_proj_qkv|in_proj_z|out_proj)|mlp\.(gate|up|down)_proj|self_attn\.(q|k|v|o)_proj)|visual\.blocks\.\d+\.(attn\.(qkv|proj)|mlp\.linear_fc[12]))\.weight$")

def te_group(k):
    for g in (256, 64, 16, 4):
        if k % g == 0:
            return g
    return None

def load_index(d, index):
    m = json.load(open(os.path.join(d, index)))["weight_map"]
    by = {}
    for k, f in m.items():
        by.setdefault(f, []).append(k)
    return by

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--component", choices=["dit", "te"], required=True)
    ap.add_argument("--mode", choices=["int8", "int4"], required=True)
    ap.add_argument("--src", required=True, help="MOSS-Image pipeline dir")
    ap.add_argument("--out", required=True)
    ap.add_argument("--device", default="cpu")
    ap.add_argument("--threads", type=int, default=32)
    a = ap.parse_args()
    torch.set_num_threads(a.threads)
    if os.path.exists(a.out):
        sys.exit(f"exists: {a.out}")
    sub, index = ("transformer", "diffusion_pytorch_model.safetensors.index.json") if a.component == "dit" else ("text_encoder", "model.safetensors.index.json")
    d = os.path.join(a.src, sub)
    rx = DIT_RE if a.component == "dit" else TE_RE
    out, layers, stats = {}, {}, []
    t0 = time.time()
    for fn, keys in sorted(load_index(d, index).items()):
        with safe_open(os.path.join(d, fn), "pt") as f:
            for k in sorted(keys):
                w = f.get_tensor(k)
                if not rx.match(k):
                    out[k] = w.contiguous()
                    continue
                name = k[: -len(".weight")]
                wf = w.to(a.device, torch.float32)
                if a.mode == "int4":
                    q, p = TensorCoreConvRotW4A4Layout.quantize(wf, convrot_groupsize=256)
                    deq = TensorCoreConvRotW4A4Layout.dequantize(q, p)
                    out[name + ".weight"] = q.cpu().contiguous()
                    out[name + ".weight_scale"] = p.scale.float().reshape(-1).cpu().contiguous()
                    layers[name] = {"format": "convrot_w4a4", "convrot_groupsize": 256, "quant_group_size": 64}
                else:
                    g = 256 if a.component == "dit" else te_group(wf.shape[1])
                    q, p = TensorWiseINT8Layout.quantize(wf, is_weight=True, per_channel=True, convrot=True, convrot_groupsize=g)
                    deq = TensorWiseINT8Layout.dequantize(q, p)
                    out[name + ".weight"] = q.cpu().contiguous()
                    out[name + ".weight_scale"] = p.scale.float().reshape(-1, 1).cpu().contiguous()
                    conf = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": g}
                    out[name + ".comfy_quant"] = torch.frombuffer(bytearray(json.dumps(conf).encode()), dtype=torch.uint8).clone()
                    layers[name] = conf
                deq = deq.float()
                rel = ((deq - wf).norm() / wf.norm()).item()
                cos = torch.nn.functional.cosine_similarity(deq.flatten(), wf.flatten(), dim=0).item()
                stats.append((name, list(wf.shape), rel, cos))
                print(f"{name} {list(wf.shape)} rel={rel:.5f} cos={cos:.6f}", flush=True)
    if a.component == "te":
        # HF / ComfyUI-core naming (model.language_model.*, model.visual.*); also accepted by the MOSS-Image plugin
        out = {("model." + k if k.startswith(("language_model.", "visual.")) else k): v for k, v in out.items()}
    meta = {"format": "pt", "model": "moss-image"}
    if a.mode == "int4":
        meta["_quantization_metadata"] = json.dumps({"format_version": "1.0", "layers": layers})
    tmp = a.out + ".partial"
    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    save_file(out, tmp, metadata=meta)
    os.rename(tmp, a.out)
    json.dump(stats, open(a.out + ".stats.json", "w"))
    rels = [s[2] for s in stats]
    print(f"DONE {a.out} layers={len(stats)} tensors={len(out)} mean_rel={sum(rels)/len(rels):.5f} max_rel={max(rels):.5f} min_cos={min(s[3] for s in stats):.6f} size={os.path.getsize(a.out)} t={time.time()-t0:.0f}s")

if __name__ == "__main__":
    main()