//! Vision model benchmarks

use crate::{BenchmarkConfig, BenchmarkMeasurement, time_benchmark_async};
use anyhow::Result;
use std::collections::HashMap;

pub struct VisionBenchmarks;

impl VisionBenchmarks {
    pub fn new() -> Self {
        Self
    }

    pub async fn run_all_benchmarks(
        &self,
        config: &BenchmarkConfig,
    ) -> Result<HashMap<String, Vec<BenchmarkMeasurement>>> {
        let mut measurements = HashMap::new();

        measurements.insert(
            "resnet50_inference".to_string(),
            self.benchmark_resnet_inference(config).await?,
        );

        measurements.insert(
            "vit_inference".to_string(),
            self.benchmark_vit_inference(config).await?,
        );

        Ok(measurements)
    }

    async fn benchmark_resnet_inference(
        &self,
        config: &BenchmarkConfig,
    ) -> Result<Vec<BenchmarkMeasurement>> {
        let mut measurements = Vec::with_capacity(config.measurement_iterations);

        for _ in 0..config.measurement_iterations {
            let measurement = time_benchmark_async("resnet50_inference".to_string(), || async {
                self.resnet_forward().await?;
                Ok(())
            })
            .await?;
            measurements.push(measurement);
        }

        Ok(measurements)
    }

    async fn benchmark_vit_inference(
        &self,
        config: &BenchmarkConfig,
    ) -> Result<Vec<BenchmarkMeasurement>> {
        let mut measurements = Vec::with_capacity(config.measurement_iterations);

        for _ in 0..config.measurement_iterations {
            let measurement = time_benchmark_async("vit_inference".to_string(), || async {
                self.vit_forward().await?;
                Ok(())
            })
            .await?;
            measurements.push(measurement);
        }

        Ok(measurements)
    }

    async fn resnet_forward(&self) -> Result<()> {
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        Ok(())
    }

    async fn vit_forward(&self) -> Result<()> {
        tokio::time::sleep(std::time::Duration::from_millis(80)).await;
        Ok(())
    }
}

impl Default for VisionBenchmarks {
    fn default() -> Self {
        Self::new()
    }
}