"""
Convert a HuggingFace GPT-2-architecture model to llm.c binary format (version 5, bf16).

Usage:
  python import_hf.py --input <hf_model_dir> --output <output.bin>

The output is compatible with train_gpt2cu's -e flag.
"""

import argparse
import math
import struct
import numpy as np
import torch
from transformers import GPT2LMHeadModel, AutoConfig, AutoModelForCausalLM


def fp32_to_bf16_int16(tensor: torch.Tensor) -> np.ndarray:
    """Convert float32 tensor to bfloat16 stored as int16 (same byte layout)."""
    bf16 = tensor.to(torch.bfloat16)
    return bf16.view(torch.int16).cpu().numpy()


def write_model(model: GPT2LMHeadModel, output_path: str) -> None:
    cfg = model.config
    V   = cfg.vocab_size
    maxT = cfg.n_positions
    L   = cfg.n_layer
    H   = cfg.n_head
    C   = cfg.n_embd

    # Pad vocab to next multiple of 128
    Vp = math.ceil(V / 128) * 128

    print(f"V={V}, Vp={Vp}, maxT={maxT}, L={L}, H={H}, C={C}")

    # Build header: 256 int32 values
    header = np.zeros(256, dtype=np.int32)
    header[0] = 20240326   # magic
    header[1] = 5          # version 5 = bf16 + padded vocab
    header[2] = maxT
    header[3] = V
    header[4] = L
    header[5] = H
    header[6] = C
    header[7] = Vp

    sd = model.state_dict()

    def get(key):
        return sd[key].float()

    with open(output_path, "wb") as f:
        f.write(header.tobytes())

        # wte: (Vp, C) — pad vocab rows with zeros
        wte = get("transformer.wte.weight")       # (V, C)
        pad = torch.zeros(Vp - V, C, dtype=torch.float32)
        wte_padded = torch.cat([wte, pad], dim=0) # (Vp, C)
        f.write(fp32_to_bf16_int16(wte_padded).tobytes())

        # wpe: (maxT, C)
        f.write(fp32_to_bf16_int16(get("transformer.wpe.weight")).tobytes())

        # Per-layer weights
        for i in range(L):
            # ln1w, ln1b: (C,)
            f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.ln_1.weight")).tobytes())
        for i in range(L):
            f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.ln_1.bias")).tobytes())

        # qkvw: (L, 3C, C) — HF stores c_attn.weight as (C, 3C), transpose to (3C, C)
        for i in range(L):
            w = get(f"transformer.h.{i}.attn.c_attn.weight")  # (C, 3C)
            f.write(fp32_to_bf16_int16(w.T.contiguous()).tobytes())  # (3C, C)
        # qkvb: (L, 3C)
        for i in range(L):
            f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.attn.c_attn.bias")).tobytes())

        # attprojw: (L, C, C) — HF stores c_proj.weight as (C, C), transpose
        for i in range(L):
            w = get(f"transformer.h.{i}.attn.c_proj.weight")  # (C, C)
            f.write(fp32_to_bf16_int16(w.T.contiguous()).tobytes())
        # attprojb: (L, C)
        for i in range(L):
            f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.attn.c_proj.bias")).tobytes())

        # ln2w, ln2b
        for i in range(L):
            f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.ln_2.weight")).tobytes())
        for i in range(L):
            f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.ln_2.bias")).tobytes())

        # fcw: (L, 4C, C) — HF c_fc.weight is (C, 4C), transpose to (4C, C)
        for i in range(L):
            w = get(f"transformer.h.{i}.mlp.c_fc.weight")  # (C, 4C)
            f.write(fp32_to_bf16_int16(w.T.contiguous()).tobytes())
        # fcb: (L, 4C)
        for i in range(L):
            f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.mlp.c_fc.bias")).tobytes())

        # fcprojw: (L, C, 4C) — HF c_proj.weight is (4C, C), transpose to (C, 4C)
        for i in range(L):
            w = get(f"transformer.h.{i}.mlp.c_proj.weight")  # (4C, C)
            f.write(fp32_to_bf16_int16(w.T.contiguous()).tobytes())
        # fcprojb: (L, C)
        for i in range(L):
            f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.mlp.c_proj.bias")).tobytes())

        # lnfw, lnfb
        f.write(fp32_to_bf16_int16(get("transformer.ln_f.weight")).tobytes())
        f.write(fp32_to_bf16_int16(get("transformer.ln_f.bias")).tobytes())

    size_mb = __import__("os").path.getsize(output_path) / 1e6
    print(f"Saved {output_path}  ({size_mb:.1f} MB)")


def main():
    parser = argparse.ArgumentParser("HF GPT-2 → llm.c bf16 binary")
    parser.add_argument("--input",  "-i", required=True, help="HF model directory")
    parser.add_argument("--output", "-o", required=True, help="Output .bin path")
    args = parser.parse_args()

    print(f"Loading model from {args.input} ...")
    # Support loading when config and weights are in different snapshot dirs.
    # If args.input has no config.json, look for one in sibling snapshot dirs.
    import os, glob
    cfg_path = os.path.join(args.input, "config.json")
    if not os.path.exists(cfg_path):
        snapshots_dir = os.path.dirname(args.input)
        candidates = sorted(glob.glob(os.path.join(snapshots_dir, "*/config.json")))
        if not candidates:
            raise FileNotFoundError(f"No config.json found near {args.input}")
        cfg_path = candidates[0]
        print(f"Using config from: {cfg_path}")
    from transformers import GPT2Config
    config = GPT2Config.from_pretrained(os.path.dirname(cfg_path))
    model = GPT2LMHeadModel(config)
    from safetensors.torch import load_file
    weights_file = os.path.join(args.input, "model.safetensors")
    if os.path.exists(weights_file):
        sd = load_file(weights_file)
        model.load_state_dict(sd, strict=False)
        print(f"Loaded safetensors weights from {weights_file}")
    else:
        model = GPT2LMHeadModel.from_pretrained(args.input, config=config)
    model.eval()

    write_model(model, args.output)


if __name__ == "__main__":
    main()