#!/usr/bin/env python3
"""Standalone honest eval of any checkpoint (or base model) on validation/test."""
import argparse
import os
import sys
import time
from pathlib import Path

BASE = Path(os.environ.get("QYVOS_HOME", "/home/z/my-project/download/qyvos"))
sys.path.insert(0, str(BASE / "Julia-1"))
sys.path.insert(0, str(BASE / "scripts"))

import torch
from transformers import AutoTokenizer

from julia.model import JuliaDecisionModel


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--split", default="validation", choices=["validation", "test"])
    ap.add_argument("--rows", type=int, default=900)
    ap.add_argument("--ckpt", default=None, help="head checkpoint to load (default: none = base)")
    ap.add_argument("--batch", type=int, default=8)
    args = ap.parse_args()

    import train_qyvos as T

    model = JuliaDecisionModel.from_pretrained(BASE / "Julia-1")
    tok = AutoTokenizer.from_pretrained(BASE / "Julia-1" / "tokenizer")
    if args.ckpt:
        payload = torch.load(args.ckpt, map_location="cpu", weights_only=False)
        model.load_state_dict(payload["model"], strict=False)
        print(f"loaded ckpt step={payload['step']} rows_seen={payload['rows_seen']}")

    class Cfg:
        eval_rows = args.rows
        eval_batch = args.batch
        max_length = 1024
        head_length = 512

    model.eval()
    t0 = time.time()
    acc, ce, detail = T.evaluate(model, tok, Cfg, split=args.split)
    print(f"split={args.split} rows={args.rows} acc={acc:.4f} softCE={ce:.4f} "
          + " ".join(f"{k}={v[0]:.4f}({v[1]})" for k, v in detail.items())
          + f" [{time.time()-t0:.0f}s]")


if __name__ == "__main__":
    main()