"""Pack a checkpoint for the browser: int8 weights with one scale per row.

    python export.py runs/chat/ckpt_final.pt ../../public/llm/models/sprout-chat.bin

File layout: b'SPRT', uint32 version, uint32 header length, JSON header,
then the tensors, each aligned to 16 bytes. A matrix W (rows x cols) is stored
as int8 q plus float32 scales s so that W[r] ≈ q[r] * s[r].
"""
import argparse
import json
import struct

import numpy as np
import torch

MAGIC, VERSION = b'SPRT', 1


def quantize_rows(w):
    scale = np.abs(w).max(axis=1) / 127.0
    scale[scale == 0] = 1.0
    q = np.clip(np.round(w / scale[:, None]), -127, 127).astype(np.int8)
    return q, scale.astype(np.float32)


def pack(state, config, meta, path, quantize=True):
    tensors, blobs, offset = [], [], 0

    def add(name, arr, dtype):
        nonlocal offset
        data = arr.astype(dtype).tobytes()
        pad = (-len(data)) % 16
        tensors.append({'name': name, 'dtype': np.dtype(dtype).name, 'shape': list(arr.shape), 'offset': offset})
        blobs.append(data + b'\0' * pad)
        offset += len(data) + pad

    for name, t in state.items():
        if name == 'head.weight':             # tied to embed.weight
            continue
        w = t.detach().float().cpu().numpy()
        if quantize and w.ndim == 2:
            q, s = quantize_rows(w)
            add(name, q, np.int8)
            add(name + '.scale', s, np.float32)
        else:
            add(name, w, np.float32)
    header = json.dumps({'config': config, 'meta': meta, 'tensors': tensors}).encode()
    header += b' ' * ((-(12 + len(header))) % 16)
    with open(path, 'wb') as f:
        f.write(MAGIC + struct.pack('<II', VERSION, len(header)) + header)
        for b in blobs:
            f.write(b)
    return 12 + len(header) + offset


def dequantized(state):
    """The weights exactly as the browser will see them, to measure the damage."""
    out = {}
    for name, t in state.items():
        w = t.detach().float().cpu()
        if w.ndim == 2 and name != 'head.weight':
            q, s = quantize_rows(w.numpy())
            w = torch.from_numpy(q.astype(np.float32) * s[:, None])
        out[name] = w
    out['head.weight'] = out['embed.weight']
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('ckpt')
    ap.add_argument('out')
    ap.add_argument('--name', default='')
    ap.add_argument('--float', action='store_true', help='keep float32 (tiny models)')
    args = ap.parse_args()
    ck = torch.load(args.ckpt, map_location='cpu')
    meta = {'name': args.name, 'step': ck.get('step'), 'tokens': ck.get('tokens'), 'kind': 'gpt'}
    size = pack(ck['model'], ck['config'], meta, args.out, quantize=not args.float)
    print(f'{args.out}: {size / 1e6:.1f} MB')


if __name__ == '__main__':
    main()