"""
Export trained RT-DETR model for OAK-D camera deployment.

This script exports your trained model to OpenVINO format, which can then be
converted to a blob for the OAK-D camera.

Usage:
    python export_rtdetr_oak.py --weights runs/rtdetr_training/training/weights/best.pt
"""

import argparse
from pathlib import Path


def export_for_oak(weights_path: Path, img_size: int = 640):
    """
    Export RT-DETR model for OAK-D deployment.
    
    Args:
        weights_path: Path to trained .pt weights
        img_size: Input image size (should match training)
    """
    from ultralytics import RTDETR
    
    if not weights_path.exists():
        raise ValueError(f"Weights not found: {weights_path}")
    
    print(f"\n{'='*60}")
    print(f"Exporting RT-DETR for OAK-D camera")
    print(f"Weights: {weights_path}")
    print(f"Image size: {img_size}x{img_size}")
    print(f"{'='*60}\n")
    
    # Load model
    model = RTDETR(weights_path)
    
    # Export to ONNX first (intermediate format)
    print("Step 1/2: Exporting to ONNX...")
    onnx_path = model.export(
        format="onnx",
        imgsz=img_size,
        simplify=True,
        opset=11,  # OAK-compatible opset
    )
    print(f"✓ ONNX exported: {onnx_path}")
    
    # Export to OpenVINO (for OAK)
    print("\nStep 2/2: Exporting to OpenVINO...")
    openvino_path = model.export(
        format="openvino",
        imgsz=img_size,
        half=False,  # Use FP32 for better compatibility
    )
    print(f"✓ OpenVINO exported: {openvino_path}")
    
    print(f"\n{'='*60}")
    print(f"✓ Export complete!")
    print(f"\nNext steps:")
    print(f"1. Test OpenVINO model: python test_openvino.py --model {openvino_path}")
    print(f"2. Convert to blob for OAK:")
    print(f"   Online: https://blobconverter.luxonis.com/")
    print(f"   Or use: blobconverter --openvino-xml {openvino_path}/model.xml")
    print(f"3. Deploy blob to OAK-D camera with DepthAI")
    print(f"{'='*60}\n")
    
    return openvino_path


def main():
    parser = argparse.ArgumentParser(description="Export RT-DETR for OAK-D deployment")
    parser.add_argument(
        "--weights",
        type=Path,
        required=True,
        help="Path to trained .pt weights file"
    )
    parser.add_argument(
        "--img-size",
        type=int,
        default=640,
        help="Input image size (should match training)"
    )
    
    args = parser.parse_args()
    
    export_for_oak(args.weights, args.img_size)


if __name__ == "__main__":
    main()