"""GPT-U-20M evaluation: per-domain val loss + bits-per-byte, sample generations, optional lm-eval-harness.

    python scripts/evaluate.py --ckpt out/ckpt/step_19840.pt            # val loss / bpb -> out/eval.json
    python scripts/evaluate.py --hf out/hf --samples                   # + 10 generations -> out/samples.md
    python scripts/evaluate.py --hf out/hf --lm-eval                   # benchmarks -> out/lm_eval/

bpb = loss_nats * targets / (bytes * ln 2), with the UTF-8 bytes of exactly the evaluated target tokens (index.json).
"""
import argparse
import json
import math
import subprocess
import sys
from pathlib import Path

import numpy as np
import torch
import torch.distributed as dist
import torch.nn.functional as F

from common import CKPT_DIR, DATA_DIR, OUT_DIR, TOKENIZER_DIR, build_model

SEQ = 1024
EVAL_BATCH = 8  # 200 batches x 8 sequences = the 1,600 fixed windows of each holdout
EOT_ID = 0
LM_EVAL_TASKS = "arc_easy,piqa,hellaswag,lambada_openai,winogrande"
TEXT_PROMPTS = [
    "The history of the Roman Empire",
    "Photosynthesis is the process by which",
    "In a small village near the mountains,",
    "The most important thing to remember when learning a new language is",
    "Climate change affects the oceans because",
]
CODE_PROMPTS = [
    "def fibonacci(n):\n",
    "import numpy as np\n\n\ndef normalize(x):\n",
    "class Stack:\n    def __init__(self):\n",
    "function debounce(fn, wait) {\n",
    "#include <stdio.h>\n\nint main(void) {\n",
]


@torch.no_grad()
def val_losses(model: torch.nn.Module, device: torch.device, rank: int = 0, world: int = 1) -> dict:
    """Val loss and bits-per-byte on each holdout; windows are split across ranks and the sums all-reduced."""
    val = json.loads((DATA_DIR / "index.json").read_text())["val"]
    was_training = model.training
    model.eval()
    results, nats_all, targets_all, bytes_all = {}, 0.0, 0.0, 0
    for domain, v in val.items():
        arr = np.fromfile(DATA_DIR / v["file"], dtype=np.uint16)  # 4 MB; avoids mmap over DrvFs
        windows = np.arange(v["eval_windows"])[rank::world]
        acc = torch.zeros(2, dtype=torch.float64, device=device)
        for i in range(0, len(windows), EVAL_BATCH):
            rows = np.stack([arr[k * SEQ:k * SEQ + SEQ + 1] for k in windows[i:i + EVAL_BATCH]]).astype(np.int64)
            x = torch.from_numpy(rows).to(device)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                logits = model(input_ids=x[:, :-1], use_cache=False).logits
            acc[0] += F.cross_entropy(logits.float().flatten(0, 1), x[:, 1:].flatten(), reduction="sum").double()
            acc[1] += x[:, 1:].numel()
        if world > 1:
            dist.all_reduce(acc)
        nats, targets = acc.tolist()
        assert int(targets) == v["eval_targets"], f"{domain}: evaluated {targets} targets"
        results[domain] = {"loss": nats / targets, "bpb": nats / (v["eval_bytes"] * math.log(2))}
        nats_all, targets_all, bytes_all = nats_all + nats, targets_all + targets, bytes_all + v["eval_bytes"]
    results["aggregate"] = {"loss": nats_all / targets_all, "bpb": nats_all / (bytes_all * math.log(2))}
    model.train(was_training)
    return results


def load_model(ckpt: Path | None, hf_dir: Path | None, device: torch.device):
    if hf_dir:
        from transformers import AutoModelForCausalLM
        return AutoModelForCausalLM.from_pretrained(hf_dir, dtype=torch.float32).to(device)
    model = build_model(init=False)
    model.load_state_dict(torch.load(ckpt, map_location="cpu", weights_only=False)["model"])
    return model.to(device)


@torch.no_grad()
def write_samples(model, device: torch.device, path: Path) -> None:
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
    torch.manual_seed(0)
    lines = ["# GPT-U-20M samples", "",
             "Sampling: `temperature=0.8`, `top_p=0.95`, `max_new_tokens=200`, seed 0. "
             "Each prompt is prefixed with `<|endoftext|>` (the BOS used in training); every sample starts with its prompt.", ""]
    model.eval()
    for kind, prompts in (("Text", TEXT_PROMPTS), ("Code", CODE_PROMPTS)):
        for prompt in prompts:
            ids = torch.tensor([[EOT_ID] + tok.encode(prompt, add_special_tokens=False)], device=device)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                out = model.generate(ids, do_sample=True, temperature=0.8, top_p=0.95, max_new_tokens=200,
                                     eos_token_id=EOT_ID, pad_token_id=EOT_ID)
            completion = tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True)
            title = prompt.strip().splitlines()[0]
            lines += [f"## {kind}: `{title}`", "", "~~~~", prompt + completion, "~~~~", ""]
    path.write_text("\n".join(lines), encoding="utf-8")


def run_lm_eval(hf_dir: Path, out_dir: Path) -> None:
    subprocess.run([sys.executable, "-m", "lm_eval", "--model", "hf",
                    "--model_args", f"pretrained={hf_dir},dtype=bfloat16,add_bos_token=True", "--tasks", LM_EVAL_TASKS,
                    "--batch_size", "32", "--output_path", str(out_dir)], check=True)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", type=Path, help="training checkpoint (default: newest in out/ckpt)")
    ap.add_argument("--hf", type=Path, help="exported Hugging Face model directory")
    ap.add_argument("--samples", action="store_true")
    ap.add_argument("--lm-eval", action="store_true")
    args = ap.parse_args()
    device = torch.device("cuda", 0)
    ckpt = args.ckpt or (None if args.hf else sorted(CKPT_DIR.glob("step_*.pt"))[-1])
    model = load_model(ckpt, args.hf, device)

    results = val_losses(model, device)
    results["source"] = str(args.hf or ckpt)
    (OUT_DIR / "eval.json").write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))
    if args.samples:
        write_samples(model, device, OUT_DIR / "samples.md")
    if args.lm_eval:
        run_lm_eval(args.hf, OUT_DIR / "lm_eval")


if __name__ == "__main__":
    main()