#!/usr/bin/env python
from __future__ import annotations
import argparse
import json
from pathlib import Path
import torch
from sacflow.utils.config import load_yaml
from sacflow.utils.misc import seed_everything, ensure_dir
from sacflow.utils.distributed import init_distributed, cleanup, is_main_process
from sacflow.data.loader import build_loader
from sacflow.models.unet3d import build_model
from sacflow.engine.train_loop import evaluate


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    ap.add_argument("--checkpoint", default=None)
    ap.add_argument("--split", default=None)
    args = ap.parse_args()
    cfg = load_yaml(args.config)
    if args.checkpoint:
        cfg.setdefault("eval", {})["checkpoint"] = args.checkpoint
    if args.split:
        cfg.setdefault("eval", {})["split"] = args.split
    seed_everything(int(cfg.get("seed", 1337)))
    device = init_distributed(cfg.get("distributed", {}).get("backend", "nccl"))
    model = build_model(cfg).to(device)
    ckpt_path = cfg.get("eval", {}).get("checkpoint") or cfg.get("train", {}).get("source_checkpoint")
    if ckpt_path is None:
        ckpt_path = str(Path(cfg["output_dir"]) / "checkpoints" / "best.pt")
    ckpt = torch.load(ckpt_path, map_location="cpu")
    model.load_state_dict(ckpt.get("model", ckpt), strict=False)
    split = cfg.get("eval", {}).get("split", "target_test")
    loader = build_loader(cfg, split=split, training=False, require_label=True)
    metrics = evaluate(model, loader, cfg, device)
    if is_main_process():
        print(json.dumps(metrics, indent=2))
        out = ensure_dir(Path(cfg["output_dir"]) / "eval")
        with open(out / f"metrics_{split}.json", "w") as f:
            json.dump(metrics, f, indent=2)
    cleanup()

if __name__ == "__main__":
    main()