//! Tests for random access functionality with normal and extra data

use ndarray::{Array1, Array2, Array3};
use taco_format::{AtomMetadata, CompressionSettings, ExtraArray, SimulationMetadata};
use taco_format::{Frame, FrameData, Reader, Writer};
use tempfile::tempdir;

#[test]
fn test_random_access_with_extra_data() {
    // Create a temporary directory for the test
    let dir = tempdir().expect("Failed to create temporary directory");
    let file_path = dir.path().join("test_random_access.taco");

    // Create metadata
    let sim_metadata = SimulationMetadata {
        name: Some("Random Access Test".to_string()),
        temperature: Some(300.0),
        ..Default::default()
    };

    let atom_metadata = AtomMetadata {
        names: Some(vec!["C".to_string(), "O".to_string()]),
        elements: Some(vec!["C".to_string(), "O".to_string()]),
        masses: Some(vec![12.01, 16.0]),
        ..Default::default()
    };

    let compression_settings = CompressionSettings::default();

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

    // Create and write 20 frames with various data
    let mut original_frames: Vec<FrameData> = Vec::new();

    for frame_idx in 0..20 {
        let time = frame_idx as f64 * 0.001;

        // Create positions that change over time
        let positions = Array2::<f32>::from_shape_vec(
            (2, 3),
            vec![
                frame_idx as f32 * 0.1,
                (frame_idx as f32 * 0.1) + 1.0,
                (frame_idx as f32 * 0.1) + 2.0,
                (frame_idx as f32 * 0.1) + 3.0,
                (frame_idx as f32 * 0.1) + 4.0,
                (frame_idx as f32 * 0.1) + 5.0,
            ],
        )
        .unwrap();

        // Create velocities
        let velocities = Array2::<f32>::from_shape_vec(
            (2, 3),
            vec![
                frame_idx as f32 * 0.01,
                (frame_idx as f32 * 0.01) + 0.1,
                (frame_idx as f32 * 0.01) + 0.2,
                (frame_idx as f32 * 0.01) + 0.3,
                (frame_idx as f32 * 0.01) + 0.4,
                (frame_idx as f32 * 0.01) + 0.5,
            ],
        )
        .unwrap();

        // Create forces
        let forces = Array2::<f32>::from_shape_vec(
            (2, 3),
            vec![
                frame_idx as f32 * 0.005,
                (frame_idx as f32 * 0.005) + 0.05,
                (frame_idx as f32 * 0.005) + 0.1,
                (frame_idx as f32 * 0.005) + 0.15,
                (frame_idx as f32 * 0.005) + 0.2,
                (frame_idx as f32 * 0.005) + 0.25,
            ],
        )
        .unwrap();

        // Create stress tensor (per atom)
        let stress = Array2::<f32>::from_shape_vec(
            (2, 6),
            vec![
                frame_idx as f32 * 0.001,
                (frame_idx as f32 * 0.001) + 0.01,
                (frame_idx as f32 * 0.001) + 0.02,
                (frame_idx as f32 * 0.001) + 0.03,
                (frame_idx as f32 * 0.001) + 0.04,
                (frame_idx as f32 * 0.001) + 0.05,
                (frame_idx as f32 * 0.001) + 0.06,
                (frame_idx as f32 * 0.001) + 0.07,
                (frame_idx as f32 * 0.001) + 0.08,
                (frame_idx as f32 * 0.001) + 0.09,
                (frame_idx as f32 * 0.001) + 0.10,
                (frame_idx as f32 * 0.001) + 0.11,
            ],
        )
        .unwrap();

        // Create extra data arrays
        let extra1d = Array1::<f32>::from(vec![
            frame_idx as f32 * 10.0,
            (frame_idx as f32 * 10.0) + 5.0,
        ]);

        let extra2d = Array2::<f32>::from_shape_vec(
            (2, 2),
            vec![
                frame_idx as f32 * 100.0,
                (frame_idx as f32 * 100.0) + 10.0,
                (frame_idx as f32 * 100.0) + 20.0,
                (frame_idx as f32 * 100.0) + 30.0,
            ],
        )
        .unwrap();

        let extra3d = Array3::<f32>::from_shape_vec(
            (2, 1, 2),
            vec![
                frame_idx as f32 * 1000.0,
                (frame_idx as f32 * 1000.0) + 100.0,
                (frame_idx as f32 * 1000.0) + 200.0,
                (frame_idx as f32 * 1000.0) + 300.0,
            ],
        )
        .unwrap();

        // Create frame data with all types of data
        let mut frame_data = FrameData::new(positions)
            .with_velocities(velocities)
            .with_forces(forces)
            .with_stress(stress);

        frame_data.set_extra_1d("density", extra1d);
        frame_data.set_extra_2d("tensor", extra2d);
        frame_data.set_extra_3d("volume", extra3d);

        // Store original frame data for comparison
        original_frames.push(frame_data.clone());

        // Create and write frame
        let frame = Frame::new(frame_idx as u64, time, frame_data)
            .with_box_dimensions([10.0, 10.0, 10.0], [90.0, 90.0, 90.0])
            .with_energy(frame_idx as f32 * -100.0, frame_idx as f32 * 50.0)
            .with_thermo(300.0 + frame_idx as f32, 1.0 + frame_idx as f32 * 0.1);

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

    // Close writer and finalize the file
    writer.finish().expect("Failed to finish writing");

    // Open reader
    let mut reader = Reader::open(&file_path).expect("Failed to open reader");

    // Test sequential access
    println!("Testing sequential access...");
    for frame_idx in 0..20 {
        let frame = reader.read_frame(frame_idx).expect("Failed to read frame");

        // Verify frame metadata
        assert_eq!(frame.frame_number, frame_idx);
        assert!((frame.time - (frame_idx as f64 * 0.001)).abs() < 1e-10);

        // Verify positions (allowing for compression precision loss)
        verify_frame_data(
            &frame.data,
            &original_frames[frame_idx as usize],
            frame_idx as usize,
            "sequential",
        );
    }

    // Test random access - read frames in non-sequential order
    println!("Testing random access...");
    let random_indices = vec![
        19, 0, 15, 7, 3, 18, 1, 12, 8, 4, 16, 2, 11, 9, 5, 17, 6, 14, 10, 13,
    ];

    for &frame_idx in &random_indices {
        let frame = reader
            .read_frame(frame_idx)
            .expect("Failed to read frame in random order");

        // Verify frame metadata
        assert_eq!(frame.frame_number, frame_idx);
        assert!((frame.time - (frame_idx as f64 * 0.001)).abs() < 1e-10);

        // Verify all data types
        verify_frame_data(
            &frame.data,
            &original_frames[frame_idx as usize],
            frame_idx as usize,
            "random",
        );
    }

    // Test repeated random access to the same frame
    println!("Testing repeated access to same frame...");
    for _ in 0..5 {
        let frame = reader
            .read_frame(10)
            .expect("Failed to read frame repeatedly");
        assert_eq!(frame.frame_number, 10);
        verify_frame_data(&frame.data, &original_frames[10], 10, "repeated");
    }

    println!("Random access test completed successfully!");
}

fn verify_frame_data(
    actual: &FrameData,
    expected: &FrameData,
    frame_idx: usize,
    access_type: &str,
) {
    let tolerance = 1e-2; // Allow for compression precision loss

    // Verify positions
    if let (Some(actual_pos), Some(expected_pos)) = (&actual.positions, &expected.positions) {
        assert_eq!(actual_pos.shape(), expected_pos.shape());
        for i in 0..actual_pos.shape()[0] {
            for j in 0..actual_pos.shape()[1] {
                let diff = (actual_pos[[i, j]] - expected_pos[[i, j]]).abs();
                assert!(
                    diff < tolerance,
                    "Position mismatch at frame {} ({}): [{}, {}] actual={}, expected={}, diff={}",
                    frame_idx,
                    access_type,
                    i,
                    j,
                    actual_pos[[i, j]],
                    expected_pos[[i, j]],
                    diff
                );
            }
        }
    } else {
        panic!("Positions missing in frame {frame_idx} ({access_type})");
    }

    // Verify velocities
    if let (Some(actual_vel), Some(expected_vel)) = (&actual.velocities, &expected.velocities) {
        assert_eq!(actual_vel.shape(), expected_vel.shape());
        for i in 0..actual_vel.shape()[0] {
            for j in 0..actual_vel.shape()[1] {
                let diff = (actual_vel[[i, j]] - expected_vel[[i, j]]).abs();
                assert!(
                    diff < tolerance,
                    "Velocity mismatch at frame {} ({}): [{}, {}] actual={}, expected={}, diff={}",
                    frame_idx,
                    access_type,
                    i,
                    j,
                    actual_vel[[i, j]],
                    expected_vel[[i, j]],
                    diff
                );
            }
        }
    } else {
        panic!("Velocities missing in frame {frame_idx} ({access_type})");
    }

    // Verify forces
    if let (Some(actual_force), Some(expected_force)) = (&actual.forces, &expected.forces) {
        assert_eq!(actual_force.shape(), expected_force.shape());
        for i in 0..actual_force.shape()[0] {
            for j in 0..actual_force.shape()[1] {
                let diff = (actual_force[[i, j]] - expected_force[[i, j]]).abs();
                assert!(
                    diff < tolerance,
                    "Force mismatch at frame {} ({}): [{}, {}] actual={}, expected={}, diff={}",
                    frame_idx,
                    access_type,
                    i,
                    j,
                    actual_force[[i, j]],
                    expected_force[[i, j]],
                    diff
                );
            }
        }
    } else {
        panic!("Forces missing in frame {frame_idx} ({access_type})");
    }

    // Verify stress
    if let (Some(actual_stress), Some(expected_stress)) = (&actual.stress, &expected.stress) {
        assert_eq!(actual_stress.shape(), expected_stress.shape());
        for i in 0..actual_stress.shape()[0] {
            for j in 0..actual_stress.shape()[1] {
                let diff = (actual_stress[[i, j]] - expected_stress[[i, j]]).abs();
                assert!(
                    diff < tolerance,
                    "Stress mismatch at frame {} ({}): [{}, {}] actual={}, expected={}, diff={}",
                    frame_idx,
                    access_type,
                    i,
                    j,
                    actual_stress[[i, j]],
                    expected_stress[[i, j]],
                    diff
                );
            }
        }
    } else {
        panic!("Stress missing in frame {frame_idx} ({access_type})");
    }

    // Verify extra data
    assert_eq!(
        actual.extra.len(),
        expected.extra.len(),
        "Extra data count mismatch in frame {frame_idx} ({access_type})"
    );

    for (key, expected_array) in &expected.extra {
        let actual_array = actual.extra.get(key).unwrap_or_else(|| {
            panic!("Extra data '{key}' missing in frame {frame_idx} ({access_type})")
        });

        match (actual_array, expected_array) {
            (ExtraArray::Array1(actual_arr), ExtraArray::Array1(expected_arr)) => {
                assert_eq!(actual_arr.shape(), expected_arr.shape());
                for i in 0..actual_arr.len() {
                    let diff = (actual_arr[i] - expected_arr[i]).abs();
                    assert!(
                        diff < tolerance,
                        "Extra 1D '{}' mismatch at frame {} ({}): [{}] actual={}, expected={}, diff={}",
                        key,
                        frame_idx,
                        access_type,
                        i,
                        actual_arr[i],
                        expected_arr[i],
                        diff
                    );
                }
            }
            (ExtraArray::Array2(actual_arr), ExtraArray::Array2(expected_arr)) => {
                assert_eq!(actual_arr.shape(), expected_arr.shape());
                for i in 0..actual_arr.shape()[0] {
                    for j in 0..actual_arr.shape()[1] {
                        let diff = (actual_arr[[i, j]] - expected_arr[[i, j]]).abs();
                        assert!(
                            diff < tolerance,
                            "Extra 2D '{}' mismatch at frame {} ({}): [{}, {}] actual={}, expected={}, diff={}",
                            key,
                            frame_idx,
                            access_type,
                            i,
                            j,
                            actual_arr[[i, j]],
                            expected_arr[[i, j]],
                            diff
                        );
                    }
                }
            }
            (ExtraArray::Array3(actual_arr), ExtraArray::Array3(expected_arr)) => {
                assert_eq!(actual_arr.shape(), expected_arr.shape());
                for i in 0..actual_arr.shape()[0] {
                    for j in 0..actual_arr.shape()[1] {
                        for k in 0..actual_arr.shape()[2] {
                            let diff = (actual_arr[[i, j, k]] - expected_arr[[i, j, k]]).abs();
                            assert!(
                                diff < tolerance,
                                "Extra 3D '{}' mismatch at frame {} ({}): [{}, {}, {}] actual={}, expected={}, diff={}",
                                key,
                                frame_idx,
                                access_type,
                                i,
                                j,
                                k,
                                actual_arr[[i, j, k]],
                                expected_arr[[i, j, k]],
                                diff
                            );
                        }
                    }
                }
            }
            _ => {
                panic!("Extra data '{key}' dimension mismatch in frame {frame_idx} ({access_type})")
            }
        }
    }
}

#[test]
fn test_random_access_performance() {
    // This test verifies that random access doesn't require reading all previous frames
    let dir = tempdir().expect("Failed to create temporary directory");
    let file_path = dir.path().join("test_performance.taco");

    // Create metadata
    let sim_metadata = SimulationMetadata::default();
    let atom_metadata = AtomMetadata {
        names: Some(vec!["H".to_string()]),
        elements: Some(vec!["H".to_string()]),
        masses: Some(vec![1.008]),
        ..Default::default()
    };

    let compression_settings = CompressionSettings::default();

    // Create writer and write many frames
    let mut writer = Writer::create(
        &file_path,
        1,     // num_atoms
        0.001, // time_step
        sim_metadata,
        atom_metadata,
        compression_settings,
    )
    .expect("Failed to create writer");

    let num_frames = 1000;
    for frame_idx in 0..num_frames {
        let positions = Array2::<f32>::from_shape_vec(
            (1, 3),
            vec![
                frame_idx as f32,
                frame_idx as f32 + 1.0,
                frame_idx as f32 + 2.0,
            ],
        )
        .unwrap();

        let mut frame_data = FrameData::new(positions);
        // Add some extra data
        frame_data.set_extra_1d("index", Array1::<f32>::from(vec![frame_idx as f32]));

        let frame = Frame::new(frame_idx as u64, frame_idx as f64 * 0.001, frame_data);
        writer.write_frame(frame).expect("Failed to write frame");
    }

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

    // Test random access performance
    let mut reader = Reader::open(&file_path).expect("Failed to open reader");

    // Read some frames in random order
    let test_indices = vec![999, 500, 100, 800, 50, 900, 10, 750];

    for &idx in &test_indices {
        let frame = reader.read_frame(idx).expect("Failed to read frame");
        assert_eq!(frame.frame_number, idx);

        // Verify extra data
        match frame.data.get_extra("index").unwrap() {
            ExtraArray::Array1(arr) => {
                assert!((arr[0] - idx as f32).abs() < 1e-5);
            }
            _ => panic!("index not Array1"),
        }
    }

    println!("Performance test completed successfully!");
}
