#!/usr/bin/env python3
"""Convert a Qwen2.5-0.5B-architecture model (e.g. the merged Devseis auditor) to WebLLM/MLC weights, q4f16_1 or q4f32_1.

Why this exists: MLC's own converter could not run on this Mac (the Intel-Mac nightly mlc_llm and TVM builds do not
match, and the stable Linux release needs an apache-tvm-ffi build that was never published). This script writes the
same format MLC's group quantization produces (4-bit symmetric, groups of 32 along the input dimension, 8 values per
uint32, one scale per group), reusing the tensor layout of MLC's published Qwen2.5-0.5B-Instruct build so WebLLM's
prebuilt runtime library for that model can run it.

It is verified, not assumed: `--verify` converts the original Qwen2.5-0.5B-Instruct and compares it with
mlc-ai/Qwen2.5-0.5B-Instruct-<quant>-MLC: identical layout, bit-identical scales, 99.6% identical 4-bit values (the rest
differ by one step at exact rounding ties, most likely because mlc-ai quantized slightly different source weights).

Usage:
  training/.venv/bin/python training/mlc_quantize.py --verify --quant q4f16_1
  training/.venv/bin/python training/mlc_quantize.py --model training/models/devseis-auditor-0.5b-v0.1 --quant q4f16_1 \\
      --out training/models/mlc/devseis-auditor-0.5b-v0.1-q4f16_1-MLC
"""
import argparse
import hashlib
import json
import os
import shutil

import numpy as np
import torch
from safetensors.torch import load_file

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MODELS = os.path.join(ROOT, "training", "models")
GROUP = 32
MAX_INT = 7  # 4-bit symmetric: values 0..14 stored, 7 = zero
COPY = ["mlc-chat-config.json", "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt"]


def load_hf(model_dir):
    weights = load_file(os.path.join(model_dir, "model.safetensors"))
    return {k: v.to(torch.float32).numpy() for k, v in weights.items()}


def quantize(w, dtype, mode=None):
    """Group quantization along the last axis. Returns (q_weight uint32 [out, in/8], q_scale [out, in/32]).

    Matches MLC: scale = max|w| * (1/7) in the model dtype (bit-exact against mlc-ai's build), values = round(w / scale)
    computed in float32 (99.6% identical to mlc-ai's build; the rest differ by one step at exact rounding ties)."""
    w = w.astype(dtype)
    out, inp = w.shape
    g = w.reshape(out, inp // GROUP, GROUP)
    scale = (np.abs(g).max(axis=-1, keepdims=True) * dtype(1 / MAX_INT)).astype(dtype)
    with np.errstate(divide="ignore", invalid="ignore"):
        x = np.where(scale == 0, 0, g.astype(np.float32) / scale.astype(np.float32))
    x = np.sign(x) * np.floor(np.abs(x) + 0.5)  # round half away from zero
    q = np.clip(x + MAX_INT, 0, 2 * MAX_INT).astype(np.uint32).reshape(out, inp // 8, 8)
    packed = np.zeros((out, inp // 8), dtype=np.uint32)
    for k in range(8):
        packed |= q[..., k] << np.uint32(4 * k)
    return packed, scale.reshape(out, inp // GROUP).astype(dtype)


def build(hf, dtype, mode=None, quantized=True):
    """Tensors named as in MLC's Qwen2 model (q0f16: unquantized float16, same fused names)."""
    t = {}

    def linear(name, w):
        if quantized:
            t[name + ".q_weight"], t[name + ".q_scale"] = quantize(w, dtype, mode)
        else:
            t[name + ".weight"] = w.astype(dtype)

    linear("model.embed_tokens", hf["model.embed_tokens.weight"])
    t["model.norm.weight"] = hf["model.norm.weight"].astype(dtype)
    layer = 0
    while f"model.layers.{layer}.input_layernorm.weight" in hf:
        p = f"model.layers.{layer}."
        t[p + "input_layernorm.weight"] = hf[p + "input_layernorm.weight"].astype(dtype)
        t[p + "post_attention_layernorm.weight"] = hf[p + "post_attention_layernorm.weight"].astype(dtype)
        qkv = np.concatenate([hf[p + f"self_attn.{x}_proj.weight"] for x in "qkv"], axis=0)
        linear(p + "self_attn.c_attn", qkv)
        t[p + "self_attn.c_attn.bias"] = np.concatenate([hf[p + f"self_attn.{x}_proj.bias"] for x in "qkv"]).astype(dtype)
        linear(p + "self_attn.o_proj", hf[p + "self_attn.o_proj.weight"])
        gate_up = np.concatenate([hf[p + "mlp.gate_proj.weight"], hf[p + "mlp.up_proj.weight"]], axis=0)
        linear(p + "mlp.gate_up_proj", gate_up)
        linear(p + "mlp.down_proj", hf[p + "mlp.down_proj.weight"])
        layer += 1
    return t


def stored(rec):
    """MLC stores float32 tensors as bfloat16 when the record's format is f32-to-bf16 (WebLLM expands them on load)."""
    return rec["dtype"] == "float32" and rec.get("format") == "f32-to-bf16"


def to_bytes(arr, rec):
    if stored(rec):
        return torch.from_numpy(np.ascontiguousarray(arr, dtype=np.float32)).to(torch.bfloat16).view(torch.int16).numpy().tobytes()
    return np.ascontiguousarray(arr).tobytes()


def from_bytes(data, rec):
    if stored(rec):
        return (np.frombuffer(data, dtype=np.uint16).astype(np.uint32) << np.uint32(16)).view(np.float32).reshape(rec["shape"])
    return np.frombuffer(data, dtype=rec["dtype"]).reshape(rec["shape"])


def records(ref_dir):
    cache = json.load(open(os.path.join(ref_dir, "tensor-cache.json")))
    return cache, [(shard, rec) for shard in cache["records"] for rec in shard["records"]]


def check_layout(tensors, recs):
    for _, rec in recs:
        arr = tensors[rec["name"]]
        assert list(arr.shape) == rec["shape"], (rec["name"], arr.shape, rec["shape"])
        assert str(arr.dtype) == rec["dtype"], (rec["name"], arr.dtype, rec["dtype"])
        assert (arr.nbytes // 2 if stored(rec) else arr.nbytes) == rec["nbytes"], rec["name"]
    assert len(tensors) == len(recs), (len(tensors), len(recs))


def write(tensors, ref_dir, out_dir):
    cache, recs = records(ref_dir)
    check_layout(tensors, recs)
    os.makedirs(out_dir, exist_ok=True)
    for shard in cache["records"]:
        data = bytearray(shard["nbytes"])
        for rec in shard["records"]:
            data[rec["byteOffset"]:rec["byteOffset"] + rec["nbytes"]] = to_bytes(tensors[rec["name"]], rec)
        with open(os.path.join(out_dir, shard["dataPath"]), "wb") as handle:
            handle.write(data)
        shard["md5sum"] = hashlib.md5(data).hexdigest()
    for name in ("tensor-cache.json", "ndarray-cache.json"):
        with open(os.path.join(out_dir, name), "w") as handle:
            json.dump(cache, handle)
    for name in COPY:
        shutil.copy(os.path.join(ref_dir, name), out_dir)


def verify(quant, ref_root=None, base=None):
    """Convert the base model and compare with mlc-ai's build: layout must match exactly, scales bit-exactly, values >= 99%."""
    ref = os.path.join(ref_root or os.path.join(MODELS, "mlc-ref"), quant)
    dtype = np.float32 if quant == "q4f32_1" else np.float16
    _, recs = records(ref)
    tensors = build(load_hf(base or os.path.join(MODELS, "Qwen2.5-0.5B-Instruct")), dtype, quantized=not quant.startswith("q0"))
    check_layout(tensors, recs)
    if quant.startswith("q0"):
        present = [(s, r) for s, r in recs if os.path.exists(os.path.join(ref, s["dataPath"]))]
        same = sum(open(os.path.join(ref, s["dataPath"]), "rb").read()[r["byteOffset"]:r["byteOffset"] + r["nbytes"]] == to_bytes(tensors[r["name"]], r) for s, r in present)
        print(f"{quant}: layout identical ({len(recs)} tensors); {same}/{len(present)} tensors in the downloaded shards byte-identical")
        return same == len(present) and len(present) > 0
    shards, scale_same, scale_total, val_same, val_total, exact = {}, 0, 0, 0, 0, 0
    for shard, rec in recs:
        if shard["dataPath"] not in shards:
            shards[shard["dataPath"]] = open(os.path.join(ref, shard["dataPath"]), "rb").read()
        want = from_bytes(shards[shard["dataPath"]][rec["byteOffset"]:rec["byteOffset"] + rec["nbytes"]], rec)
        got = from_bytes(to_bytes(tensors[rec["name"]], rec), rec)
        exact += bool(np.array_equal(got, want))
        if rec["name"].endswith("q_scale"):
            scale_same += int((got == want).sum()); scale_total += got.size
        elif rec["name"].endswith("q_weight"):
            a = np.stack([(got >> np.uint32(4 * k)) & np.uint32(15) for k in range(8)], -1)
            b = np.stack([(want >> np.uint32(4 * k)) & np.uint32(15) for k in range(8)], -1)
            val_same += int((a == b).sum()); val_total += a.size
            assert int(np.abs(a.astype(np.int32) - b.astype(np.int32)).max()) <= 1, rec["name"]
    print(f"{quant}: layout identical ({len(recs)} tensors, {exact} byte-identical); scales {scale_same / scale_total:.4%} identical; "
          f"4-bit values {val_same / val_total:.4%} identical, all others within one step")
    return scale_same == scale_total and val_same / val_total >= 0.99


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--quant", choices=["q4f16_1", "q4f32_1", "q0f16"], required=True)
    parser.add_argument("--verify", action="store_true", help="compare a conversion of the base model with mlc-ai's build")
    parser.add_argument("--model", help="merged model folder (model.safetensors)")
    parser.add_argument("--out", help="output folder")
    parser.add_argument("--ref", help="folder with mlc-ai reference builds per quantization (default training/models/mlc-ref, 0.5B)")
    parser.add_argument("--base", help="--verify: original base model folder matching --ref")
    args = parser.parse_args()
    if args.verify:
        raise SystemExit(0 if verify(args.quant, args.ref, args.base) else 1)
    dtype = np.float32 if args.quant == "q4f32_1" else np.float16
    tensors = build(load_hf(args.model), dtype, quantized=not args.quant.startswith("q0"))
    write(tensors, os.path.join(args.ref or os.path.join(MODELS, "mlc-ref"), args.quant), args.out)
    print(f"wrote {args.out}")


if __name__ == "__main__":
    main()