use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
use ndarray::Array2;
use taco_format::{PrecisionMode, compress_tensor, decompress_tensor};

fn create_test_data(size: usize) -> Array2<f32> {
    let mut data = Array2::<f32>::zeros((size, 3));
    for i in 0..size {
        for j in 0..3 {
            data[[i, j]] = (i as f32) + (j as f32) * 0.1 + 0.05 * (i % 7) as f32;
        }
    }
    data
}

fn benchmark_compression(c: &mut Criterion) {
    let mut group = c.benchmark_group("tensor_compression");
    group
        .sample_size(25) // Reduce the number of iterations
        .warm_up_time(std::time::Duration::from_millis(300)) // 0.3s instead of 3s
        .measurement_time(std::time::Duration::from_millis(500)); // 0.5s instead of 5s

    // Test different sizes of tensors
    for size in [100, 1_000, 10_000] {
        let data = create_test_data(size);

        // Test with different precision modes
        for &precision in &[PrecisionMode::Full, PrecisionMode::Half] {
            let precision_name = match precision {
                PrecisionMode::Full => "full",
                PrecisionMode::Half => "half",
            };

            let id = BenchmarkId::new(format!("zstd_{precision_name}"), size);

            group.bench_with_input(id, &data, |b, data| {
                b.iter(|| {
                    compress_tensor(
                        black_box(data),
                        black_box(precision),
                        black_box(3), // Use default compression level
                    )
                });
            });
        }
    }

    group.finish();
}

fn benchmark_decompression(c: &mut Criterion) {
    let mut group = c.benchmark_group("tensor_decompression");
    group
        .sample_size(25) // Reduce the number of iterations
        .warm_up_time(std::time::Duration::from_millis(300)) // 0.3s instead of 3s
        .measurement_time(std::time::Duration::from_millis(500)); // 0.5s instead of 5s

    // Test different sizes of tensors
    for size in [100, 1_000, 10_000] {
        let data = create_test_data(size);

        // Test with different precision modes
        for &precision in &[PrecisionMode::Full, PrecisionMode::Half] {
            let precision_name = match precision {
                PrecisionMode::Full => "full",
                PrecisionMode::Half => "half",
            };

            let compressed = compress_tensor(&data, precision, 3).unwrap();

            let id = BenchmarkId::new(format!("zstd_{precision_name}"), size);

            group.bench_with_input(id, &compressed, |b, compressed_data| {
                b.iter(|| decompress_tensor(black_box(compressed_data), black_box(precision)));
            });
        }
    }

    group.finish();
}

criterion_group!(
    compression_benches,
    benchmark_compression,
    benchmark_decompression
);
criterion_main!(compression_benches);
