"""Evaluate a checkpoint and print every seeded story continuation, unedited.

    python eval.py checkpoints/main/ckpt.pt
    python eval.py checkpoints/vocab8192/ckpt.pt --data data/vocab8192
"""

import argparse
import hashlib
import json

import torch

import tokenizer
from model import Sprout
from train import DEVICE, SEED, autocast, fingerprint, split, token_bytes, val_metrics

PROMPTS = (
    "Once upon a time",
    "Lily found a small red box under the old tree.",
    "Tom wanted to fly, but he did not have wings.",
    "The little rabbit was afraid of the dark.",
    "One rainy day, Anna heard a knock at the door.",
    "Ben and Mia were best friends. Today, they had an argument.",
    "A tiny dragon could only blow bubbles instead of fire.",
    "Sara planted a seed and waited for it to grow.",
    "Max lost his favorite blue hat at the park.",
    "The puppy saw its own face in the mirror.",
    "Lucy promised her mother that she would not open the door.",
    "Jack had three apples. He gave one apple to his sister.",
    "The snowman wanted to see what summer was like.",
    "Every morning, the little bird sang the same song. But today",
    "Nora found a key, but she did not know what it could open.",
    "The toy boat sank when Sam put a heavy stone inside it.",
    "A fox and a duck wanted to cross the river together.",
    "Emma was sad because nobody had remembered her birthday.",
    "Leo wore a red coat. His sister Ada wore a yellow coat. They traded coats.",
    "Rose put her marble in a green cup. While Rose was outside, Finn moved the "
    "marble into a blue cup. Rose came back and looked for her marble.",
)

def load(path, data, device):
    ckpt = torch.load(path, map_location="cpu", weights_only=False)
    assert ckpt["fingerprint"] == fingerprint(data), "checkpoint and data do not match"
    cfg = ckpt["config"]
    tok = tokenizer.load(f"{data}/tokenizer.model")
    assert len(tok.vocab) == cfg["vocab"]
    model = Sprout(cfg["n_layer"], cfg["dim"], cfg["n_head"], cfg["ctx"],
                   cfg["vocab"], cfg.get("pos", "rope"))
    model.load_state_dict(ckpt["model"])
    return model.to(device).eval(), tok, cfg, ckpt["step"]

@torch.inference_mode()
def continuation(model, tok, prompt, args, seed):
    torch.manual_seed(seed)
    if args.device == "mps":
        torch.mps.manual_seed(seed)
    eot = tok.special["<|endoftext|>"]
    prefix = [eot] + tok.encode(prompt)
    assert len(prefix) <= model.ctx, "prompt exceeds the model context"
    idx = torch.tensor([prefix], device=args.device)
    with autocast(args.device):
        ids = model.generate(idx, args.max_tokens, args.temperature, eot, args.top_k)
    ids = ids[0, len(prefix):].tolist()
    ended = eot in ids
    if ended:
        ids = ids[:ids.index(eot)]
    return tok.decode(ids), ids, ended

def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("checkpoint")
    p.add_argument("--data", default="data")
    p.add_argument("--device", choices=("cpu", "cuda", "mps"), default=DEVICE)
    p.add_argument("--micro-batch", type=int, default=8)
    p.add_argument("--seed", type=int, default=SEED)
    p.add_argument("--temperature", type=float, default=0.8)
    p.add_argument("--top-k", type=int, default=50)
    p.add_argument("--max-tokens", type=int, default=256)
    args = p.parse_args()
    assert args.micro_batch > 0 and args.max_tokens > 0
    assert args.temperature > 0 and args.top_k > 0
    model, tok, cfg, step = load(args.checkpoint, args.data, args.device)
    assert args.top_k <= cfg["vocab"]
    args.ctx = cfg["ctx"]
    with open(args.checkpoint, "rb") as f:
        digest = hashlib.file_digest(f, "sha256").hexdigest()
    print(json.dumps(dict(checkpoint=args.checkpoint, sha256=digest, step=step,
                          config=cfg, fingerprint=fingerprint(args.data),
                          torch=torch.__version__, device=args.device,
                          dtype="bfloat16" if args.device == "cuda" else "float32",
                          micro_batch=args.micro_batch)), flush=True)
    data = split(args.data, "val")
    nbytes = token_bytes(tok)
    loss, bpb = val_metrics(model, data, nbytes, args)
    print(json.dumps(dict(val_loss=loss, val_bpb=bpb, targets=len(data) - 1,
                          bytes=int(nbytes[data[1:]].sum()), context=args.ctx)), flush=True)
    for i, prompt in enumerate(PROMPTS):
        seed = args.seed + i
        text, ids, ended = continuation(model, tok, prompt, args, seed)
        if i == 0:
            repeated = continuation(model, tok, prompt, args, seed)
            assert (text, ids, ended) == repeated, "seeded sampling did not reproduce"
        print(json.dumps(dict(prompt_id=i + 1, prompt=prompt, continuation=text,
                              token_ids=ids, seed=seed, temperature=args.temperature,
                              top_k=args.top_k, max_tokens=args.max_tokens, use_cache=True,
                              generated_tokens=len(ids), ended=ended), ensure_ascii=False),
              flush=True)

if __name__ == "__main__":
    main()