use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
use ndarray::Array2;
use pprof::criterion::{Output, PProfProfiler};
use std::path::PathBuf;
use taco_format::{
    AtomMetadata, CompressionSettings, Frame, FrameData, Reader, SimulationMetadata, Writer,
};
use tempfile::{TempDir, tempdir};

/// Create a temporary file path for benchmarking
fn create_temp_file() -> (TempDir, PathBuf) {
    let dir = tempdir().expect("Failed to create temp directory");
    let file_path = dir.path().join("benchmark.taco");
    (dir, file_path)
}

/// Create a test trajectory file with the given number of atoms and frames
fn create_test_trajectory(num_atoms: usize, num_frames: usize) -> (TempDir, PathBuf) {
    let (dir, file_path) = create_temp_file();

    let simulation_metadata = SimulationMetadata::default();
    let atom_metadata = AtomMetadata::default();
    let compression_settings = CompressionSettings::default();

    // Create writer
    let mut writer = Writer::create(
        &file_path,
        num_atoms as u32,
        0.001, // time_step
        simulation_metadata,
        atom_metadata,
        compression_settings,
    )
    .expect("Failed to create writer");

    // Create and write frames
    for i in 0..num_frames {
        // Generate positions
        let mut positions = Array2::<f32>::zeros((num_atoms, 3));
        for a in 0..num_atoms {
            for j in 0..3 {
                positions[[a, j]] = (a as f32) * 0.1 + (j as f32) * 0.01 + (i as f32) * 0.001;
            }
        }

        // Create frame with positions
        let frame_data = FrameData::new(positions);
        let frame = Frame::new(i as u64, i as f64 * 0.001, frame_data);

        // Write frame
        writer.write_frame(frame).expect("Failed to write frame");
    }

    // Finish writing
    writer.finish().expect("Failed to finish writing");

    (dir, file_path)
}

fn benchmark_trajectory_writing(c: &mut Criterion) {
    let mut group = c.benchmark_group("trajectory_writing");
    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 trajectories
    for num_atoms in [100, 1_000] {
        for num_frames in [10, 50] {
            let simulation_metadata = SimulationMetadata::default();
            let atom_metadata = AtomMetadata::default();
            let compression_settings = CompressionSettings::default();

            // Prepare frames
            let frames: Vec<Frame> = (0..num_frames)
                .map(|i| {
                    let mut positions = Array2::<f32>::zeros((num_atoms, 3));
                    for a in 0..num_atoms {
                        for j in 0..3 {
                            positions[[a, j]] =
                                (a as f32) * 0.1 + (j as f32) * 0.01 + (i as f32) * 0.001;
                        }
                    }

                    let frame_data = FrameData::new(positions);
                    Frame::new(i as u64, i as f64 * 0.001, frame_data)
                })
                .collect();

            group.bench_with_input(
                BenchmarkId::new(format!("{num_atoms}atoms"), num_frames),
                &(num_atoms, num_frames, frames.clone()),
                |b, (num_atoms, _, frames)| {
                    b.iter(|| {
                        let (_dir, file_path) = create_temp_file();
                        let mut writer = Writer::create(
                            &file_path,
                            *num_atoms as u32,
                            0.001,
                            simulation_metadata.clone(),
                            atom_metadata.clone(),
                            compression_settings.clone(),
                        )
                        .unwrap();

                        for frame in frames {
                            writer.write_frame(frame.clone()).unwrap();
                        }

                        writer.finish().unwrap();
                    });
                },
            );

            // Benchmark parallel writing
            group.bench_with_input(
                BenchmarkId::new(format!("{num_atoms}atoms_parallel"), num_frames),
                &(num_atoms, num_frames, frames.clone()),
                |b, (num_atoms, _, frames)| {
                    b.iter(|| {
                        let (_dir, file_path) = create_temp_file();
                        let mut writer = Writer::create(
                            &file_path,
                            *num_atoms as u32,
                            0.001,
                            simulation_metadata.clone(),
                            atom_metadata.clone(),
                            compression_settings.clone(),
                        )
                        .unwrap();

                        writer.write_frames(frames.clone()).unwrap();
                        writer.finish().unwrap();
                    });
                },
            );
        }
    }

    group.finish();
}

fn benchmark_trajectory_reading(c: &mut Criterion) {
    let mut group = c.benchmark_group("trajectory_reading");
    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 trajectories
    for num_atoms in [100, 1_000] {
        for num_frames in [10, 50] {
            let (_dir, file_path) = create_test_trajectory(num_atoms, num_frames);

            // Sequential reading
            group.bench_with_input(
                BenchmarkId::new(format!("{num_atoms}atoms_sequential"), num_frames),
                &file_path,
                |b, path| {
                    b.iter(|| {
                        let mut reader = Reader::open(path).unwrap();
                        for i in 0..num_frames {
                            let frame = reader.read_frame(i as u64).unwrap();
                            black_box(frame);
                        }
                    });
                },
            );

            // Parallel reading of all frames
            group.bench_with_input(
                BenchmarkId::new(format!("{num_atoms}atoms_parallel"), num_frames),
                &file_path,
                |b, path| {
                    b.iter(|| {
                        let mut reader = Reader::open(path).unwrap();
                        let frames = reader.read_frame_range(0, num_frames as u64).unwrap();
                        black_box(frames);
                    });
                },
            );

            // Random access pattern (every other frame)
            let indices: Vec<u64> = (0..num_frames as u64).step_by(2).collect();
            group.bench_with_input(
                BenchmarkId::new(format!("{num_atoms}atoms_random"), num_frames),
                &(file_path.clone(), indices.clone()),
                |b, (path, indices)| {
                    b.iter(|| {
                        let mut reader = Reader::open(path).unwrap();
                        let frames = reader.read_frames(indices).unwrap();
                        black_box(frames);
                    });
                },
            );
        }
    }

    group.finish();
}

criterion_group!(
    name = io_benches;
    config = Criterion::default().with_profiler(PProfProfiler::new(100, Output::Flamegraph(None)));
    targets = benchmark_trajectory_writing, benchmark_trajectory_reading
);
criterion_main!(io_benches);
