//! Integration tests for file I/O, full/delta frames, and PBC settings

#[cfg(test)]
mod integration_tests {
    use ndarray::Array2;
    use taco_format::{AtomMetadata, CompressionSettings, SimulationMetadata};
    use taco_format::{Frame, FrameData};
    use taco_format::{Reader, Writer};
    use tempfile::tempdir;

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

        // Create metadata
        let sim_metadata = SimulationMetadata {
            name: Some("Test Simulation".to_string()),
            temperature: Some(300.0),
            pressure: Some(1.0),
            ensemble: Some("NPT".to_string()),
            ..Default::default()
        };

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

        // Create compression settings - will use default half precision for delta frames
        let compression_settings = CompressionSettings::default();

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

        // Create 10 frames with a simple trajectory
        let mut positions = Array2::<f32>::zeros((3, 3));
        for frame_idx in 0..10 {
            // Update positions
            for i in 0..3 {
                for j in 0..3 {
                    positions[[i, j]] = (i as f32) + (j as f32) * 0.1 + (frame_idx as f32) * 0.01;
                }
            }

            // Create frame
            let frame_data = FrameData::new(positions.clone());
            let frame = Frame::new(frame_idx, frame_idx as f64 * 0.002, frame_data)
                .with_box_dimensions([10.0, 10.0, 10.0], [90.0, 90.0, 90.0])
                .with_thermo(300.0, 1.0);

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

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

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

        // Read header
        let header = reader.header();
        assert_eq!(header.num_atoms, 3);
        assert_eq!(header.num_frames, 10);
        assert_eq!(header.time_step, 0.002);
        assert_eq!(
            header.simulation_metadata.name,
            Some("Test Simulation".to_string())
        );

        // Read frames
        for i in 0..10 {
            let frame = reader.read_frame(i).expect("Failed to read frame");

            assert_eq!(frame.frame_number, i);
            assert!((frame.time - i as f64 * 0.002).abs() < 0.0001);

            // Check positions
            if let Some(pos) = &frame.data.positions {
                for atom_idx in 0..3 {
                    for coord_idx in 0..3 {
                        let expected =
                            (atom_idx as f32) + (coord_idx as f32) * 0.1 + (i as f32) * 0.01;
                        assert!((pos[[atom_idx, coord_idx]] - expected).abs() < 0.001);
                    }
                }
            } else {
                panic!("Positions missing in frame {i}");
            }
        }
    }

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

        println!("===== Starting test_full_and_delta_frames =====");

        // Create metadata
        let sim_metadata = SimulationMetadata::default();
        let atom_metadata = AtomMetadata::default();

        // Create compression settings
        let compression_settings = CompressionSettings::default();

        // Create header with custom full_frame_interval
        let mut header = taco_format::Header::new(
            2,     // num_atoms
            0.001, // time_step
            sim_metadata,
            atom_metadata,
            compression_settings,
        );
        header.full_frame_interval = 3; // Set a full frame every 3 frames

        println!(
            "Using full_frame_interval = {} for test",
            header.full_frame_interval
        );

        // Create writer with custom header
        let mut writer =
            Writer::create_with_header(&file_path, header).expect("Failed to create writer");

        // Create 10 frames
        let mut positions = Array2::<f32>::zeros((2, 3));
        let mut expected_positions = Vec::with_capacity(10);

        for frame_idx in 0..10 {
            // Update positions
            for i in 0..2 {
                for j in 0..3 {
                    positions[[i, j]] = (i as f32) + (j as f32) * 0.1 + (frame_idx as f32) * 0.01;
                }
            }

            // Store expected positions for later verification
            expected_positions.push(positions.clone());

            println!("Writing frame {frame_idx}: positions = {positions:?}");

            // Create frame
            let frame_data = FrameData::new(positions.clone());
            let frame = Frame::new(frame_idx, frame_idx 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");

        println!("Finished writing frames. Reading back frames...");

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

        // Check that we have the expected number of frames
        assert_eq!(reader.header().num_frames, 10);
        assert_eq!(reader.header().full_frame_interval, 3);
        println!(
            "Verified header has {} frames and full_frame_interval={}",
            reader.header().num_frames,
            reader.header().full_frame_interval
        );

        // Read all frames and check full/delta status
        for i in 0..10 {
            let frame = reader.read_frame(i).expect("Failed to read frame");

            // Verify frame number and time
            assert_eq!(frame.frame_number, i);
            assert!((frame.time - i as f64 * 0.001).abs() < 0.0001);

            // Print info about the frame we're reading
            println!(
                "Reading frame {}: frame_type={}",
                i,
                if frame.is_full_frame { "full" } else { "delta" }
            );

            // Check positions against expected values with detailed diagnostics
            if let Some(pos) = &frame.data.positions {
                println!("  Frame {} positions shape: {:?}", i, pos.shape());
                println!("  Frame {i} positions: {pos:?}");

                // Compare with expected positions
                let expected = &expected_positions[i as usize];

                for atom_idx in 0..2 {
                    for coord_idx in 0..3 {
                        let expected_val = expected[[atom_idx, coord_idx]];
                        let actual_val = pos[[atom_idx, coord_idx]];
                        let diff = (actual_val - expected_val).abs();

                        if diff >= 0.001 {
                            println!(
                                "  ERROR at atom {atom_idx}, coord {coord_idx}: expected={expected_val}, actual={actual_val}, diff={diff}"
                            );
                        }

                        assert!(
                            diff < 0.001,
                            "Frame {i} position mismatch at [{atom_idx},{coord_idx}]: expected {expected_val}, got {actual_val}, diff={diff}"
                        );
                    }
                }
            } else {
                panic!("Positions missing in frame {i}");
            }

            // Frames 0, 3, 6, 9 should be full frames
            let expected_full_frame = i % 3 == 0;

            assert_eq!(
                frame.is_full_frame,
                expected_full_frame,
                "Frame {} should be {} but was {}",
                i,
                if expected_full_frame {
                    "a full frame"
                } else {
                    "a delta frame"
                },
                if frame.is_full_frame {
                    "a full frame"
                } else {
                    "a delta frame"
                }
            );
        }

        println!("===== test_full_and_delta_frames completed successfully =====");
    }

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

        println!("===== Starting test_pbc_settings =====");

        // Create metadata
        let sim_metadata = SimulationMetadata::default();
        let atom_metadata = AtomMetadata::default();

        // Create compression settings
        let compression_settings = CompressionSettings::default();

        // Create header with custom PBC settings
        let mut header = taco_format::Header::new(
            3,     // num_atoms
            0.001, // time_step
            sim_metadata.clone(),
            atom_metadata.clone(),
            compression_settings.clone(),
        );

        // Set specific PBC settings - periodic in x and z, non-periodic in y
        header.pbc = [true, false, true];
        println!("Setting PBC to [true, false, true]");

        // Create writer with custom header
        let mut writer =
            Writer::create_with_header(&file_path, header).expect("Failed to create writer");

        // Create a simple frame with 3 atoms
        let mut positions = Array2::<f32>::zeros((3, 3));
        for i in 0..3 {
            for j in 0..3 {
                positions[[i, j]] = (i as f32) + (j as f32) * 0.1;
            }
        }

        // Create frame
        let frame_data = FrameData::new(positions.clone());
        let frame = Frame::new(0, 0.0, frame_data)
            .with_box_dimensions([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);

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

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

        println!("Finished writing frame with PBC [true, false, true]. Reading back...");

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

        // Verify PBC settings in header
        let header = reader.header();
        println!(
            "Read PBC: [{}, {}, {}]",
            header.pbc[0], header.pbc[1], header.pbc[2]
        );
        assert_eq!(header.pbc, [true, false, true], "PBC settings do not match");

        // Next, test different PBC settings
        println!("Testing PBC [false, true, false]");

        // Create a new file with different PBC settings
        let file_path2 = dir.path().join("test_pbc2.taco");

        // Create header with different PBC settings
        let mut header2 = taco_format::Header::new(
            3,     // num_atoms
            0.001, // time_step
            sim_metadata.clone(),
            atom_metadata.clone(),
            compression_settings.clone(),
        );

        // Set different PBC settings - non-periodic in x and z, periodic in y
        header2.pbc = [false, true, false];

        // Create writer with custom header
        let mut writer2 =
            Writer::create_with_header(&file_path2, header2).expect("Failed to create writer");

        // Create a simple frame
        let frame_data = FrameData::new(positions.clone());
        let frame = Frame::new(0, 0.0, frame_data)
            .with_box_dimensions([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);

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

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

        // Open file for reading
        let reader2 = Reader::open(&file_path2).expect("Failed to open file for reading");

        // Verify PBC settings in header
        let header2 = reader2.header();
        assert_eq!(
            header2.pbc,
            [false, true, false],
            "PBC settings do not match"
        );

        // Test all non-periodic
        println!("Testing PBC [false, false, false]");

        // Create a new file with all non-periodic settings
        let file_path3 = dir.path().join("test_pbc3.taco");

        // Create header with non-periodic settings
        let mut header3 = taco_format::Header::new(
            3,     // num_atoms
            0.001, // time_step
            sim_metadata.clone(),
            atom_metadata.clone(),
            compression_settings.clone(),
        );

        // Set non-periodic in all directions
        header3.pbc = [false, false, false];

        // Create writer with custom header
        let mut writer3 =
            Writer::create_with_header(&file_path3, header3).expect("Failed to create writer");

        // Create a simple frame
        let frame_data = FrameData::new(positions.clone());
        let frame = Frame::new(0, 0.0, frame_data)
            .with_box_dimensions([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);

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

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

        // Open file for reading
        let reader3 = Reader::open(&file_path3).expect("Failed to open file for reading");

        // Verify PBC settings in header
        let header3 = reader3.header();
        assert_eq!(
            header3.pbc,
            [false, false, false],
            "PBC settings do not match"
        );

        println!("===== test_pbc_settings completed successfully =====");
    }
}
