#!/usr/bin/env python
"""训练入口。

用法:
    python train.py --config configs/train/baseline.yaml
    python train.py --config configs/train/baseline.yaml --set train.epochs=10 train.optimizer.lr=5e-4
    python train.py --config configs/train/baseline.yaml --device mps

设备由配置的 `device` 决定，默认 `auto`：按 cuda > mps > cpu 挑最快的可用后端。
显式指定但不可用时直接报错，不会悄悄退回 CPU。
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

from datasets import build_dataloader, build_dataset, load_classes
from engine import Trainer
from models import build_model
from utils import check_label_mode, load_config, seed_everything, 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, help="yaml 配置路径")
    parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE", help="覆盖配置项")
    parser.add_argument("--work-dir", default=None, help="覆盖配置中的 work_dir")
    parser.add_argument("--device", default=None, help="覆盖配置中的 device：auto / cuda / mps / cpu")
    parser.add_argument("--resume", default=None, help="从检查点恢复训练")
    return parser.parse_args()


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

    work_dir = Path(args.work_dir or cfg.get("work_dir", "runs/default"))
    work_dir.mkdir(parents=True, exist_ok=True)
    logger = setup_logging(work_dir / "train.log")

    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))

    seed_everything(int(cfg.get("seed", 42)), bool(cfg.get("deterministic", False)))

    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).to(device)
    logger.info(
        "模型 %s 输出 stride=%d，参数量 %.2fM",
        model_cfg["name"],
        model.out_stride,
        sum(p.numel() for p in model.parameters()) / 1e6,
    )

    loaders = {}
    for split in ("train", "val"):
        split_cfg = cfg["data"].get(split)
        if not split_cfg:
            loaders[split] = None
            continue
        dataset = build_dataset(
            split_cfg["dataset"], classes=classes, stride=model.out_stride, label_mode=label_mode
        )
        loaders[split] = build_dataloader(dataset, **dict(split_cfg.get("loader", {})))
        logger.info("%s 集样本数: %d", split, len(dataset))

    trainer = Trainer(
        model=model,
        train_loader=loaders["train"],
        val_loader=loaders["val"],
        cfg=cfg.get("train", {}),
        device=device,
        work_dir=work_dir,
        full_config=cfg.to_dict(),
    )

    if args.resume:
        trainer.resume(args.resume)

    cfg.dump(work_dir / "config.yaml")
    best = trainer.fit()
    logger.info("训练结束，最佳 %s=%.4f", trainer.monitor, best)
    return 0


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