use ndarray::Array2;
use std::path::PathBuf;
use taco_format::{
    AtomMetadata, CompressionSettings, Frame, FrameData, Reader, SimulationMetadata, Writer,
};
use tempfile::{TempDir, tempdir};

/// Create a temporary file path for examples
fn create_temp_file() -> (TempDir, PathBuf) {
    let dir = tempdir().expect("Failed to create temp directory");
    let file_path = dir.path().join("example.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)
}

/// Example function for trajectory writing
fn example_trajectory_writing(num_atoms: usize, num_frames: usize) {
    println!("Writing trajectory with {num_atoms} atoms and {num_frames} frames");

    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();

    // Serial writing
    println!("Serial writing...");
    for _ in 0..100 {
        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();
    }

    println!("Serial writing completed.");
}

/// Example function for trajectory reading
fn example_trajectory_reading(num_atoms: usize, num_frames: usize) {
    println!("Reading trajectory with {num_atoms} atoms and {num_frames} frames");

    let (_dir, file_path) = create_test_trajectory(num_atoms, num_frames);

    // Sequential reading
    println!("Sequential reading...");
    for _ in 0..10000 {
        let mut reader = Reader::open(&file_path).unwrap();
        for i in 0..num_frames {
            let frame = reader.read_frame(i as u64).unwrap();
            // Process frame (just accessing data here to ensure it's not optimized away)
            std::hint::black_box(frame.data.positions.unwrap()[[0, 0]]);
            // println!("Frame {} read", i);
        }
    }

    // // Parallel reading of all frames
    // println!("Parallel reading...");
    // {
    //     let mut reader = Reader::open(&file_path).unwrap();
    //     let frames = reader.read_frame_range(0, num_frames as u64).unwrap();
    //     // Process frames (just accessing data here to ensure it's not optimized away)
    //     for frame in &frames {
    //         let _ = frame.data.positions.as_ref().unwrap()[[0, 0]];
    //     }
    // }

    // // Random access pattern (every other frame)
    // println!("Random access reading...");
    // {
    //     let indices: Vec<u64> = (0..num_frames as u64).step_by(2).collect();
    //     let mut reader = Reader::open(&file_path).unwrap();
    //     let frames = reader.read_frames(&indices).unwrap();
    //     // Process frames (just accessing data here to ensure it's not optimized away)
    //     for frame in &frames {
    //         let _ = frame.data.positions.as_ref().unwrap()[[0, 0]];
    //     }
    // }

    println!("Reading completed.");
}

fn main() {
    // Default sizes, can be modified via command line arguments
    let num_atoms = 1000;
    let num_frames = 1000;

    // Parse command line arguments if provided
    let args: Vec<String> = std::env::args().collect();
    let write_mode = args.contains(&"--write".to_string());

    let (atoms, frames) = if args.len() > 2 {
        let a = args[1].parse().unwrap_or(num_atoms);
        let f = args[2].parse().unwrap_or(num_frames);
        (a, f)
    } else {
        (num_atoms, num_frames)
    };

    if write_mode {
        example_trajectory_writing(atoms, frames);
    } else {
        example_trajectory_reading(atoms, frames);
    }
}
