# /// script
# requires-python = ">=3.10"
# dependencies = ["ultralytics>=8.0"]
# ///
"""Export a YOLOv8-worldv2 model to ONNX with N open-vocabulary class slots.

The exported model has two inputs:
  images     (1, 3, imgsz, imgsz)  float32
  text_feats (1, N, 512)           float32  ← CLIP embeddings, set at runtime

N is fixed per export. Use export_clip_text_encoder.py to generate the text
encoder and vocab files used to compute embeddings at runtime in the C++ node.

Usage:
    uv run scripts/export_yoloworld.py
    uv run scripts/export_yoloworld.py --model m --num-classes 20
    uv run scripts/export_yoloworld.py --model l --num-classes 5 --output-dir models/
"""

import argparse
import shutil
import sys
from pathlib import Path

MODELS = {
    "s": "yolov8s-worldv2",
    "m": "yolov8m-worldv2",
    "l": "yolov8l-worldv2",
    "x": "yolov8x-worldv2",
}


def parse_args():
    p = argparse.ArgumentParser(
        description="Export YOLOv8-worldv2 to ONNX with open-vocabulary text inputs"
    )
    p.add_argument(
        "--model", "-m",
        choices=list(MODELS),
        default="s",
        help="Model size: s (~14 MB), m (~30 MB), l (~50 MB), x (~100 MB). Default: s",
    )
    p.add_argument(
        "--num-classes", "-n",
        type=int,
        default=10,
        help="Number of concurrent class slots (must match number of labels at runtime). Default: 10",
    )
    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_inputs(onnx_path: Path, expected_inputs: int):
    try:
        import onnx
        model = onnx.load(str(onnx_path))
        actual = len(model.graph.input)
        if actual != expected_inputs:
            print(f"\nWarning: expected {expected_inputs} inputs, got {actual}.")
            print("The visual model may have baked text embeddings (ultralytics version dependent).")
            print("In that case, use a fixed vocabulary instead of the text encoder.")
        else:
            inputs = [(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim])
                      for i in model.graph.input]
            print(f"  Inputs: {inputs}")
    except ImportError:
        print("  (install 'onnx' to verify model inputs)")


def main():
    args = parse_args()

    try:
        from ultralytics import YOLO
    except ImportError:
        print("Error: ultralytics not installed. Run: uv run scripts/export_yoloworld.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}_n{args.num_classes}_i{args.imgsz}.onnx"
    dst = out_dir / out_name

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

    model = YOLO(f"{model_name}.pt")
    # Set N dummy class slots — actual label meanings are provided at runtime via text_feats
    model.set_classes(["object"] * args.num_classes)
    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_inputs(dst, expected_inputs=2)

    print(f"\nExported: {dst.resolve()}")
    print(f"\nNext step — export the CLIP text encoder (once, reused for all model sizes):")
    print(f"  uv run scripts/export_clip_text_encoder.py --output-dir models/")
    print(f"\nRun with:")
    print(f"  ZED_ENGINE=open_vocab \\")
    print(f"  ZED_ONNX_MODEL_PATH={dst.resolve()} \\")
    print(f"  ZED_ONNX_TEXT_ENCODER_PATH=models/clip_text_encoder.onnx \\")
    print(f'  ZED_ONNX_LABELS="person,car,bag,...  (exactly {args.num_classes} labels)" \\')
    print(f"  dora start dataflow.yml")


if __name__ == "__main__":
    main()