"""Evaluate a trained task checkpoint on its labelled public held-out split.

Accepts either a complete training checkpoint or an inference-only release.
No optimizer or original pretrained checkpoint is needed for evaluation.
"""
from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
import sys

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT / "native"))
import torch
from model_bio import GPT, GPTConfig
from task_runtime import check_gpu
from train_tasks import TaskModel, evaluate


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--split", choices=["validation", "test"])
    parser.add_argument("--batch-size", type=int)
    parser.add_argument("--max-process-gib", type=float, default=9.0)
    parser.add_argument("--output", type=Path, help="Optional JSON with scores and every prediction.")
    args = parser.parse_args()
    if args.max_process_gib <= 0 or args.batch_size is not None and args.batch_size <= 0:
        parser.error("Memory cap and batch size must be positive.")
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA required; CPU model execution is disabled.")
    check_gpu()
    torch.set_num_threads(2)
    torch.cuda.set_device(0)
    free, total = torch.cuda.mem_get_info()
    budget = min(int(args.max_process_gib * 1024**3), total - int(2.5 * 1024**3),
                 free - int(0.5 * 1024**3))
    if budget < 2 * 1024**3:
        raise RuntimeError("Less than 2 GiB is available after reserving desktop memory.")
    torch.cuda.set_per_process_memory_fraction(budget / total, 0)
    torch.set_float32_matmul_precision("high")
    saved = torch.load(args.checkpoint, map_location="cpu", weights_only=False, mmap=True)
    task = saved["task"]
    if task not in ("sst2", "race_middle"):
        raise ValueError(f"Unknown task: {task}")
    split = args.split or ("validation" if task == "sst2" else "test")
    if task == "sst2" and split == "test":
        raise ValueError("SST-2 public test labels are hidden. Use validation for labelled evaluation.")
    source_hash = saved.get("source_hashes", saved.get("provenance", {}).get("source_hashes", {})).get("model")
    if source_hash:
        actual_hash = hashlib.sha256((ROOT / "native" / "model_bio.py").read_bytes()).hexdigest()
        if source_hash != actual_hash:
            raise ValueError("Native model source differs from the checkpoint's recorded source.")
    config = GPTConfig(**saved["model_args"])
    config.gradient_checkpointing = False
    dropout = saved.get("config", {}).get("classifier_dropout", 0.1)
    with torch.device("cuda:0"):
        base = GPT(config)
        model = TaskModel(base, task, dropout)
    model.load_state_dict(saved["model"], strict=True)
    del saved
    cache_path = ROOT / "data" / f"{task}_{split}.pt"
    cache = torch.load(cache_path, map_location="cpu", weights_only=False, mmap=True)
    data = {key: cache[key].to("cuda:0") for key in ("input_ids", "lengths", "labels")}
    if torch.any(data["labels"] < 0):
        raise ValueError("This split has withheld labels and cannot be scored.")
    result = evaluate(model, data, args.batch_size or (16 if task == "sst2" else 4))
    result.update(task=task, split=split, checkpoint=str(args.checkpoint),
                  example_ids=cache.get("example_ids"), source_sha256=cache.get("source_sha256"),
                  device=torch.cuda.get_device_name(0))
    print(json.dumps({key: value for key, value in result.items()
                      if key not in ("predictions", "example_ids")}, indent=2))
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()