"""Full holdout evaluation (single GPU).

python scripts/08_eval.py --rows data/manifests/split_test.jsonl --ckpt runs/round0/best.pt
"""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import argparse
import json
import torch
from src.common import ROOT, load_jsonl, save_jsonl
from src import training


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--rows", required=True)
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--patch", type=int, default=96)
    ap.add_argument("--out", default=None)
    args = ap.parse_args()
    device = "cuda:0"
    model, sd = training.load_model(args.ckpt, device)
    rows = load_jsonl(args.rows)
    d, per = training.evaluate(model, rows, device, args.patch, tta=True)
    res = {"ckpt": args.ckpt, "n": len(per), "dice": d, "epoch": sd.get("epoch"),
           "per_row": per}
    print(json.dumps({k: res[k] for k in ("ckpt", "n", "dice", "epoch")}, indent=1))
    out = args.out or os.path.join(ROOT, "results", os.path.basename(os.path.dirname(args.ckpt)) + "_test.json")
    os.makedirs(os.path.dirname(out), exist_ok=True)
    with open(out, "w") as f:
        json.dump(res, f, indent=1)
    save_jsonl(per, out.replace(".json", "_per_row.jsonl"))
    print("saved", out)


if __name__ == "__main__":
    main()