#!/usr/bin/env python3
"""Requantize selected GGUF tensors (default: refiner FFN/attention weights) to Q4_K.

The reference Z-Image-Turbo Q4_K GGUF stores the two noise-refiner and two
context-refiner blocks as Q8_0. On a memory-bound CPU GEMM those 8-bit weights
make the refiner matmuls ~2x slower than the Q4_K main layers even though the
token count is smaller. This tool rewrites just the matching Q8_0 tensors as
Q4_K, leaving every other tensor byte-for-byte untouched.

Requires the small `q4k` helper (build with build.sh or see README) which wraps
ggml's own Q4_K quantizer.
"""
import argparse
import os
import subprocess
import sys
import tempfile

import numpy as np
import gguf
from gguf import GGUFReader, GGUFWriter, GGUFValueType

Q4K = os.environ.get("Q4K", "/tmp/q4k")


def q4k_quantize(f32: np.ndarray, helper: str) -> np.ndarray:
    nrows, n_per_row = f32.shape
    p = subprocess.run([helper, str(nrows), str(n_per_row)],
                       input=np.ascontiguousarray(f32, dtype=np.float32).tobytes(),
                       stdout=subprocess.PIPE, check=True)
    return np.frombuffer(p.stdout, dtype=np.uint8)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("src")
    ap.add_argument("dst")
    ap.add_argument("--pattern", default="refiner")
    ap.add_argument("--only-kinds", default="attention,feed_forward",
                    help="comma list of substrings a tensor name must contain (empty = any)")
    ap.add_argument("--helper", default=Q4K)
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    r = GGUFReader(args.src)
    arch = str(r.fields["general.architecture"].contents()) if "general.architecture" in r.fields else "zimage"

    # pick targets
    kinds = [k for k in args.only_kinds.split(",") if k]
    targets = []
    for t in r.tensors:
        if args.pattern in t.name and t.tensor_type == gguf.GGMLQuantizationType.Q8_0:
            if kinds and not any(k in t.name for k in kinds):
                continue
            targets.append(t.name)
    print(f"targets ({len(targets)}):")
    for n in targets:
        print("  ", n)
    if args.dry_run or not targets:
        return 0

    w = GGUFWriter(args.dst, arch, use_temp_file=True)
    for name, field in r.fields.items():
        if name == "general.architecture" or name.startswith("GGUF."):
            continue
        v = field.contents()
        w.add_key_value(name, v, field.types[0],
                        sub_type=field.types[-1] if len(field.types) > 1 else None)

    tset = set(targets)
    for t in r.tensors:
        if t.name in tset:
            f32 = gguf.quants.dequantize(t.data, t.tensor_type)
            raw = q4k_quantize(f32, args.helper)
            tgt = int(t.shape[0]) // 256 * 144  # Q4_K bytes per row (ne0 = input width)
            w.add_tensor(t.name, raw.reshape(-1, tgt),
                         raw_dtype=gguf.GGMLQuantizationType.Q4_K)
            print(f"  requantized {t.name} {tuple(t.shape)} Q8_0->Q4_K "
                  f"{t.n_bytes} -> {raw.nbytes} bytes")
        else:
            w.add_tensor(t.name, t.data, raw_dtype=t.tensor_type)
    w.write_header_to_file()
    w.write_kv_data_to_file()
    w.write_tensors_to_file()
    w.close()
    print("wrote", args.dst)
    return 0


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