"""
Generate text from a checkpoint with constant-memory decode.

    python sample.py --run runs/hoard_small_XXXX --prompt "The dragon" --tokens 100
"""
import argparse, json, os
import mlx.core as mx
from model import HOARD, HoardConfig


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", required=True)
    ap.add_argument("--prompt", default="\n")
    ap.add_argument("--tokens", type=int, default=100)
    ap.add_argument("--temperature", type=float, default=0.8)
    ap.add_argument("--top_k", type=int, default=50)
    ap.add_argument("--loops", type=int, default=None, help="try fewer/more loops than trained")
    ap.add_argument("--seed", type=int, default=0)
    a = ap.parse_args()

    meta = json.load(open(os.path.join(a.run, "config.json")))
    cfg = HoardConfig.from_dict(meta["config"])
    model = HOARD(cfg)
    model.load_weights(os.path.join(a.run, "model.safetensors"))
    model.eval()
    mx.random.seed(a.seed)

    import tiktoken
    enc = tiktoken.get_encoding("gpt2")
    ids = mx.array([enc.encode_ordinary(a.prompt)])
    out = model.generate(ids, max_new_tokens=a.tokens, temperature=a.temperature,
                         top_k=a.top_k, n_loops=a.loops)
    print(enc.decode(out[0].tolist()))


if __name__ == "__main__":
    main()