"""Generate from a trained checkpoint."""
import argparse
from pathlib import Path

import torch
from tokenizers import Tokenizer

from config import ModelConfig
from model import GPT


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--ckpt", type=str, default="checkpoints/best.pt")
    p.add_argument("--tokenizer", type=str, default="data/tokenizer.json")
    p.add_argument("--prompt", type=str, default="Once upon a time")
    p.add_argument("--max-new-tokens", type=int, default=256)
    p.add_argument("--temperature", type=float, default=0.8)
    p.add_argument("--top-k", type=int, default=200)
    p.add_argument("--num-samples", type=int, default=1)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--device", type=str, default=None)
    args = p.parse_args()

    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(args.seed)

    ckpt = torch.load(args.ckpt, map_location=device, weights_only=False)
    cfg_dict = ckpt["model_cfg"]
    valid = {f for f in ModelConfig.__dataclass_fields__}
    cfg = ModelConfig(**{k: v for k, v in cfg_dict.items() if k in valid})

    model = GPT(cfg).to(device).eval()
    model.load_state_dict(ckpt["model"])

    tok = Tokenizer.from_file(args.tokenizer)
    eot = tok.token_to_id("<|endoftext|>")

    ids = tok.encode(args.prompt).ids
    if not ids:
        ids = [eot]
    x = torch.tensor([ids], dtype=torch.long, device=device)

    for s in range(args.num_samples):
        out = model.generate(
            x, max_new_tokens=args.max_new_tokens,
            temperature=args.temperature, top_k=args.top_k, eos_id=eot,
        )[0].tolist()
        text = tok.decode(out)
        print(f"\n--- sample {s + 1} ---")
        print(text)


if __name__ == "__main__":
    main()