"""Evaluate Dory checkpoints with vLLM (Dory support from vllm-project/vllm#60149) on Modal.

  modal run modal_app/eval_dory.py --models base --loops 2 --limit 20      # smoke
  modal run modal_app/eval_dory.py --models base,runs/r1/final --loops 1,2,4

One H100 container per (model, loops) config, run in parallel. Each sampled answer is
scored with scripts/verify.py; per-sample records go to the beam-dory-data Volume under
eval/results/<model>/loops<L>.jsonl. Reports accuracy and output tokens per benchmark.
"""

import json
import os
import shutil
import sys
import time
from collections import defaultdict
from pathlib import Path

import modal

APP = modal.App("beam-dory-eval")
weights = modal.Volume.from_name("beam-weights", create_if_missing=True, version=2)
data = modal.Volume.from_name("beam-dory-data", create_if_missing=True, version=2)
ckpt = modal.Volume.from_name("beam-dory-ckpt", create_if_missing=True, version=2)
W, D, C = "/weights", "/data", "/ckpt"
VLLM_PR_SHA = "ce9dd0ed884b7cb4014396c1aa715a8b7c9fc4a4"  # vllm-project/vllm#60149 head, 2026-10-22

image = (
    # devel image: FlashInfer JIT-compiles kernels on first use and needs nvcc (CUDA 13, matching torch 2.14)
    modal.Image.from_registry("nvidia/cuda:13.0.3-devel-ubuntu24.04", add_python="3.12")
    .apt_install("git")
    .env({"VLLM_USE_PRECOMPILED": "1"})
    .pip_install(f"vllm @ git+https://github.com/vllm-project/vllm.git@{VLLM_PR_SHA}",
                 "math-verify[antlr4_13_2]")
    .pip_install("langdetect", "immutabledict", "nltk>=3.9.1")  # IFEval checkers (third_party/ifeval)
    .env({"PYTHONUNBUFFERED": "1", "TOKENIZERS_PARALLELISM": "false",
          "FLASHINFER_WORKSPACE_BASE": "/data/cache/flashinfer"})  # JIT cache persists on the Volume
    .add_local_file("scripts/verify.py", "/proj/verify.py")
    .add_local_dir("third_party/ifeval", "/proj/ifeval")
)

TOKENIZER_FILES = ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "chat_template.jinja")


def model_dir(name):
    return f"{W}/dory-4b" if name == "base" else f"{C}/{name}"


def load_llm(model, loops, max_tokens):
    from vllm import LLM

    path = model_dir(model)
    # Tokenizer-only dir: keeps AutoTokenizer from parsing Dory's (unregistered) model config.
    tok_dir = "/tmp/dory-tok"
    os.makedirs(tok_dir, exist_ok=True)
    for f in TOKENIZER_FILES:
        if os.path.exists(f"{path}/{f}"):
            shutil.copy(f"{path}/{f}", tok_dir)
    return LLM(model=path, tokenizer=tok_dir, max_model_len=max_tokens + 4096, gpu_memory_utilization=0.92,
               max_num_seqs=256, enable_prefix_caching=True,
               hf_overrides={"n_recurrent_loops": loops, "recurrent_kv_cache_mode": "per_loop"})


@APP.function(image=image, gpu="H100", volumes={W: weights, D: data, C: ckpt}, timeout=4 * 3600,
              cpu=8, memory=65536)
def evaluate(model: str, loops: int, limit: int = 0, max_tokens: int = 32768, medium_effort: bool = False,
             benches: str = "", shard: int = 0, num_shards: int = 1, reasoning_effort: str = ""):
    """num_shards > 1 evaluates every num_shards-th problem (interleaved, so each shard mixes benchmarks).
    reasoning_effort (low/medium/high) is passed to Dory-Beam's chat template, which tags the user turn."""
    sys.path.insert(0, "/proj")
    from verify import check
    from vllm import SamplingParams

    rows = [json.loads(l) for l in open(f"{D}/eval/evalsets.jsonl")]
    if benches:
        rows = [r for r in rows if r["bench"] in benches.split(",")]
    if limit:  # first `limit` problems of each benchmark
        per, kept = defaultdict(int), []
        for r in rows:
            per[r["bench"]] += 1
            if per[r["bench"]] <= limit:
                kept.append(r)
        rows = kept
    if num_shards > 1:
        rows = rows[shard::num_shards]

    t0 = time.time()
    llm = load_llm(model, loops, max_tokens)
    load_s = time.time() - t0

    by_k = defaultdict(list)
    for r in rows:
        by_k[r["samples"]].append(r)
    records = []
    t1 = time.time()
    kwargs = {"enable_thinking": True, "medium_effort": medium_effort}
    if reasoning_effort:
        kwargs["reasoning_effort"] = reasoning_effort
    for k, group in sorted(by_k.items()):
        sp = SamplingParams(n=k, temperature=1.0, top_p=0.95, max_tokens=max_tokens)
        outs = llm.chat([[{"role": "user", "content": r["prompt"]}] for r in group], sp,
                        chat_template_kwargs=kwargs, use_tqdm=False)
        for r, out in zip(group, outs):
            for j, o in enumerate(out.outputs):
                text = o.text
                answer = text.split("</think>")[-1] if "</think>" in text else ""
                ok, pred = check(answer, r["answer"], r["answer_type"])
                records.append({"bench": r["bench"], "id": r["id"], "sample": j, "correct": bool(ok),
                                "pred": pred, "tokens": len(o.token_ids), "finish": o.finish_reason,
                                "thinking_closed": "</think>" in text, "text": text})
    gen_s = time.time() - t1

    out_dir = Path(D) / "eval" / "results" / model.replace("/", "_")
    out_dir.mkdir(parents=True, exist_ok=True)
    tag = ("eff" if medium_effort else "") + (f"-{reasoning_effort}" if reasoning_effort else "")
    # split runs must not overwrite each other
    part = (f"-{benches.replace(',', '+')}" if benches else "") + (f"-s{shard}of{num_shards}" if num_shards > 1 else "")
    with open(out_dir / f"loops{loops}{tag}{'-limit' if limit else ''}{part}.jsonl", "w") as f:
        for rec in records:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
    data.commit()

    summary = {"model": model, "loops": loops, "medium_effort": medium_effort,
               "reasoning_effort": reasoning_effort, "load_s": round(load_s),
               "gen_s": round(gen_s), "tok_per_s": round(sum(r["tokens"] for r in records) / max(gen_s, 1)),
               "benches": {}}
    for b in sorted({r["bench"] for r in records}):
        rs = [r for r in records if r["bench"] == b]
        summary["benches"][b] = {"n": len(rs), "acc": round(sum(r["correct"] for r in rs) / len(rs), 4),
                                 "mean_tokens": round(sum(r["tokens"] for r in rs) / len(rs)),
                                 "truncated": round(sum(r["finish"] == "length" for r in rs) / len(rs), 3)}
    print("SUMMARY", json.dumps(summary), flush=True)
    return summary


@APP.function(image=image, gpu="H100", volumes={W: weights, D: data, C: ckpt}, timeout=4 * 3600,
              cpu=8, memory=65536)
def evaluate_ifeval(model: str, loops: int, samples: int = 2, max_tokens: int = 16384):
    """IFEval (google/IFEval, 541 prompts) with lm-evaluation-harness's checkers, scoring the text after
    </think>. Strict and loose, prompt- and instruction-level; a response with no closed think block fails."""
    sys.path.insert(0, "/proj")
    from huggingface_hub import hf_hub_download
    from ifeval.utils import InputExample, test_instruction_following_loose, test_instruction_following_strict
    from vllm import SamplingParams

    rows = [json.loads(l) for l in open(hf_hub_download("google/IFEval", "ifeval_input_data.jsonl",
                                                        repo_type="dataset"))]
    llm = load_llm(model, loops, max_tokens)
    t1 = time.time()
    sp = SamplingParams(n=samples, temperature=1.0, top_p=0.95, max_tokens=max_tokens)
    outs = llm.chat([[{"role": "user", "content": r["prompt"]}] for r in rows], sp,
                    chat_template_kwargs={"enable_thinking": True}, use_tqdm=False)
    records = []
    for r, out in zip(rows, outs):
        inp = InputExample(key=r["key"], instruction_id_list=r["instruction_id_list"], prompt=r["prompt"],
                           kwargs=r["kwargs"])
        for j, o in enumerate(out.outputs):
            answer = o.text.split("</think>")[-1].strip() if "</think>" in o.text else ""
            strict, loose = test_instruction_following_strict(inp, answer), test_instruction_following_loose(inp, answer)
            records.append({"key": r["key"], "sample": j, "tokens": len(o.token_ids), "finish": o.finish_reason,
                            "strict_prompt": strict.follow_all_instructions,
                            "strict_inst": strict.follow_instruction_list,
                            "loose_prompt": loose.follow_all_instructions,
                            "loose_inst": loose.follow_instruction_list, "text": o.text})
    gen_s = time.time() - t1

    out_dir = Path(D) / "eval" / "results" / model.replace("/", "_")
    out_dir.mkdir(parents=True, exist_ok=True)
    with open(out_dir / f"ifeval_loops{loops}.jsonl", "w") as f:
        for rec in records:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
    data.commit()

    inst = lambda k: [x for rec in records for x in rec[k]]
    summary = {"model": model, "loops": loops, "gen_s": round(gen_s), "n": len(records),
               "prompt_strict": round(sum(r["strict_prompt"] for r in records) / len(records), 4),
               "inst_strict": round(sum(inst("strict_inst")) / len(inst("strict_inst")), 4),
               "prompt_loose": round(sum(r["loose_prompt"] for r in records) / len(records), 4),
               "inst_loose": round(sum(inst("loose_inst")) / len(inst("loose_inst")), 4),
               "mean_tokens": round(sum(r["tokens"] for r in records) / len(records)),
               "truncated": round(sum(r["finish"] == "length" for r in records) / len(records), 3)}
    print("SUMMARY", json.dumps(summary), flush=True)
    return summary


@APP.function(image=image, gpu="H100", volumes={W: weights, D: data, C: ckpt}, timeout=4 * 3600,
              cpu=8, memory=65536)
def continue_truncated(model: str, loops: int, records: list, kind: str = "bench", max_total: int = 65536):
    """Resume answers that hit the original max_tokens from where they stopped, up to max_total output tokens.
    Same sampling settings, so each result is distributed as if the run had used max_total from the start.
    `records` are the truncated per-sample records from evaluate (kind="bench") or evaluate_ifeval
    (kind="ifeval"); returns and saves them with the continued text, re-scored."""
    sys.path.insert(0, "/proj")
    from vllm import SamplingParams, TokensPrompt

    if kind == "bench":
        from verify import check
        gold = {(r["bench"], r["id"]): r for r in map(json.loads, open(f"{D}/eval/evalsets.jsonl"))}
        prompt_of = lambda rec: gold[(rec["bench"], rec["id"])]["prompt"]
    else:
        from huggingface_hub import hf_hub_download
        from ifeval.utils import InputExample, test_instruction_following_loose, test_instruction_following_strict
        gold = {r["key"]: r for r in map(json.loads, open(hf_hub_download(
            "google/IFEval", "ifeval_input_data.jsonl", repo_type="dataset")))}
        prompt_of = lambda rec: gold[rec["key"]]["prompt"]

    llm = load_llm(model, loops, max_total)
    tok = llm.get_tokenizer()
    prompts, params, retok = [], [], []
    for rec in records:
        prefix = tok.apply_chat_template([{"role": "user", "content": prompt_of(rec)}], tokenize=False,
                                         add_generation_prompt=True, enable_thinking=True, medium_effort=False)
        gen_ids = tok.encode(rec["text"], add_special_tokens=False)
        retok.append(len(gen_ids) - rec["tokens"])  # re-tokenizing the text should give back ~the same tokens
        ids = tok.encode(prefix, add_special_tokens=False) + gen_ids
        prompts.append(TokensPrompt(prompt_token_ids=ids))
        params.append(SamplingParams(n=1, temperature=1.0, top_p=0.95, max_tokens=max_total - rec["tokens"]))
    t1 = time.time()
    outs = llm.generate(prompts, params, use_tqdm=False)
    gen_s = time.time() - t1

    out = []
    for rec, o in zip(records, outs):
        o = o.outputs[0]
        text = rec["text"] + o.text
        new = {**rec, "tokens": rec["tokens"] + len(o.token_ids), "finish": o.finish_reason, "text": text,
               "continued_tokens": len(o.token_ids)}
        answer = text.split("</think>")[-1] if "</think>" in text else ""
        if kind == "bench":
            ok, pred = check(answer, gold[(rec["bench"], rec["id"])]["answer"],
                             gold[(rec["bench"], rec["id"])]["answer_type"])
            new.update(correct=bool(ok), pred=pred, thinking_closed="</think>" in text)
        else:
            g = gold[rec["key"]]
            inp = InputExample(key=g["key"], instruction_id_list=g["instruction_id_list"], prompt=g["prompt"],
                               kwargs=g["kwargs"])
            strict = test_instruction_following_strict(inp, answer.strip())
            loose = test_instruction_following_loose(inp, answer.strip())
            new.update(strict_prompt=strict.follow_all_instructions, strict_inst=strict.follow_instruction_list,
                       loose_prompt=loose.follow_all_instructions, loose_inst=loose.follow_instruction_list)
        out.append(new)

    out_dir = Path(D) / "eval" / "results" / model.replace("/", "_")
    out_dir.mkdir(parents=True, exist_ok=True)
    name = f"{'ifeval_' if kind == 'ifeval' else ''}loops{loops}-cont{max_total // 1024}k.jsonl"
    with open(out_dir / name, "w") as f:
        for rec in out:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
    data.commit()
    summary = {"model": model, "loops": loops, "kind": kind, "n": len(out), "gen_s": round(gen_s),
               "still_truncated": sum(r["finish"] == "length" for r in out),
               "retok_diff_max": max(map(abs, retok), default=0), "prefix_tail": prefix[-40:]}
    print("SUMMARY", json.dumps(summary), flush=True)
    return summary


@APP.local_entrypoint()
def main(models: str = "base", loops: str = "2", limit: int = 0, medium_effort: bool = False, benches: str = ""):
    configs = [(m, int(l)) for m in models.split(",") for l in loops.split(",")]
    results = list(evaluate.starmap([(m, l, limit, 32768, medium_effort, benches) for m, l in configs]))
    Path("data/eval").mkdir(parents=True, exist_ok=True)
    stamp = time.strftime("%Y%m%d-%H%M%S")
    Path(f"data/eval/summary-{stamp}.json").write_text(json.dumps(results, indent=1))
    for s in results:
        line = "  ".join(f"{b}={v['acc']:.3f}/{v['mean_tokens']}t" for b, v in s["benches"].items())
        print(f"{s['model']:>24s} loops={s['loops']}: {line}")