#!/usr/bin/env python3
"""
Run the HuggingFace-exported Teensy model (safetensors + config.json).

This script does NOT depend on transformers; it loads the exported weights
back into the native TeensyLM architecture defined in model.py.

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

import os
import json
import argparse
import secrets
from contextlib import nullcontext

import torch
import tiktoken
from safetensors.torch import load_file

from model import TeensyConfig, TeensyLM


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 load_model(export_dir, device, dtype):
    device_type = 'cuda' if 'cuda' in device else 'mps' if 'mps' in device else 'cpu'

    with open(os.path.join(export_dir, 'config.json')) as f:
        cfg = TeensyConfig(**json.load(f))

    model = TeensyLM(cfg)
    state_dict = load_file(os.path.join(export_dir, 'model.safetensors'))
    model.load_state_dict(state_dict)
    model.eval()
    model.to(device)

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

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

    return model, encode, decode, ctx, eos_token_id


def generate(prompt, model, encode, decode, ctx, args, eos_token_id):
    ids = encode(prompt)
    x = torch.tensor(ids, dtype=torch.long, device=args.device)[None, ...]
    with torch.no_grad():
        with ctx:
            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,
            )
            return decode(y[0].tolist())


def main():
    parser = argparse.ArgumentParser(description='Run exported Teensy safetensors model')
    parser.add_argument('--export_dir', default='exported', help='Directory containing config.json and model.safetensors')
    parser.add_argument('--prompt', default='Hello', help='Input prompt')
    parser.add_argument('--num_samples', type=int, default=1, help='Number of samples')
    parser.add_argument('--temperature', type=float, default=0.8, help='Sampling temperature')
    parser.add_argument('--max_new_tokens', type=int, default=200, help='Number of tokens to generate')
    parser.add_argument('--top_k', type=int, default=200, help='Top-k sampling parameter')
    parser.add_argument('--top_p', type=float, default=0.9, help='Nucleus (top-p) sampling parameter')
    parser.add_argument('--device', default='cpu', help='Device to run on (cpu/cuda/mps; cpu recommended)')
    parser.add_argument('--dtype', default='float16', help='Data type for model')
    parser.add_argument('--seed', type=int, default=None, help='Random seed (omit for non-deterministic sampling)')
    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)

    print(f"Loading exported model from {args.export_dir}...")
    model, encode, decode, ctx, eos_token_id = load_model(args.export_dir, args.device, args.dtype)
    print(f"Model loaded! Running on {args.device}")
    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"number of parameters: {n_params/1e6:.2f}M")

    for sample_idx in range(args.num_samples):
        if args.num_samples > 1:
            print(f"\n--- sample {sample_idx + 1} ---")
        print(generate(args.prompt, model, encode, decode, ctx, args, eos_token_id))


if __name__ == '__main__':
    main()