#!/usr/bin/env python
"""评估入口。

用法:
    python eval.py --config configs/eval/baseline.yaml --checkpoint runs/baseline/best.pth

设备同样支持 auto / cuda / mps / cpu；部署前回归可以显式 `--device cpu`。
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

from datasets import build_dataloader, build_dataset, load_classes
from engine import Evaluator
from models import build_model
from utils import check_label_mode, load_config, setup_logging
from utils.device import DeviceNotAvailableError, describe_device, resolve_device


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="薄膜缺陷模型评估")
    parser.add_argument("--config", required=True)
    parser.add_argument("--checkpoint", default=None, help="覆盖配置中的 checkpoint")
    parser.add_argument("--split", default=None, help="评估哪个 data 划分，默认取配置的 eval.split")
    parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE")
    parser.add_argument("--device", default=None, help="覆盖配置中的 device：auto / cuda / mps / cpu")
    parser.add_argument("--out", default=None, help="把指标写入 json 文件")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    cfg = load_config(args.config, args.set)
    logger = setup_logging()

    eval_cfg = dict(cfg.get("eval", {}))
    split = args.split or eval_cfg.pop("split", "val")
    eval_cfg.pop("split", None)

    checkpoint = args.checkpoint or cfg.get("checkpoint")
    if not checkpoint:
        logger.error("必须通过 --checkpoint 或配置中的 checkpoint 指定权重")
        return 1

    try:
        device = resolve_device(args.device or cfg.get("device", "auto"))
    except DeviceNotAvailableError as exc:
        logger.error("%s", exc)
        return 1
    logger.info("评估设备: %s", describe_device(device))
    classes = load_classes(cfg["classes"])

    # label_mode 由 yaml 决定，模型与数据集必须一致，这里统一下发。
    label_mode = check_label_mode(str(cfg.get("label_mode", "multilabel")))
    model_cfg = dict(cfg["model"])
    model_cfg.setdefault("num_classes", len(classes))
    model_cfg.setdefault("label_mode", label_mode)
    model = build_model(model_cfg)

    from engine import load_checkpoint

    load_checkpoint(checkpoint, model, map_location=device)
    model.to(device).eval()

    split_cfg = cfg["data"][split]
    dataset = build_dataset(
        split_cfg["dataset"], classes=classes, stride=model.out_stride, label_mode=label_mode
    )
    loader = build_dataloader(dataset, **{**dict(split_cfg.get("loader", {})), "shuffle": False})
    logger.info("评估 %s 集，共 %d 张图，权重 %s", split, len(dataset), checkpoint)

    metrics = Evaluator(model, loader, device, **eval_cfg).evaluate()
    for key, value in metrics.items():
        logger.info("%-16s %.4f", key, value)

    if args.out:
        Path(args.out).parent.mkdir(parents=True, exist_ok=True)
        Path(args.out).write_text(json.dumps(metrics, indent=2, ensure_ascii=False), encoding="utf-8")
        logger.info("指标已写入 %s", args.out)
    return 0


if __name__ == "__main__":
    sys.exit(main())