#!/usr/bin/env python
"""导出 ONNX，供实时设备（TensorRT / OpenVINO）部署。

    python scripts/export_onnx.py --config configs/train/baseline.yaml \
        --checkpoint runs/baseline/best.pth --output runs/baseline/model.onnx

导出在 CPU 上完成，不需要 GPU；需要额外依赖：
    uv sync --group export
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import torch  # noqa: E402

from datasets import load_classes  # noqa: E402
from engine import load_checkpoint  # noqa: E402
from models import build_model  # noqa: E402
from utils import load_config  # noqa: E402


def main() -> int:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--config", required=True)
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--opset", type=int, default=17)
    parser.add_argument("--dynamic-batch", action="store_true")
    args = parser.parse_args()

    cfg = load_config(args.config)
    model_cfg = dict(cfg["model"])
    model_cfg.setdefault("num_classes", len(load_classes(cfg["classes"])))
    model_cfg.setdefault("label_mode", str(cfg.get("label_mode", "multilabel")))
    model = build_model(model_cfg)
    load_checkpoint(args.checkpoint, model, map_location="cpu")
    model.eval()

    height, width = cfg.get("image_size", [512, 512])
    channels = 1 if cfg.get("gray") else 3
    dummy = torch.randn(1, channels, int(height), int(width))

    output_names = ["heatmap", "wh", "offset"] + (["logits"] if model.cls_head is not None else [])
    dynamic_axes = {"images": {0: "batch"}} if args.dynamic_batch else None
    if dynamic_axes:
        dynamic_axes.update({name: {0: "batch"} for name in output_names})

    try:
        import onnxscript  # noqa: F401
    except ImportError:
        print("缺少导出依赖，请先执行: uv sync --group export", file=sys.stderr)
        return 1

    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    torch.onnx.export(
        model,
        dummy,
        str(out_path),
        input_names=["images"],
        output_names=output_names,
        opset_version=args.opset,
        dynamic_axes=dynamic_axes,
    )
    print(f"已导出 {out_path}（输入 {tuple(dummy.shape)}，输出 stride={model.out_stride}）")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())