use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use ndarray::Array2;
use std::io::Cursor;
use taco_format::{CompressionSettings, Frame, FrameData};

fn create_test_frame(num_atoms: usize, frame_number: u64) -> Frame {
    // Generate positions
    let mut positions = Array2::<f32>::zeros((num_atoms, 3));
    for i in 0..num_atoms {
        for j in 0..3 {
            positions[[i, j]] =
                (i as f32) * 0.1 + (j as f32) * 0.01 + (frame_number as f32) * 0.001;
        }
    }

    // Generate velocities
    let mut velocities = Array2::<f32>::zeros((num_atoms, 3));
    for i in 0..num_atoms {
        for j in 0..3 {
            velocities[[i, j]] =
                (i as f32) * 0.05 + (j as f32) * 0.02 - (frame_number as f32) * 0.002;
        }
    }

    // Create frame with positions and velocities
    let frame_data = FrameData::new(positions).with_velocities(velocities);

    Frame::new(frame_number, frame_number as f64 * 0.001, frame_data)
        .with_box_dimensions([10.0, 10.0, 10.0], [90.0, 90.0, 90.0])
        .with_thermo(300.0, 1.0)
}

fn benchmark_frame_serialization(c: &mut Criterion) {
    let mut group = c.benchmark_group("frame_serialization");
    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

    let compression_settings = CompressionSettings::default();

    // Test different numbers of atoms
    for num_atoms in [100, 1_000, 10_000] {
        let frame = create_test_frame(num_atoms, 0);

        group.bench_with_input(
            BenchmarkId::new("write_frame", num_atoms),
            &frame,
            |b, frame| {
                b.iter(|| {
                    let mut buffer = Cursor::new(Vec::new());
                    frame.write(&mut buffer, &compression_settings).unwrap();
                    buffer.into_inner()
                });
            },
        );
    }

    group.finish();
}

fn benchmark_frame_deserialization(c: &mut Criterion) {
    let mut group = c.benchmark_group("frame_deserialization");
    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

    let compression_settings = CompressionSettings::default();

    // Test different numbers of atoms
    for num_atoms in [100, 1_000, 10_000] {
        let frame = create_test_frame(num_atoms, 0);
        let mut buffer = Cursor::new(Vec::new());
        frame.write(&mut buffer, &compression_settings).unwrap();
        let data = buffer.into_inner();

        group.bench_with_input(
            BenchmarkId::new("read_frame", num_atoms),
            &data,
            |b, data| {
                b.iter(|| {
                    let mut cursor = Cursor::new(data);
                    Frame::read(&mut cursor, &compression_settings).unwrap()
                });
            },
        );
    }

    group.finish();
}

fn benchmark_delta_encoding(c: &mut Criterion) {
    let mut group = c.benchmark_group("delta_encoding");
    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 numbers of atoms
    for num_atoms in [100, 1_000, 10_000] {
        // Create two consecutive frames
        let frame1 = create_test_frame(num_atoms, 0);
        let frame2 = create_test_frame(num_atoms, 1);

        group.bench_with_input(
            BenchmarkId::new("create_delta", num_atoms),
            &(frame1.data.clone(), frame2.data.clone()),
            |b, (data1, data2)| {
                b.iter(|| data2.delta_from(data1).unwrap());
            },
        );

        // Create a delta
        let delta = frame2.data.delta_from(&frame1.data).unwrap();

        group.bench_with_input(
            BenchmarkId::new("apply_delta", num_atoms),
            &(frame1.data.clone(), delta.clone()),
            |b, (data, delta)| {
                b.iter(|| data.apply_delta(delta).unwrap());
            },
        );
    }

    group.finish();
}

criterion_group!(
    frame_benches,
    benchmark_frame_serialization,
    benchmark_frame_deserialization,
    benchmark_delta_encoding
);
criterion_main!(frame_benches);
