"""Arm C checkpoint -> armC.bin for cgemma.c (CPU kernel for the Gemma-3-270M shape with a ternary BitNet body and the
frozen 512-bit table). Every ternary tensor is packed exactly as the trainer's BitLinear rounds it (round-trip checked).

  python kernel/export_arm_c.py finals/C/step_07820.pt kernel/armC.bin
"""
import struct, sys, numpy as np, torch
sys.path.insert(0, ".")
import train

V, BITS, D, L, NH, NKV, HD, FFN = 262144, 512, 640, 18, 4, 1, 256, 2048
ckpt, out = sys.argv[1], sys.argv[2]
sd = torch.load(ckpt, map_location="cpu", weights_only=False)["model"]
f = open(out, "wb")
f.write(b"CGEM"); f.write(struct.pack("<8i", V, BITS, D, L, NH, NKV, HD, FFN))

def f32(t): np.ascontiguousarray(t.detach().float().numpy(), dtype=np.float32).tofile(f)
def tern(w):
    """BitLinear: sw = 1/mean|w|; code = round(w*sw).clamp(-1,1); kernel value = code / sw. Packed 2 bits, code+1, row-major."""
    w = w.detach().float(); sw = 1.0 / w.abs().mean().clamp_min(1e-5)
    code = (w * sw).round().clamp(-1, 1)
    assert torch.allclose(code / sw, train.ste(w, (w * sw).round().clamp(-1, 1) / sw), atol=1e-6, rtol=0), "round-trip"   # trainer value = w + (q - w) in fp32
    q = (code + 1).to(torch.uint8).numpy().reshape(-1)
    packed = (q[0::4] | (q[1::4] << 2) | (q[2::4] << 4) | (q[3::4] << 6)).astype(np.uint8)
    packed.tofile(f); np.full(w.shape[0], float(1.0 / sw), np.float32).tofile(f)

bits = ((sd["table"] > 0).to(torch.uint8).numpy())                     # +1 -> bit 1
np.packbits(bits, axis=1, bitorder="little").tofile(f)                 # [V, 64] bytes; byte b bit j <-> dim 8b+j
f32(sd["seal"]); f32(sd["inp.weight"]); f32(sd["out.weight"])          # seal [V]; inp [640,512]; out [512,640]
f32(1.0 + sd["body.norm.weight"])
for l in range(L):
    p = f"body.layers.{l}."
    for n in ("input_layernorm", "post_attention_layernorm", "pre_feedforward_layernorm", "post_feedforward_layernorm"):
        f32(1.0 + sd[p + n + ".weight"])
    f32(1.0 + sd[p + "self_attn.q_norm.weight"]); f32(1.0 + sd[p + "self_attn.k_norm.weight"])
    for n in ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj"):
        tern(sd[p + n + ".weight"])
f.close()
import os; print("wrote", out, f"{os.path.getsize(out)/1e6:.1f} MB")