//! Criterion benchmark for tensor operations

use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use rtx_bench::{BenchmarkConfig, core::tensor_ops::TensorOperationBenchmarks};
use tokio::runtime::Runtime;

fn benchmark_tensor_creation(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let benchmarks = TensorOperationBenchmarks::new();
    let config = BenchmarkConfig::default()
        .with_measurement_iterations(10)
        .with_warmup_iterations(3);

    let mut group = c.benchmark_group("tensor_creation");

    for size in [64, 256, 1024].iter() {
        group.bench_with_input(BenchmarkId::new("zeros", size), size, |b, &size| {
            b.to_async(&rt).iter(|| async {
                benchmarks
                    .benchmark_tensor_zeros(&config, size, size)
                    .await
                    .unwrap();
            });
        });

        group.bench_with_input(BenchmarkId::new("random", size), size, |b, &size| {
            b.to_async(&rt).iter(|| async {
                benchmarks
                    .benchmark_tensor_random(&config, size, size)
                    .await
                    .unwrap();
            });
        });
    }

    group.finish();
}

fn benchmark_tensor_arithmetic(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let benchmarks = TensorOperationBenchmarks::new();
    let config = BenchmarkConfig::default()
        .with_measurement_iterations(10)
        .with_warmup_iterations(3);

    let mut group = c.benchmark_group("tensor_arithmetic");

    for size in [256, 512, 1024].iter() {
        group.bench_with_input(BenchmarkId::new("add", size), size, |b, &size| {
            b.to_async(&rt).iter(|| async {
                benchmarks
                    .benchmark_tensor_add(&config, size, size)
                    .await
                    .unwrap();
            });
        });

        group.bench_with_input(BenchmarkId::new("matmul", size), size, |b, &size| {
            b.to_async(&rt).iter(|| async {
                benchmarks
                    .benchmark_tensor_matmul(&config, size, size)
                    .await
                    .unwrap();
            });
        });
    }

    group.finish();
}

criterion_group!(
    benches,
    benchmark_tensor_creation,
    benchmark_tensor_arithmetic
);
criterion_main!(benches);