#!/usr/bin/env python3
"""Export PyTorch tern_tc model to TC02 binary format for C inference."""

import argparse
import struct
import torch
import sys

def quantize_ternary(w):
    g = w.abs().mean().clamp(min=1e-5)
    wq = (w / g).round().clamp(-1, 1).to(torch.int8)
    return wq, g.item()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--ckpt", required=True, help="Path to PyTorch .pt model")
    parser.add_argument("--out", required=True, help="Path to output .bin file")
    args = parser.parse_args()

    print(f"Loading {args.ckpt}...")
    sd = torch.load(args.ckpt, map_location="cpu")
    if "model" in sd:
        sd = sd["model"]
        
    # Determine config from weights
    vocab_size, d_model = sd["wte.weight"].shape
    n_layer = max(int(k.split(".")[1]) for k in sd.keys() if k.startswith("blocks.")) + 1
    # Check layer 0 shapes
    q_out, q_in = sd["blocks.0.attn.q.weight"].shape
    k_out, k_in = sd["blocks.0.attn.k.weight"].shape
    gate_out, gate_in = sd["blocks.0.mlp.gate.weight"].shape
    
    assert q_in == d_model
    d_ff = gate_out
    
    # Calculate n_head and n_kv_head
    # We know from tern_tc: n_head=5, n_kv_head=1, head_dim = d_model / n_head
    # But wait, head_dim is usually d_model / n_head or explicit. 
    # Let's assume head_dim = 64 (standard for 320 d_model / 5)
    head_dim = 64
    n_head = q_out // head_dim
    n_kv_head = k_out // head_dim

    print(f"Detected Config: n_layer={n_layer}, n_head={n_head}, n_kv_head={n_kv_head}")
    print(f"                 d_model={d_model}, d_ff={d_ff}, head_dim={head_dim}, vocab={vocab_size}")

    with open(args.out, "wb") as f:
        # Magic: TC02
        f.write(b"TC02")
        # Header: n_layer, n_head, n_kv_head, d_model, d_ff, head_dim, vocab_size
        f.write(struct.pack("<7i", n_layer, n_head, n_kv_head, d_model, d_ff, head_dim, vocab_size))

        def write_f32(t):
            f.write(t.float().numpy().tobytes())

        def write_ternary(name):
            w = sd[name]
            wq, g = quantize_ternary(w)
            f.write(wq.numpy().tobytes())
            f.write(struct.pack("<f", g))

        print("Writing embeddings...")
        write_f32(sd["wte.weight"])

        for i in range(n_layer):
            print(f"Writing layer {i}...")
            prefix = f"blocks.{i}"
            
            write_f32(sd[f"{prefix}.n1.weight"])
            
            write_ternary(f"{prefix}.attn.q.weight")
            write_ternary(f"{prefix}.attn.k.weight")
            write_ternary(f"{prefix}.attn.v.weight")
            write_ternary(f"{prefix}.attn.o.weight")
            
            write_f32(sd[f"{prefix}.n2.weight"])
            
            write_ternary(f"{prefix}.mlp.gate.weight")
            write_ternary(f"{prefix}.mlp.up.weight")
            write_ternary(f"{prefix}.mlp.down.weight")

        print("Writing final norm...")
        write_f32(sd["norm.weight"])
        
    print(f"Exported to {args.out} successfully.")

if __name__ == "__main__":
    main()