#!/usr/bin/env python3
"""Repack a student GGUF's expert gate/up tensors from Q2_0 to llama.cpp's TQ1_0 (1.69 bits per weight).

    python scripts/tq1_repack.py <in.gguf> <out.gguf>

TQ1_0 holds ternary codes for blocks of 256 weights with ONE fp16 scale, and needs rows whose length is a
multiple of 256. Our groups are 128 wide with their own scales, so each TQ1_0 block covers two groups:
- the trained codes are kept exactly;
- the block's scale is the least-squares fit to both groups' values, mean |w| over its nonzero codes;
- ffn_down_exps (rows of 384) cannot be TQ1_0 and stays Q2_0, as does everything else.
Prints the size change and the scale approximation's relative weight error.
"""
from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "src"))
from lowbit.export.gguf import dequant_q2_0, gguf_module, write_gguf  # noqa: E402

LLAMA = Path("/path/to/llama.cpp")
TQ_BLOCK = 256


def to_tq1(raw: np.ndarray, gguf) -> tuple[np.ndarray, float]:
    """Q2_0 bytes [..., nb · 18] → TQ1_0 bytes [..., cols / 256 · 54], and the relative L2 weight error."""
    lead = raw.shape[:-1]
    w = dequant_q2_0(raw)                                   # [..., cols] float32, codes · group scale
    cols = w.shape[-1]
    if cols % TQ_BLOCK:
        raise ValueError(f"rows of {cols} are not a multiple of {TQ_BLOCK}")
    blk = w.reshape(-1, TQ_BLOCK)
    nz = blk != 0
    d = np.abs(blk).sum(-1) / np.maximum(nz.sum(-1), 1)     # least squares with the codes fixed
    codes = np.sign(blk)
    x = (codes * d[:, None]).astype(np.float32)             # TQ1_0's own quantiser then keeps d and the codes
    err = float(np.linalg.norm(x - blk) / max(np.linalg.norm(blk), 1e-30))
    q = gguf.quants.quantize(x.reshape(*lead, cols), gguf.GGMLQuantizationType.TQ1_0)
    back = gguf.quants.dequantize(q, gguf.GGMLQuantizationType.TQ1_0).reshape(-1, TQ_BLOCK)
    if not np.array_equal(np.sign(back), codes):
        raise RuntimeError("TQ1_0 changed a code")
    return q, err


def main() -> None:
    src, dst = sys.argv[1], sys.argv[2]
    gguf = gguf_module(LLAMA)
    r = gguf.GGUFReader(src)
    T = gguf.GGMLQuantizationType
    tensors, errs, moved = {}, [], 0
    for t in r.tensors:
        data = np.asarray(t.data)
        if t.tensor_type == T.Q2_0 and (t.name.endswith("ffn_gate_exps.weight") or t.name.endswith("ffn_up_exps.weight")):
            q, e = to_tq1(data, gguf)
            tensors[t.name] = (T.TQ1_0, q)
            errs.append(e)
            moved += data.nbytes - q.nbytes
        else:
            tensors[t.name] = (t.tensor_type, data)
    info = write_gguf(src, dst, tensors, gguf)
    out = {"in_bytes": Path(src).stat().st_size, "out_bytes": info["bytes"],
           "saved_pct": round(100 * (1 - info["bytes"] / Path(src).stat().st_size), 1),
           "tq1_tensors": len(errs), "scale_rel_err_mean": round(float(np.mean(errs)), 4),
           "scale_rel_err_max": round(float(np.max(errs)), 4)}
    print(json.dumps(out, indent=1))


if __name__ == "__main__":
    main()