//! YOLOv8n with TensorRT backend + OpenCV for image I/O
//!
//! This example demonstrates the `tensorrt-opencv-compat` feature:
//! - OpenCV handles image loading (imread) and drawing
//! - TensorRT handles neural network inference
//!
//! Run with: cargo run --example yolo_v8_n_tensorrt_opencv --no-default-features --features "tensorrt-opencv-compat" --release

use std::time::Instant;

use opencv::{
    imgcodecs,
    imgproc,
    core::Scalar,
    prelude::*,
};
use od_opencv::{Model, ModelTrait};

fn main() {
    let classes_labels: Vec<&str> = vec![
        "person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck", "boat",
        "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
        "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
        "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
        "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
        "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
        "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
        "sofa", "pottedplant", "bed", "diningtable", "toilet", "tvmonitor", "laptop", "mouse",
        "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator",
        "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush",
    ];

    // Create TensorRT model
    let mut model = Model::tensorrt(
        "pretrained/yolov8n.engine",
    ).expect("Failed to load engine");

    // Load image using OpenCV
    let mut img = imgcodecs::imread("images/dog.jpg", imgcodecs::IMREAD_COLOR)
        .expect("Failed to load image");

    println!("Image size: {}x{}", img.cols(), img.rows());

    // Warm up
    for _ in 0..10 {
        let _ = ModelTrait::forward(&mut model, &img, 0.25, 0.4).expect("Warm-up failed");
    }

    // Benchmark
    let start = Instant::now();
    let n_runs = 20;
    let mut result = None;
    for _ in 0..n_runs {
        result = Some(ModelTrait::forward(&mut model, &img, 0.25, 0.4).expect("Inference failed"));
    }
    let avg = start.elapsed() / n_runs;
    println!("Average inference time ({} runs): {:?}", n_runs, avg);

    let (bboxes, class_ids, confidences) = result.unwrap();

    // Draw results on image
    for (i, bbox) in bboxes.iter().enumerate() {
        let label = format!("{}: {:.2}", classes_labels[class_ids[i]], confidences[i]);

        imgproc::rectangle(
            &mut img,
            *bbox,
            Scalar::new(0.0, 255.0, 0.0, 0.0),
            2,
            imgproc::LINE_8,
            0,
        ).unwrap();

        imgproc::put_text(
            &mut img,
            &label,
            opencv::core::Point::new(bbox.x, bbox.y - 5),
            imgproc::FONT_HERSHEY_SIMPLEX,
            0.5,
            Scalar::new(0.0, 255.0, 0.0, 0.0),
            1,
            imgproc::LINE_8,
            false,
        ).unwrap();

        println!("Class: {}", classes_labels[class_ids[i]]);
        println!("\tBounding box: x={}, y={}, w={}, h={}", bbox.x, bbox.y, bbox.width, bbox.height);
        println!("\tConfidence: {:.2}", confidences[i]);
    }
}
