use taco_format::{AsyncReader, AsyncWriter, CompressionSettings, Frame, FrameData, SimulationMetadata, AtomMetadata, Result};
use ndarray::Array2;
use std::collections::HashMap;
use tempfile::NamedTempFile;

#[tokio::main]
async fn main() -> Result<()> {
    // Recreate the failing test but with debug output
    let temp_file = NamedTempFile::new().unwrap();
    let path = temp_file.path().to_path_buf();

    let num_atoms = 4;
    let time_step = 0.001;
    let simulation_metadata = SimulationMetadata::default();
    let atom_metadata = AtomMetadata::default();
    let compression_settings = CompressionSettings::default();

    // Create frames
    let mut frames = Vec::new();
    for i in 0..6 {
        let positions_data = vec![(i * 10) as f32; num_atoms as usize * 3];
        let positions = Array2::from_shape_vec((num_atoms as usize, 3), positions_data).unwrap();

        let frame_data = FrameData {
            positions: Some(positions),
            velocities: None,
            forces: None,
            stress: None,
            extra: HashMap::new(),
        };

        frames.push(Frame::new(i as u64, i as f64 * time_step, frame_data));
    }

    println!("Writing frames with values:");
    for (i, frame) in frames.iter().enumerate() {
        if let Some(pos) = &frame.data.positions {
            println!("Frame {}: [{}, {}, {}]", i, pos[[0,0]], pos[[0,1]], pos[[0,2]]);
        }
    }

    // Write frames
    {
        let mut writer = AsyncWriter::create(
            &path,
            num_atoms,
            time_step,
            simulation_metadata,
            atom_metadata,
            compression_settings,
        ).await?;

        writer.write_frames(frames).await?;
        writer.finish().await?;
    }

    // Read sequentially
    let mut sequential_frames = Vec::new();
    {
        let mut reader = AsyncReader::open(&path).await?;
        reader.reset().await?;
        while let Some(frame) = reader.next_frame().await? {
            sequential_frames.push(frame);
        }
    }

    println!("\nSequential read:");
    for (i, frame) in sequential_frames.iter().enumerate() {
        if let Some(pos) = &frame.data.positions {
            println!("Frame {}: [{}, {}, {}]", i, pos[[0,0]], pos[[0,1]], pos[[0,2]]);
        }
    }

    // Read randomly
    let mut random_frames = Vec::new();
    {
        let mut reader = AsyncReader::open(&path).await?;
        let indices = vec![5, 1, 3, 0, 4, 2]; // Random order
        for &idx in &indices {
            let frame = reader.read_frame(idx).await?;
            println!("Read frame {} randomly: [{}, {}, {}]", 
                     idx, 
                     frame.data.positions.as_ref().unwrap()[[0,0]], 
                     frame.data.positions.as_ref().unwrap()[[0,1]], 
                     frame.data.positions.as_ref().unwrap()[[0,2]]);
            random_frames.push(frame);
        }
        // Sort back to original order
        random_frames.sort_by_key(|f| f.frame_number);
    }

    println!("\nRandom read (sorted):");
    for (i, frame) in random_frames.iter().enumerate() {
        if let Some(pos) = &frame.data.positions {
            println!("Frame {}: [{}, {}, {}]", frame.frame_number, pos[[0,0]], pos[[0,1]], pos[[0,2]]);
        }
    }

    // Compare frame 1
    println!("\nComparing frame 1:");
    if let (Some(seq_pos), Some(rand_pos)) = (&sequential_frames[1].data.positions, &random_frames[1].data.positions) {
        println!("Sequential: {:?}", seq_pos);
        println!("Random:     {:?}", rand_pos);
        println!("Equal: {}", seq_pos == rand_pos);
    }

    Ok(())
}
