#!/usr/bin/env python3
"""
Generate text samples from a trained Teensy checkpoint.

Copyright (c) 2025 Pankaj Doharey
Modified from NanoGPT (Andrej Karpathy).
"""

import os
import argparse
import secrets
from contextlib import nullcontext

import torch
import tiktoken

from model import TeensyConfig, TeensyLM, adapt_nanogpt_weights


def load_model(checkpoint_path, device):
    checkpoint = torch.load(checkpoint_path, map_location=device)
    cfg = TeensyConfig(**checkpoint['model_args'])
    model = TeensyLM(cfg)

    state_dict = checkpoint['model']
    unwanted_prefix = '_orig_mod.'
    for k, v in list(state_dict.items()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
    state_dict = adapt_nanogpt_weights(state_dict)
    model.load_state_dict(state_dict)

    model.eval()
    model.to(device)
    return model


def resolve_device(requested):
    """Pick an available device and warn about MPS quality issues."""
    if requested == "cuda":
        if torch.cuda.is_available():
            return "cuda"
        fallback = "mps" if torch.backends.mps.is_available() else "cpu"
        print(f"CUDA not available; falling back to {fallback}.")
        return fallback
    if requested == "mps":
        if torch.backends.mps.is_available():
            print("WARNING: MPS inference for this model can produce degraded output")
            print("         (garbled text and stray <|endoftext|> tokens). CPU is recommended.")
            return "mps"
        print("MPS not available; falling back to cpu.")
        return "cpu"
    return "cpu"


def main():
    parser = argparse.ArgumentParser(description="Generate text from a Teensy checkpoint")
    parser.add_argument("--out_dir", default="checkpoints", help="Checkpoint directory")
    parser.add_argument("--start", default="\n", help="Prompt text")
    parser.add_argument("--num_samples", type=int, default=10, help="Number of samples")
    parser.add_argument("--max_new_tokens", type=int, default=500, help="Tokens per sample")
    parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature")
    parser.add_argument("--top_k", type=int, default=50, help="Top-k sampling")
    parser.add_argument("--top_p", type=float, default=0.9, help="Nucleus (top-p) sampling")
    parser.add_argument("--device", default="cpu", help="Device (cpu/cuda/mps; cpu recommended)")
    parser.add_argument("--dtype", default="float16", help="Torch dtype")
    parser.add_argument("--compile", action="store_true", help="torch.compile the model")
    parser.add_argument("--seed", type=int, default=None, help="Random seed")
    args = parser.parse_args()

    args.device = resolve_device(args.device)
    device_type = "cuda" if "cuda" in args.device else "mps" if "mps" in args.device else "cpu"

    seed = args.seed if args.seed is not None else secrets.randbelow(2**32)
    torch.manual_seed(seed)
    if device_type == "cuda":
        torch.cuda.manual_seed(seed)

    ptdtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[args.dtype]
    ctx = nullcontext() if device_type in ["cpu", "mps"] else torch.amp.autocast(device_type=device_type, dtype=ptdtype)

    ckpt_path = os.path.join(args.out_dir, "teensy-0.pt")
    model = load_model(ckpt_path, args.device)

    if args.compile and device_type == "cuda":
        model = torch.compile(model)

    enc = tiktoken.get_encoding("gpt2")
    encode = lambda s: enc.encode(s, allowed_special={"<|endoftext|>"})
    decode = lambda ids: enc.decode(ids)
    eos_token_id = enc.eot_token

    if args.start.startswith("FILE:"):
        with open(args.start[5:], "r", encoding="utf-8") as f:
            args.start = f.read()

    start_ids = encode(args.start)
    x = torch.tensor(start_ids, dtype=torch.long, device=args.device)[None, ...]

    with torch.no_grad():
        with ctx:
            for _ in range(args.num_samples):
                y = model.generate(x, args.max_new_tokens, temperature=args.temperature, top_k=args.top_k, top_p=args.top_p, eos_token_id=eos_token_id)
                print(decode(y[0].tolist()))
                print("---------------")


if __name__ == "__main__":
    main()