"""Generate text from a trained checkpoint to verify English quality."""
import torch
import tiktoken
from config import Config
from model import build_model

def generate(model, cfg, prompt_ids, max_tokens=200, temperature=0.8, top_k=40):
    """Greedy or sampled generation from the base model (no Medusa, no cache)."""
    model.eval()
    enc = tiktoken.get_encoding("gpt2")
    ids = torch.tensor([prompt_ids], device="cuda", dtype=torch.long)
    
    with torch.no_grad():
        for _ in range(max_tokens):
            # Use last 1024 tokens as context (sliding window)
            x = ids[:, -1024:]
            h = model(x)
            logits = model.lm_head(h[:, -1])  # [1, V]
            
            if temperature == 0:
                next_id = logits.argmax(dim=-1, keepdim=True)
            else:
                logits = logits / temperature
                if top_k > 0:
                    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                    logits[logits < v[:, [-1]]] = float('-inf')
                probs = torch.softmax(logits, dim=-1)
                next_id = torch.multinomial(probs, 1)
            
            ids = torch.cat([ids, next_id], dim=1)
    
    return ids[0].tolist()

if __name__ == "__main__":
    import sys
    ckpt_path = sys.argv[1] if len(sys.argv) > 1 else "checkpoints/base/best.pt"
    
    ckpt = torch.load(ckpt_path, map_location="cuda")
    cfg = Config.v5_500m()
    model = build_model(cfg, "cuda")
    model.load_state_dict(ckpt["model"])
    model.eval()
    
    print(f"Loaded: {ckpt_path} (step {ckpt.get('step', '?')}, loss {ckpt.get('loss', '?'):.4f})")
    
    enc = tiktoken.get_encoding("gpt2")
    
    prompts = [
        "The quick brown fox",
        "In a galaxy far far away",
        "The capital of France is",
        "Once upon a time",
    ]
    
    for prompt in prompts:
        prompt_ids = enc.encode_ordinary(prompt)
        out_ids = generate(model, cfg, prompt_ids, max_tokens=100, temperature=0.7, top_k=40)
        text = enc.decode(out_ids)
        print(f"\n--- Prompt: {prompt!r} ---")
        print(text)