#!/usr/bin/env python3
"""Convert OLMoE HuggingFace checkpoint to colibri int4 format.

Downloads or converts a local OLMoE checkpoint (e.g., allenai/OLMoE-1B-7B-0125-Instruct).
Dense weights stay as-is (engine reads BF16/F16 → F32 on load).
Expert weights get row-wise symmetric quantization to --ebits bits (default 4)
with float32 scales. Storage stays one value per int8 byte regardless of bits,
matching the engine's expert layout (olmoe.c quantize_rows) — for 4 bits the
values are simply confined to [-8, 7] with scales computed against qmax=7.

Usage:
  python tools/convert_olmoe.py --repo allenai/OLMoE-1B-7B-0125-Instruct --out ./olmoe_i4
  python tools/convert_olmoe.py --model ./OLMoE-1B-7B-0125-Instruct --out ./olmoe_i4
"""

import argparse, json, math, os, struct, sys
from pathlib import Path

# Windows: force UTF-8 output
if sys.platform == "win32":
    for s in (sys.stdout, sys.stderr):
        try: s.reconfigure(encoding="utf-8")
        except (AttributeError, OSError): pass

try:
    import torch
    from safetensors.torch import load_file, save_file
except ImportError as exc:
    sys.exit(f"Missing dependencies: {exc}. Install: pip install torch safetensors")


EXPERT_KEY_RE = r"model\.layers\.\d+\.mlp\.experts\.\d+\.(gate_proj|up_proj|down_proj)\.weight"


def quantize_row(w: torch.Tensor, bits: int = 8) -> tuple[torch.Tensor, torch.Tensor]:
    """Row-wise symmetric quantization to `bits` (2..8).

    Returns (int8_weights, float32_scales). Storage is one value per int8 byte
    for every bit width — the engine dequantizes as q*scale and never assumes
    the full int8 range — mirroring olmoe.c quantize_rows():
        qmax  = 2**(bits-1) - 1        (8 -> 127, 4 -> 7, 2 -> 1)
        scale = amax(|w|, row) / qmax
        q     = clamp(round(w / scale), -qmax-1, qmax)
    """
    qmax = (1 << (bits - 1)) - 1
    w_f32 = w.float()
    row_max = w_f32.abs().amax(dim=1, keepdim=True).clamp(min=1e-12)
    scales = row_max / qmax
    q = (w_f32 / scales).round().clamp(-qmax - 1, qmax).to(torch.int8)
    return q, scales.squeeze(1)


def is_expert_weight(name: str) -> bool:
    import re
    return bool(re.search(EXPERT_KEY_RE, name))


def main():
    ap = argparse.ArgumentParser(description="Convert OLMoE HF checkpoint -> colibri int4")
    src = ap.add_mutually_exclusive_group(required=True)
    src.add_argument("--repo", help="HuggingFace repo ID")
    src.add_argument("--model", help="Local HF checkpoint directory")
    ap.add_argument("--out", required=True, help="Output directory for int4 model")
    ap.add_argument("--ebits", type=int, default=4,
                    help="Expert quant bits (2..8, default 4)")
    args = ap.parse_args()

    if not 2 <= args.ebits <= 8:   # storage is int8_t; engine rejects the same range (olmoe.c)
        sys.exit(f"--ebits must be 2..8 (got {args.ebits})")

    if args.repo:
        from huggingface_hub import snapshot_download
        from huggingface_hub.errors import LocalEntryNotFoundError
        print(f"Downloading {args.repo}...")
        try:
            src_dir = snapshot_download(args.repo, local_files_only=True, max_workers=4)
        except LocalEntryNotFoundError:
            src_dir = None
        if src_dir is None or not any(Path(src_dir).glob("*.safetensors")):
            print("Downloading safetensors...")
            src_dir = snapshot_download(args.repo, max_workers=4)
    else:
        src_dir = args.model

    src = Path(src_dir)
    if not src.is_dir():
        sys.exit(f"Model directory not found: {src}")
    if not (src / "config.json").is_file():
        sys.exit(f"config.json missing in {src}")

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    # Copy config.json
    import shutil
    shutil.copy2(src / "config.json", out / "config.json")
    print(f"config.json -> {out}")

    # Process safetensors
    shards = sorted(src.glob("*.safetensors"))
    if not shards:
        sys.exit(f"No safetensors found in {src}")

    expert_count = 0
    total_expert_f32 = total_expert_q = 0

    for si, shard in enumerate(shards, 1):
        print(f"[{si}/{len(shards)}] {shard.name}...", end=" ", flush=True)
        tensors = load_file(str(shard))
        out_tensors = {}
        for name, tensor in tensors.items():
            if is_expert_weight(name):
                expert_count += 1
                q, scales = quantize_row(tensor, args.ebits)
                total_expert_f32 += tensor.numel() * tensor.element_size()
                total_expert_q += q.numel() * 1 + scales.numel() * 4
                out_tensors[name] = q
                out_tensors[name + ".qs"] = scales
            else:
                out_tensors[name] = tensor

        out_shard = out / shard.name
        save_file(out_tensors, str(out_shard))
        ratio = total_expert_q / max(total_expert_f32, 1) * 100
        print(f"ok")

    print(f"\nDone. {expert_count} expert tensors quantized to int{args.ebits}.")
    print(f"Expert storage: {total_expert_f32/1e9:.1f} GB -> {total_expert_q/1e9:.1f} GB ({ratio:.0f}%)")
    print(f"Model ready at: {out}")
    print(f"\nRun: SNAP={out} ./olmoe.exe 32 4 16")


if __name__ == "__main__":
    main()