"""
Stage 5 — EVALUATION (lecture: Evaluation, for pretraining AND post-training).

Compare base vs SFT vs DPO on three axes:
  1. Perplexity on the pretraining validation set  -> pure "language modeling" ability.
     Usually rises a bit after SFT/DPO (the "alignment tax").
  2. Instruction following on held-out prompts     -> the goal of post-training.
     Metrics: keyword hit rate (required words used), all-3-words rate, finish rate (story closed with EOS).
  3. Reward-hacking detectors — things the reward does NOT see:
     - fluency ppl: perplexity of the generated story under the frozen BASE model (lower = more natural text)
     - mentions/word: how often each required word appears (~1-2 natural, higher = keyword stuffing)
     - repeated trigrams: fraction of word trigrams that repeat (looping text)

Output: results/eval.json + results/eval.md (table + example generations).

Run:  python evaluate.py
"""
from __future__ import annotations

import json
import math
import os

import numpy as np
import torch

from common import (autocast_ctx, ensure_dir, get_device, keyword_hit_rate, keyword_mentions, load_model,
                    repeated_trigram_rate, set_seed)
from config import get_config
from data import TokenDataset
from generation import sample_responses
from pretrain import evaluate_loss
from tokenizer import BPETokenizer


@torch.no_grad()
def fluency_nll(base_model, tok: BPETokenizer, outs: list[dict], device: str, batch_size: int = 32) -> float:
    """Mean per-token NLL of the generated stories under the frozen base model (story only, no prompt)."""
    from dpo import seq_logprob

    ctx, _ = autocast_ctx(device)
    total, n_tok = 0.0, 0
    resp = [o["ids"] for o in outs if o["ids"]]
    for i in range(0, len(resp), batch_size):
        b = resp[i:i + batch_size]
        lp = seq_logprob(base_model, [[tok.eos_id]] * len(b), b, tok.eos_id, device, ctx)
        total += -lp.sum().item()
        n_tok += sum(len(r) for r in b)
    return total / max(n_tok, 1)


def evaluate_model(model, tok, cfg, val_prompts, base_model, device, n_samples_kept: int = 3):
    e = cfg.eval
    val_ds = TokenDataset(os.path.join(cfg.data.data_dir, "val.bin"), model.cfg.context_len)
    vl = evaluate_loss(model, val_ds, cfg.pretrain.batch_size, cfg.pretrain.eval_batches, device)

    set_seed(cfg.seed)  # identical sampling across models -> fair comparison
    _, outs = sample_responses(model, tok, val_prompts, 1, e.gen_max_tokens, e.temperature, e.top_k, device)
    outs = [o[0] for o in outs]
    hits = [keyword_hit_rate(o["text"], w) for o, w in zip(outs, val_prompts)]
    metrics = {
        "val_loss": vl,
        "val_ppl": math.exp(vl),
        "keyword_hit_rate": float(np.mean(hits)),
        "all_keywords_rate": float(np.mean([h == 1.0 for h in hits])),
        "finish_rate": float(np.mean([o["finished"] for o in outs])),
        "avg_tokens": float(np.mean([len(o["ids"]) for o in outs])),
        "fluency_ppl": math.exp(fluency_nll(base_model, tok, outs, device)),
        "mentions_per_word": float(np.mean([keyword_mentions(o["text"], w) for o, w in zip(outs, val_prompts)])),
        "repeated_trigrams": float(np.mean([repeated_trigram_rate(o["text"]) for o in outs])),
    }
    samples = [{"words": w, "text": o["text"]} for w, o in zip(val_prompts[:n_samples_kept], outs[:n_samples_kept])]
    return metrics, samples


def format_row(name: str, r: dict) -> str:
    return (f"| {name} | {r['val_ppl']:.2f} | {r['keyword_hit_rate']:.1%} | {r['all_keywords_rate']:.1%} | "
            f"{r['finish_rate']:.1%} | {r['fluency_ppl']:.2f} | {r['mentions_per_word']:.2f} | "
            f"{r['repeated_trigrams']:.1%} | {r['avg_tokens']:.0f} |")


TABLE_HEADER = ["| Model | Val ppl ↓ | Keyword hit ↑ | All 3 words ↑ | Finished ↑ | Fluency ppl ↓ | "
                "Mentions/word | Repeated trigrams ↓ | Avg tokens |",
                "|---|---|---|---|---|---|---|---|---|"]


def load_val_prompts(cfg) -> list[list[str]]:
    with open(os.path.join(cfg.data.data_dir, "sft_val.json")) as f:
        return [x["words"] for x in json.load(f)][: cfg.eval.n_prompts]


def main() -> None:
    cfg = get_config()
    e = cfg.eval
    device = get_device()
    tok = BPETokenizer.load(os.path.join(cfg.data.data_dir, "tokenizer.json"))
    val_prompts = load_val_prompts(cfg)
    base_model, _ = load_model(cfg.pretrain.ckpt_path, device)
    base_model.eval()

    ckpts = {"base": cfg.pretrain.ckpt_path, "sft": cfg.sft.ckpt_path, "dpo": cfg.dpo.ckpt_path}
    results, samples = {}, {}
    for name, path in ckpts.items():
        if not os.path.exists(path):
            print(f"Skipping {name}: {path} not found")
            continue
        model, _ = load_model(path, device)
        model.eval()
        results[name], samples[name] = evaluate_model(model, tok, cfg, val_prompts, base_model, device)
        r = results[name]
        print(f"{name:5s} | ppl {r['val_ppl']:6.2f} | keyword hit {r['keyword_hit_rate']:.1%} | "
              f"all-3-words {r['all_keywords_rate']:.1%} | finished {r['finish_rate']:.1%} | "
              f"fluency ppl {r['fluency_ppl']:.2f} | mentions/word {r['mentions_per_word']:.2f} | "
              f"rep-trigrams {r['repeated_trigrams']:.1%} | {r['avg_tokens']:.0f} tok")
        del model
        if device == "cuda":
            torch.cuda.empty_cache()

    ensure_dir("results/eval.json")
    with open("results/eval.json", "w") as f:
        json.dump({"n_prompts": len(val_prompts), "temperature": e.temperature, "top_k": e.top_k,
                   "metrics": results, "samples": samples}, f, indent=2)

    lines = ["# Evaluation results", "",
             f"Held-out prompts: {len(val_prompts)} | temperature {e.temperature} | top-k {e.top_k}", "",
             *TABLE_HEADER]
    lines += [format_row(n, r) for n, r in results.items()]
    lines += ["", "## Example generations", ""]
    for n, ss in samples.items():
        lines.append(f"### {n}")
        for s in ss:
            lines += [f"**Required words:** {', '.join(s['words'])}", "", f"> {s['text'].strip()[:600]}", ""]
    with open("results/eval.md", "w") as f:
        f.write("\n".join(lines))
    print("\nSaved: results/eval.json, results/eval.md")


if __name__ == "__main__":
    main()