"""Evaluate a saved decision model folder on rows from a jsonl file, with the base read from local disk.

    python extra_eval.py --folder /content/out/final --base /content/models/Qwen3.8-27B \
        --rows /content/out/heldout2_rows.jsonl --name final_heldout2
"""

import argparse
import json
import os
import shutil
import sys
import time

parser = argparse.ArgumentParser()
parser.add_argument("--folder", default = "/content/out/final")
parser.add_argument("--base", default = "/content/models/Qwen3.8-27B")
parser.add_argument("--rows", required = True)
parser.add_argument("--name", required = True)
parser.add_argument("--out", default = "/content/out")
args = parser.parse_args()

# A copy whose configs name the local base, so nothing is downloaded again.
local = args.folder.rstrip("/") + "_local"
shutil.rmtree(local, ignore_errors = True)
shutil.copytree(args.folder, local)
for name, key in (("unsloth_decision_config.json", "base_model"), ("adapter_config.json", "base_model_name_or_path")):
    path = os.path.join(local, name)
    config = json.load(open(path, encoding = "utf-8"))
    config[key] = args.base
    json.dump(config, open(path, "w", encoding = "utf-8"), indent = 2)

from unsloth import FastDecisionModel  # noqa: E402

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from td_metrics import parsed, score_cases  # noqa: E402

start = time.time()
model, tokenizer = FastDecisionModel.from_pretrained(local)
FastDecisionModel.for_inference(model)
print("loaded", f"{time.time() - start:.0f}s", "temperatures", model.decision_config.get("temperature"), flush = True)
rows = [json.loads(line) for line in open(args.rows, encoding = "utf-8") if line.strip()]
cases = []
for i, row in enumerate(rows, 1):
    answers = FastDecisionModel.predict(model, tokenizer, parsed(row["state"]), parsed(row["questions"]))
    cases.append(
        {
            "id": row["id"],
            "workflow": row["workflow"],
            "questions": row["questions"],
            "gold": row["gold"],
            "pred": {qid: {k: v for k, v in a.items() if k != "legend"} for qid, a in answers.items()},
        }
    )
    if i % 100 == 0:
        print(f"predicted {i}/{len(rows)}", flush = True)
metrics = score_cases(cases)
with open(f"{args.out}/{args.name}_preds.jsonl", "w", encoding = "utf-8") as f:
    for case in cases:
        f.write(json.dumps(case) + "\n")
with open(f"{args.out}/{args.name}_metrics.json", "w", encoding = "utf-8") as f:
    json.dump(metrics, f, indent = 1)
print("EXTRA_METRICS " + json.dumps(metrics), flush = True)