#!/usr/bin/env python3
"""
export_kronos_weights.py
========================
Companion to kronos_reference_capture.py. Dumps EVERY parameter and buffer of
Kronos-small and Kronos-Tokenizer-base to flat little-endian float32 .bin files,
plus a manifest (tensor name -> filename, shape), for loading in MQL5 via
FileReadArray. One-time, offline; no Python at inference.

    python export_kronos_weights.py --kronos_repo /path/to/Kronos --out ./kronos_weights

It also prints the full name+shape listing to stdout. Paste that back and the
MQL5 weight loader + stacks can be written against the exact layout.
"""

import argparse
import json
import os
import re
import sys

import numpy as np


def safe_name(name):
    """PyTorch names use dots; make a filesystem- and MQL5-friendly filename."""
    return re.sub(r"[^A-Za-z0-9]", "_", name)


def dump_module(mod, outdir, label):
    os.makedirs(outdir, exist_ok=True)
    manifest = {"label": label, "tensors": {}}
    tensors = list(mod.named_parameters()) + list(mod.named_buffers())
    total = 0
    print(f"\n=== {label} : {len(tensors)} tensors ===")
    for name, t in tensors:
        arr = np.ascontiguousarray(t.detach().cpu().float().numpy().astype(np.float32))
        fn = safe_name(name) + ".bin"
        arr.tofile(os.path.join(outdir, fn))
        manifest["tensors"][name] = {
            "file": fn,
            "shape": list(arr.shape),
            "dtype": "float32",
            "count": int(arr.size),
        }
        total += int(arr.size)
        print(f"  {name:55s} shape={tuple(arr.shape)}")
    manifest["total_params"] = total
    with open(os.path.join(outdir, "manifest.json"), "w") as f:
        json.dump(manifest, f, indent=2)
    print(f"  -> {total:,} values, manifest at {os.path.join(outdir, 'manifest.json')}")
    return manifest


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--kronos_repo", default=".")
    ap.add_argument("--out", default="./kronos_weights")
    ap.add_argument("--tokenizer_id", default="NeoQuasar/Kronos-Tokenizer-base")
    ap.add_argument("--model_id", default="NeoQuasar/Kronos-small")
    args = ap.parse_args()

    sys.path.insert(0, os.path.abspath(args.kronos_repo))
    import torch
    from model import Kronos, KronosTokenizer

    torch.set_grad_enabled(False)
    tok = KronosTokenizer.from_pretrained(args.tokenizer_id).float().eval()
    mdl = Kronos.from_pretrained(args.model_id).float().eval()

    os.makedirs(args.out, exist_ok=True)
    dump_module(tok, os.path.join(args.out, "tokenizer"), "Kronos-Tokenizer-base")
    dump_module(mdl, os.path.join(args.out, "predictor"), "Kronos-small")

    print("\nAll weights exported (float32, row-major, little-endian).")
    print("Send back: config_tokenizer.json (from the capture script) + the two")
    print("manifest.json name/shape listings, and the MQL5 stacks can be written.")


if __name__ == "__main__":
    main()