#!/usr/bin/env python3
"""Evaluate the released multi-task checkpoint on one CPU or one GPU."""

import argparse
import json
import os
import sys
from types import SimpleNamespace

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
REPO_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, ".."))
SRC_ROOT = os.path.join(REPO_ROOT, "src")
if SRC_ROOT not in sys.path:
    sys.path.insert(0, SRC_ROOT)
if SCRIPT_DIR not in sys.path:
    sys.path.insert(0, SCRIPT_DIR)

import torch

from clean_eval_utils import evaluate_clean_full, prepare_clean_records
from multi_task_model import MultiTaskGAT


def resolve_device(requested):
    if requested == "cpu":
        return torch.device("cpu")
    if not torch.cuda.is_available():
        if requested == "cuda":
            raise RuntimeError("CUDA was requested but is unavailable.")
        return torch.device("cpu")
    try:
        probe = torch.empty(1, device="cuda:0")
        del probe
        return torch.device("cuda:0")
    except RuntimeError as error:
        if requested == "cuda":
            raise
        print("CUDA initialization failed; using CPU: {}".format(error), flush=True)
        return torch.device("cpu")


def load_checkpoint(path, device):
    artifact = torch.load(path, map_location="cpu")
    if not isinstance(artifact, dict) or artifact.get("artifact_format") != "icflownet_model_v1":
        raise RuntimeError("Expected an icflownet_model_v1 checkpoint.")
    config = dict(artifact.get("model_config") or {})
    config.pop("architecture_version", None)
    model = MultiTaskGAT(**config)
    model.load_state_dict(artifact["model_state_dict"], strict=True)
    model.to(device)
    model.eval()
    return model, config


def evaluate_scope(args, scope, device):
    clean_root = os.path.abspath(args.clean_root)
    records, preflight = prepare_clean_records(
        clean_root, os.path.join(clean_root, "resources"), scope
    )
    if args.max_records is not None:
        records = records[: args.max_records]
    model, model_config = load_checkpoint(os.path.abspath(args.checkpoint), device)
    output_dir = os.path.join(os.path.abspath(args.output_dir), scope)
    os.makedirs(output_dir, exist_ok=True)
    metrics = evaluate_clean_full(
        SimpleNamespace(module=model),
        records,
        device,
        rank=0,
        world_size=1,
        output_dir=output_dir,
        log_every=args.log_every,
        amp_enabled=False,
    )
    result = {
        "format": "icflownet_cleantest_eval_v1",
        "scope": scope,
        "checkpoint": os.path.abspath(args.checkpoint),
        "checkpoint_artifact_format": "icflownet_model_v1",
        "threshold": 0.5,
        "clean_root": clean_root,
        "world_size": 1,
        "device": str(device),
        "precision": "fp32",
        "amp_enabled": False,
        "model_config": model_config,
        "preflight": preflight,
        "metrics": metrics,
    }
    result_path = os.path.join(output_dir, "result.json")
    with open(result_path, "w") as handle:
        json.dump(result, handle, indent=2)
        handle.write("\n")
    print(json.dumps({"result": result_path, "scope": scope, "metrics": metrics}, indent=2))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--clean_root", required=True)
    parser.add_argument("--scope", choices=("overall", "long_range", "both"), default="both")
    parser.add_argument("--output_dir", required=True)
    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
    parser.add_argument("--max_records", type=int)
    parser.add_argument("--log_every", type=int, default=10)
    args = parser.parse_args()
    if args.max_records is not None and args.max_records < 1:
        parser.error("--max_records must be at least 1")
    device = resolve_device(args.device)
    scopes = ("overall", "long_range") if args.scope == "both" else (args.scope,)
    for scope in scopes:
        evaluate_scope(args, scope, device)


if __name__ == "__main__":
    main()