# /// script
# requires-python = ">=3.10"
# dependencies = ["ultralytics>=8.3"]
# ///
"""Export a standard YOLOv11 model to ONNX for closed-vocabulary detection.

The exported model has one input:
  images  (1, 3, imgsz, imgsz)  float32

Class names are fixed at training time. Provide them at runtime via
ZED_ONNX_LABELS or ZED_ONNX_LABELS_FILE.

Usage:
    uv run scripts/export_yolo.py
    uv run scripts/export_yolo.py --model s --output-dir models/
    uv run scripts/export_yolo.py --model m --imgsz 1280
"""

import argparse
import shutil
import sys
from pathlib import Path

MODELS = {
    "n": "yolo11n",
    "s": "yolo11s",
    "m": "yolo11m",
    "l": "yolo11l",
    "x": "yolo11x",
}


def parse_args():
    p = argparse.ArgumentParser(
        description="Export YOLOv11 to ONNX for closed-vocabulary detection"
    )
    p.add_argument(
        "--model", "-m",
        choices=list(MODELS),
        default="n",
        help="Model size: n (~6 MB), s (~18 MB), m (~38 MB), l (~49 MB), x (~109 MB). Default: n",
    )
    p.add_argument(
        "--imgsz",
        type=int,
        default=640,
        help="Input image size in pixels (square). Default: 640",
    )
    p.add_argument(
        "--output-dir", "-o",
        default=".",
        help="Output directory for the .onnx file. Default: current directory",
    )
    return p.parse_args()


def verify_output(onnx_path: Path):
    try:
        import onnx
        model = onnx.load(str(onnx_path))
        inputs  = [(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim])
                   for i in model.graph.input]
        outputs = [(o.name, [d.dim_value for d in o.type.tensor_type.shape.dim])
                   for o in model.graph.output]
        print(f"  Inputs:  {inputs}")
        print(f"  Outputs: {outputs}")
    except ImportError:
        print("  (install 'onnx' to verify model structure)")


def main():
    args = parse_args()

    try:
        from ultralytics import YOLO
    except ImportError:
        print("Error: ultralytics not installed. Run: uv run scripts/export_yolo.py", file=sys.stderr)
        sys.exit(1)

    model_name = MODELS[args.model]
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    out_name = f"{model_name}_i{args.imgsz}.onnx"
    dst = out_dir / out_name

    print(f"Model:      {model_name}.pt")
    print(f"Input size: {args.imgsz}")
    print(f"Output:     {dst}")
    print()

    model = YOLO(f"{model_name}.pt")
    exported = Path(model.export(format="onnx", imgsz=args.imgsz))

    if exported.resolve() != dst.resolve():
        shutil.move(str(exported), str(dst))

    print(f"\nVerifying exported model...")
    verify_output(dst)

    print(f"\nExported: {dst.resolve()}")
    print(f"\nRun with (COCO labels example):")
    print(f"  ZED_ENGINE=onnx \\")
    print(f"  ZED_ONNX_MODEL_PATH={dst.resolve()} \\")
    print(f"  ZED_ONNX_LABELS_FILE=models/coco80.txt \\")
    print(f"  dora start dataflow.yml")
    print(f"\nOr with inline labels:")
    print(f"  ZED_ONNX_LABELS=\"person,car,chair\" \\")


if __name__ == "__main__":
    main()