"""Rust eval: pass@1 vs k with the frozen-prompt loop (transfer or trained)."""

import argparse
import json
import os
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import torch

from loop_common import BandLooper, MergeAdapter
from rust_common import DIRECT_SUFFIX, extract_rust, run_rust_tests, rust_prompt

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from jlens.core import load_model  # noqa: E402

OUT = Path(os.environ.get("LOOP_OUT",
    Path(__file__).resolve().parent.parent / "results-loop"))


@torch.no_grad()
def pass1(looper, adapter, tok, items, k, batch=8, max_new=300):
    codes = []
    for i in range(0, len(items), batch):
        torch.cuda.empty_cache()
        chunk = items[i : i + batch]
        enc = tok([rust_prompt(tok, it, DIRECT_SUFFIX) for it in chunk],
                  return_tensors="pt", padding=True,
                  add_special_tokens=False).to("cuda")
        gen = looper.generate_frozen_prompt(adapter, tok, enc["input_ids"], k,
                                            max_new_tokens=max_new,
                                            attention_mask=enc["attention_mask"])
        for j in range(len(chunk)):
            codes.append(extract_rust(
                tok.decode(gen[j, enc["input_ids"].shape[1]:],
                           skip_special_tokens=True)))
    with ThreadPoolExecutor(16) as ex:
        oks = list(ex.map(lambda ci: run_rust_tests(ci[0], ci[1]),
                          zip(codes, items)))
    by = {}
    for it, ok in zip(items, oks):
        d = by.setdefault(it["label"], [0, 0])
        d[0] += ok
        d[1] += 1
    per_item = [{"task_id": it["task_id"], "ok": bool(ok)}
                for it, ok in zip(items, oks)]
    return sum(oks) / len(items), {l: c / n for l, (c, n) in by.items()}, per_item


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--adapter", default=None)
    ap.add_argument("--tag", default="rust")
    ap.add_argument("--ks", default="0,2,4")
    ap.add_argument("--split", default="test")
    args = ap.parse_args()
    ks = [int(x) for x in args.ks.split(",")]

    model, tok = load_model(dtype=torch.bfloat16)
    tok.padding_side = "left"
    looper = BandLooper(model)
    adapter = MergeAdapter(
        d=model.config.get_text_config().hidden_size).cuda()
    if args.adapter:
        adapter.load_state_dict(torch.load(args.adapter, map_location="cuda"))
        adapter.eval()

    items = [it for it in json.load(open(OUT / "rust_data.json"))
             if it["split"] == args.split]
    print(f"[{args.tag}] Rust eval on {len(items)} items, ks={ks}", flush=True)
    res = {"tag": args.tag, "ks": {}, "n": len(items)}
    for k in ks:
        t0 = time.time()
        acc, by, per_item = pass1(looper, adapter, tok, items, k)
        res["ks"][k] = {"acc": acc, "by_label": by, "per_item": per_item}
        print(f"k={k}: pass@1={acc:.3f}  "
              f"by_label={ {l: round(v,3) for l,v in by.items()} }"
              f"  ({time.time()-t0:.0f}s)", flush=True)
    json.dump(res, open(OUT / f"eval_rust_{args.tag}.json", "w"), indent=1)
    print("wrote", OUT / f"eval_rust_{args.tag}.json")


if __name__ == "__main__":
    main()