"""HumanEval (Python coding): complete each function, run its unit tests, pass@1 greedy."""
import argparse
import os
import subprocess
import sys
import tempfile

from datasets import load_dataset

from ar.data.tokenizer import Tokenizer
from ar.eval.batch_gen import generate_batch
from ar.eval.generate import load_model

# a new top-level statement means the function body is finished
STOP = ["\ndef ", "\nclass ", "\nif __name__", "\nprint(", "\n#"]


def passes(program, timeout=10):
    # run the model-written code in a separate process with a time limit,
    # so a crash or infinite loop can't take down the evaluation
    with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f:
        f.write(program)
        path = f.name
    try:
        return subprocess.run([sys.executable, path], capture_output=True,
                              timeout=timeout).returncode == 0
    except subprocess.TimeoutExpired:
        return False
    finally:
        os.remove(path)


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

    device = "cuda"
    tok = Tokenizer()
    model = load_model(args.model, device)
    ds = load_dataset("openai/openai_humaneval", split="test")

    solved = 0
    for i in range(0, len(ds), args.batch):
        chunk = ds.select(range(i, min(i + args.batch, len(ds))))
        outs = generate_batch(model, tok, [ex["prompt"] for ex in chunk], seq_len=1536,
                              max_new=384, stop=STOP, device=device)
        for ex, completion in zip(chunk, outs):
            program = (ex["prompt"] + completion + "\n\n" + ex["test"]
                       + f"\n\ncheck({ex['entry_point']})\n")
            solved += passes(program)
        done = min(i + args.batch, len(ds))
        print(f"[{done}/{len(ds)}] pass@1 so far: {100 * solved / done:.1f}%", flush=True)
        if i == 0:
            print("--- example completion ---\n" + chunk[0]["prompt"] + outs[0]
                  + "\n--------------------------", flush=True)

    print(f"HumanEval pass@1 (greedy): {100 * solved / len(ds):.1f}% ({solved}/{len(ds)})")