#!/usr/bin/env python3
"""Packed unary converter: uint8 magnitudes + bitpacked signs + per-row scales."""
import os, json, sys, time
import numpy as np
from pathlib import Path

def load_safetensors(model_dir):
    from safetensors.torch import load_file
    tensors = {}
    for f in sorted(Path(model_dir).glob("*.safetensors")):
        print(f"  Loading {f.name}...")
        for k, v in load_file(str(f)).items():
            tensors[k] = v.float().numpy()
    return tensors

def quantize_packed(w, n_levels=7):
    out_dim, in_dim = w.shape
    chunks = (in_dim + 63) // 64
    padded = chunks * 64
    row_max = np.max(np.abs(w), axis=1, keepdims=True)
    row_max = np.where(row_max == 0, 1.0, row_max)
    scales = (row_max.flatten() / n_levels).astype(np.float32)
    mags = np.clip(np.round(np.abs(w / scales[:, None])), 0, n_levels).astype(np.uint8)
    signs = (w < 0)
    rmm = np.max(mags, axis=1).astype(np.uint8)
    if in_dim < padded:
        sp = np.zeros((out_dim, padded), dtype=bool)
        sp[:, :in_dim] = signs
    else:
        sp = signs
    bit_pos = np.uint64(1) << np.arange(64, dtype=np.uint64)
    sign_bits = np.bitwise_or.reduce(sp.reshape(out_dim, chunks, 64).astype(np.uint64) * bit_pos, axis=2)
    return mags, sign_bits, scales, rmm, np.mean(mags), np.mean(mags == 0)

def convert(tensors, output_dir, n_levels=7):
    os.makedirs(output_dir, exist_ok=True)
    config = {"hidden_size":1536,"intermediate_size":8960,"num_attention_heads":12,
              "num_key_value_heads":2,"num_hidden_layers":28,"vocab_size":151936,
              "head_dim":128,"rope_theta":1000000.0,"rms_norm_eps":1e-6,
              "n_levels":n_levels,"quant_type":"packed_unary"}
    linear_keys = [k for k in tensors if any(p in k for p in
        ['q_proj.weight','k_proj.weight','v_proj.weight','o_proj.weight',
         'gate_proj.weight','up_proj.weight','down_proj.weight'])]
    other_keys = [k for k in tensors if k not in linear_keys]
    with open(os.path.join(output_dir, "config.json"), "w") as f:
        json.dump(config, f, indent=2)
    total_packed = total_orig = 0
    all_avg = []
    for key in linear_keys:
        w = tensors[key]; total_orig += w.nbytes
        t0 = time.time()
        mags, sb, sc, rmm, am, sp = quantize_packed(w, n_levels)
        dt = time.time() - t0
        pfx = os.path.join(output_dir, key.replace(".", "_"))
        mags.tofile(pfx+".mags"); sb.tofile(pfx+".signs")
        sc.tofile(pfx+".scales"); rmm.tofile(pfx+".rmm")
        ub = mags.nbytes + sb.nbytes + sc.nbytes + rmm.nbytes
        total_packed += ub; all_avg.append(am)
        print(f"  {key}: {w.shape} -> {ub/1024:.0f}KB (avg_mag={am:.2f}, {dt:.1f}s)")
    total_fp16 = 0
    for key in other_keys:
        w = tensors[key].astype(np.float16)
        pfx = os.path.join(output_dir, key.replace(".", "_"))
        w.tofile(pfx+".fp16"); total_fp16 += w.nbytes
    manifest = {"packed":{k:list(tensors[k].shape) for k in linear_keys},
                "fp16":{k:list(tensors[k].shape) for k in other_keys}}
    with open(os.path.join(output_dir, "manifest.json"), "w") as f:
        json.dump(manifest, f, indent=2)
    print(f"\n=== PACKED UNARY ===")
    print(f"Packed linear: {total_packed/1e6:.1f} MB | FP16 other: {total_fp16/1e6:.1f} MB")
    print(f"Total: {(total_packed+total_fp16)/1e6:.1f} MB | Avg mag: {np.mean(all_avg):.3f}")
    print(f"Expected speedup vs 7-plane: {7/np.mean(all_avg):.1f}x")

if __name__ == "__main__":
    model_dir = sys.argv[1] if len(sys.argv) > 1 else "deepseek-r1-1.5b-hf"
    output_dir = sys.argv[2] if len(sys.argv) > 2 else "deepseek-r1-1.5b-packed"
    tensors = load_safetensors(model_dir)
    convert(tensors, output_dir)
    print("Done!")