"""Generate text from an exported checkpoint (quick quality check)."""
import argparse

import torch

from ar.model.config import ARConfig
from ar.model.model import ARModel
from ar.model.mask import build_doc_block_mask
from ar.model.rope import build_rope_cache
from ar.data.tokenizer import Tokenizer

SEQ = 256  # fixed length, so FlexAttention compiles once instead of for every new length


def load_model(path, device):
    cfg = ARConfig()
    with torch.device(device):
        model = ARModel(cfg)                     # real init also builds correct RoPE tables
    model.load_state_dict(torch.load(path, map_location=device))
    model = model.to(torch.bfloat16).eval()    
    return model


@torch.no_grad()
def generate(model, tok, prompt, device, max_new=150, temperature=0.8, top_k=50):
    ids = tok.encode(prompt)
    n = len(ids)
    # fixed-size buffer; slots after the current position are placeholders
    # that the causal mask hides from earlier tokens
    buf = torch.full((1, SEQ), tok.eot_id, dtype=torch.long, device=device)
    buf[0, :n] = torch.tensor(ids, device=device)
    seg = torch.zeros(1, SEQ, dtype=torch.int32, device=device)   # one document
    mask = build_doc_block_mask(seg, SEQ, device=device)

    pos = n
    for pos in range(n, min(SEQ, n + max_new)):
        logits = model(buf, mask)[0, pos - 1].float() / temperature
        # keep only the top_k most likely tokens, then sample one of them
        kth = torch.topk(logits, top_k).values[-1]
        logits[logits < kth] = -float("inf")
        nxt = torch.multinomial(torch.softmax(logits, -1), 1)
        buf[0, pos] = nxt
        if nxt.item() == tok.eot_id:              # model decided the text is over
            break
    return tok.decode(buf[0, : pos + 1].tolist())


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True)
    args = p.parse_args()

    device = "cuda"
    tok = Tokenizer()
    model = load_model(args.model, device)
    torch.manual_seed(0)

    prompts = [
        "The capital of France is",
        "Once upon a time, there was a little",
        "def fibonacci(n):",
        "Question: What is 12 + 7?\nAnswer:",
    ]
    for pr in prompts:
        print("=" * 60, flush=True)
        print(generate(model, tok, pr, device), flush=True)