use criterion::{Criterion, black_box, criterion_group, criterion_main};
use rtx_multimodal::*;
use rtx_runtime::Device;
use rtx_tensor::Tensor;

fn benchmark_vit_forward(c: &mut Criterion) {
    let device = Device::cuda(0).unwrap_or(Device::default());
    let config = vision::vit::ViTConfig {
        image_size: 224,
        patch_size: 16,
        in_channels: 3,
        embed_dim: 768,
        depth: 12,
        num_heads: 12,
        mlp_ratio: 4.0,
        num_classes: 1000,
        dropout: 0.0,
        attention_dropout: 0.0,
    };

    let model = vision::vit::VisionTransformer::new(&config, &device).unwrap();
    let input = Tensor::randn(&[1, 3, 224, 224], &device).unwrap();

    c.bench_function("vit_forward", |b| {
        b.iter(|| {
            let _output = model.forward(black_box(&input)).unwrap();
        })
    });
}

fn benchmark_conformer_forward(c: &mut Criterion) {
    let device = Device::cuda(0).unwrap_or(Device::default());
    let config = audio::conformer::ConformerEncoderConfig {
        input_dim: 80,
        d_model: 512,
        num_layers: 6,
        num_heads: 8,
        feed_forward_expansion_factor: 4,
        conv_kernel_size: 31,
        dropout: 0.0,
    };

    let model = audio::conformer::ConformerEncoder::new(&config, &device).unwrap();
    let input = Tensor::randn(&[1, 100, 80], &device).unwrap();

    c.bench_function("conformer_forward", |b| {
        b.iter(|| {
            let _output = model.forward(black_box(&input)).unwrap();
        })
    });
}

fn benchmark_timesformer_forward(c: &mut Criterion) {
    let device = Device::cuda(0).unwrap_or(Device::default());
    let config = video::timesformer::TimeSformerConfig {
        image_size: 224,
        patch_size: 16,
        in_channels: 3,
        embed_dim: 768,
        depth: 12,
        num_heads: 12,
        mlp_ratio: 4.0,
        num_frames: 8,
        num_classes: 400,
        dropout: 0.0,
        attention_dropout: 0.0,
    };

    let model = video::timesformer::TimeSformerModel::new(&config, &device).unwrap();
    let input = Tensor::randn(&[1, 3, 8, 224, 224], &device).unwrap();

    c.bench_function("timesformer_forward", |b| {
        b.iter(|| {
            let _output = model.forward(black_box(&input)).unwrap();
        })
    });
}

fn benchmark_cross_modal_attention(c: &mut Criterion) {
    let device = Device::cuda(0).unwrap_or(Device::default());
    let config = fusion::CrossModalAttentionConfig {
        embed_dim: 512,
        num_heads: 8,
        dropout: 0.0,
    };

    let attention = fusion::CrossModalAttention::new(&config, &device).unwrap();
    let vision = Tensor::randn(&[1, 197, 512], &device).unwrap();
    let audio = Tensor::randn(&[1, 100, 512], &device).unwrap();

    c.bench_function("cross_modal_attention", |b| {
        b.iter(|| {
            let _output = attention
                .forward(black_box(&vision), black_box(&audio))
                .unwrap();
        })
    });
}

criterion_group!(
    benches,
    benchmark_vit_forward,
    benchmark_conformer_forward,
    benchmark_timesformer_forward,
    benchmark_cross_modal_attention
);
criterion_main!(benches);