//! Criterion benchmarks for model architectures

use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use rtx_bench::{BenchmarkConfig, models::transformers::TransformerBenchmarks};
use tokio::runtime::Runtime;

fn benchmark_transformer_models(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let benchmarks = TransformerBenchmarks::new();
    let config = BenchmarkConfig::default()
        .with_measurement_iterations(5)
        .with_warmup_iterations(2);

    let mut group = c.benchmark_group("transformer_models");
    group.sample_size(10);

    // Test with smaller configurations for fast benchmarking
    for batch_size in [1, 4, 8].iter() {
        group.bench_with_input(
            BenchmarkId::new("bert_forward", batch_size),
            batch_size,
            |b, &batch_size| {
                b.to_async(&rt).iter(|| async {
                    let model_config = rtx_bench::models::transformers::ModelConfig {
                        model_type: rtx_bench::models::transformers::ModelType::Bert,
                        hidden_size: 256,
                        num_layers: 4,
                        num_heads: 8,
                        sequence_length: 128,
                        vocab_size: 10000,
                    };
                    benchmarks
                        .benchmark_forward_pass(&config, &model_config, batch_size)
                        .await
                        .unwrap();
                });
            },
        );
    }

    group.finish();
}

criterion_group!(benches, benchmark_transformer_models);
criterion_main!(benches);