"""
Try a model by hand.

  python chat.py --ckpt checkpoints/dpo.pt --words dragon,cookie,rain
  python chat.py --ckpt checkpoints/base.pt --text "Once upon a time"
"""
from __future__ import annotations

import argparse
import os

import torch

from common import build_prompt, get_device
from common import load_model
from config import get_config
from tokenizer import BPETokenizer


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="checkpoints/dpo.pt")
    ap.add_argument("--words", help="required words, comma-separated (instruction mode)")
    ap.add_argument("--text", help="prompt text (continuation mode, suited to the base model)")
    ap.add_argument("--n", type=int, default=2)
    ap.add_argument("--temperature", type=float, default=0.7)
    ap.add_argument("--max_tokens", type=int, default=200)
    args = ap.parse_args()

    cfg = get_config()
    device = get_device()
    tok = BPETokenizer.load(os.path.join(cfg.data.data_dir, "tokenizer.json"))
    model, mcfg = load_model(args.ckpt, device)
    model.eval()

    if args.words:
        prompt = build_prompt([w.strip() for w in args.words.split(",") if w.strip()])
    else:
        prompt = args.text or "Once upon a time"
    ids = tok.encode(prompt, allow_special=bool(args.words))
    x = torch.tensor([ids] * args.n, device=device)
    budget = min(args.max_tokens, mcfg.context_len - len(ids))
    out, _ = model.generate(x, budget, temperature=args.temperature, top_k=50, eos_id=tok.eos_id, valid_vocab=tok.vocab_size)
    for i in range(args.n):
        text = tok.decode([t for t in out[i].tolist() if t != tok.eos_id])
        print(f"\n--- sample {i + 1} ---\n{text.strip()}")


if __name__ == "__main__":
    main()