//! Frame serialization, delta encoding, and related tests

#[cfg(test)]
mod frame_tests {
    use ndarray::{Array1, Array2};
    use std::io::Cursor;
    use taco_format::CompressionSettings;
    use taco_format::{ExtraArray, Frame, FrameData};

    #[test]
    fn test_frame_serialization_roundtrip() {
        // Create a compression settings
        let compression_settings = CompressionSettings::default();

        // Create test positions
        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 velocities
        let mut velocities = Array2::<f32>::zeros((3, 3));
        for i in 0..3 {
            for j in 0..3 {
                velocities[[i, j]] = (i as f32) * 0.5 + (j as f32) * 0.05;
            }
        }

        // Create frame data
        let frame_data = FrameData::new(positions.clone()).with_velocities(velocities.clone());

        // Create frame
        let frame = Frame::new(0, 0.0, frame_data)
            .with_box_dimensions([10.0, 10.0, 10.0], [90.0, 90.0, 90.0])
            .with_thermo(300.0, 1.0);

        // Serialize
        let mut buffer = Cursor::new(Vec::new());
        let extra_specs: std::collections::HashMap<String, taco_format::ExtraDatasetSpec> =
            std::collections::HashMap::new();
        let bytes_written = frame
            .write(&mut buffer, &compression_settings, &extra_specs)
            .expect("Failed to write frame");

        // Make sure some data was written
        assert!(bytes_written > 0);

        // Deserialize
        buffer.set_position(0);
        let frame2 = Frame::read(&mut buffer, &compression_settings, &extra_specs)
            .expect("Failed to read frame");

        // Compare metadata
        assert_eq!(frame.frame_number, frame2.frame_number);
        assert_eq!(frame.time, frame2.time);
        assert_eq!(frame.box_dimensions, frame2.box_dimensions);
        assert_eq!(frame.temperature, frame2.temperature);

        // Compare positions (might have small differences due to compression)
        if let (Some(pos1), Some(pos2)) = (&frame.data.positions, &frame2.data.positions) {
            assert_eq!(pos1.shape(), pos2.shape());
            for i in 0..pos1.shape()[0] {
                for j in 0..pos1.shape()[1] {
                    assert!((pos1[[i, j]] - pos2[[i, j]]).abs() < 0.001);
                }
            }
        } else {
            panic!("Positions missing after deserialization");
        }
    }

    #[test]
    fn test_frame_delta_encoding() {
        // Create two frames with slightly different coordinates
        let mut positions1 = Array2::<f32>::zeros((3, 3));
        let mut positions2 = Array2::<f32>::zeros((3, 3));

        for i in 0..3 {
            for j in 0..3 {
                positions1[[i, j]] = (i as f32) + (j as f32) * 0.1;
                positions2[[i, j]] = positions1[[i, j]] + 0.01; // Small change
            }
        }

        let frame_data1 = FrameData::new(positions1);
        let frame_data2 = FrameData::new(positions2);

        // Calculate delta
        let delta = frame_data2
            .delta_from(&frame_data1)
            .expect("Failed to calculate delta");

        // Apply delta to frame1 to get frame2
        let reconstructed = frame_data1
            .apply_delta(&delta)
            .expect("Failed to apply delta");

        // Compare
        if let (Some(pos_orig), Some(pos_recon)) =
            (&frame_data2.positions, &reconstructed.positions)
        {
            assert_eq!(pos_orig.shape(), pos_recon.shape());
            for i in 0..pos_orig.shape()[0] {
                for j in 0..pos_orig.shape()[1] {
                    assert!((pos_orig[[i, j]] - pos_recon[[i, j]]).abs() < 0.0001);
                }
            }
        } else {
            panic!("Positions missing after delta application");
        }
    }

    #[test]
    fn test_frame_delta_encoding_with_stress_and_extra() {
        use ndarray::{Array1, Array2, Array3};
        use taco_format::ExtraArray;

        // Create two frames with slightly different coordinates, stress, and extra data
        let mut positions1 = Array2::<f32>::zeros((3, 3));
        let mut positions2 = Array2::<f32>::zeros((3, 3));

        for i in 0..3 {
            for j in 0..3 {
                positions1[[i, j]] = (i as f32) + (j as f32) * 0.1;
                positions2[[i, j]] = positions1[[i, j]] + 0.01; // Small change
            }
        }

        // Create stress tensors
        let stress1 = Array2::<f32>::from_shape_vec(
            (3, 6),
            vec![
                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, //
                7.0, 8.0, 9.0, 10.0, 11.0, 12.0, //
                13.0, 14.0, 15.0, 16.0, 17.0, 18.0,
            ],
        )
        .unwrap();
        let stress2 = stress1.mapv(|x| x + 0.1); // Slightly different stress

        // Create extra data
        let extra1d_1 = Array1::<f32>::from(vec![1.0, 2.0, 3.0]);
        let extra1d_2 = extra1d_1.mapv(|x| x + 0.05);

        let extra2d_1 =
            Array2::<f32>::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let extra2d_2 = extra2d_1.mapv(|x| x + 0.02);

        let extra3d_1 =
            Array3::<f32>::from_shape_vec((2, 2, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
                .unwrap();
        let extra3d_2 = extra3d_1.mapv(|x| x + 0.03);

        // Create frame data
        let mut frame_data1 = FrameData::new(positions1).with_stress(stress1.clone());
        frame_data1.set_extra_1d("extra1d", extra1d_1.clone());
        frame_data1.set_extra_2d("extra2d", extra2d_1.clone());
        frame_data1.set_extra_3d("extra3d", extra3d_1.clone());

        let mut frame_data2 = FrameData::new(positions2).with_stress(stress2.clone());
        frame_data2.set_extra_1d("extra1d", extra1d_2.clone());
        frame_data2.set_extra_2d("extra2d", extra2d_2.clone());
        frame_data2.set_extra_3d("extra3d", extra3d_2.clone());

        // Calculate delta
        let delta = frame_data2
            .delta_from(&frame_data1)
            .expect("Failed to calculate delta");

        // Apply delta to frame1 to get frame2
        let reconstructed = frame_data1
            .apply_delta(&delta)
            .expect("Failed to apply delta");

        // Compare positions
        if let (Some(pos_orig), Some(pos_recon)) =
            (&frame_data2.positions, &reconstructed.positions)
        {
            assert_eq!(pos_orig.shape(), pos_recon.shape());
            for i in 0..pos_orig.shape()[0] {
                for j in 0..pos_orig.shape()[1] {
                    assert!((pos_orig[[i, j]] - pos_recon[[i, j]]).abs() < 0.0001);
                }
            }
        } else {
            panic!("Positions missing after delta application");
        }

        // Compare stress
        if let (Some(stress_orig), Some(stress_recon)) =
            (&frame_data2.stress, &reconstructed.stress)
        {
            assert_eq!(stress_orig.shape(), stress_recon.shape());
            for i in 0..stress_orig.shape()[0] {
                for j in 0..stress_orig.shape()[1] {
                    assert!((stress_orig[[i, j]] - stress_recon[[i, j]]).abs() < 0.0001);
                }
            }
        } else {
            panic!("Stress missing after delta application");
        }

        // Compare extra 1D data
        match reconstructed.get_extra("extra1d").unwrap() {
            ExtraArray::Array1(arr_recon) => {
                for i in 0..extra1d_2.len() {
                    assert!((extra1d_2[i] - arr_recon[i]).abs() < 0.0001);
                }
            }
            _ => panic!("extra1d not Array1"),
        }

        // Compare extra 2D data
        match reconstructed.get_extra("extra2d").unwrap() {
            ExtraArray::Array2(arr_recon) => {
                for i in 0..extra2d_2.shape()[0] {
                    for j in 0..extra2d_2.shape()[1] {
                        assert!((extra2d_2[[i, j]] - arr_recon[[i, j]]).abs() < 0.0001);
                    }
                }
            }
            _ => panic!("extra2d not Array2"),
        }

        // Compare extra 3D data
        match reconstructed.get_extra("extra3d").unwrap() {
            ExtraArray::Array3(arr_recon) => {
                for i in 0..extra3d_2.shape()[0] {
                    for j in 0..extra3d_2.shape()[1] {
                        for k in 0..extra3d_2.shape()[2] {
                            assert!((extra3d_2[[i, j, k]] - arr_recon[[i, j, k]]).abs() < 0.0001);
                        }
                    }
                }
            }
            _ => panic!("extra3d not Array3"),
        }
    }

    #[test]
    fn test_frame_delta_encoding_new_extra_data() {
        // Test when delta frame has extra data that reference frame doesn't have
        let positions1 = Array2::<f32>::zeros((2, 3));
        let positions2 = Array2::<f32>::ones((2, 3));

        let frame_data1 = FrameData::new(positions1);
        let mut frame_data2 = FrameData::new(positions2);

        // Add extra data only to frame2
        let extra_data = ndarray::Array1::<f32>::from(vec![10.0, 20.0]);
        frame_data2.set_extra_1d("new_extra", extra_data.clone());

        // Calculate delta
        let delta = frame_data2
            .delta_from(&frame_data1)
            .expect("Failed to calculate delta");

        // Apply delta to frame1
        let reconstructed = frame_data1
            .apply_delta(&delta)
            .expect("Failed to apply delta");

        // Check that new extra data is present
        match reconstructed.get_extra("new_extra").unwrap() {
            ExtraArray::Array1(arr) => {
                assert_eq!(arr, &extra_data);
            }
            _ => panic!("new_extra not Array1"),
        }
    }

    #[test]
    fn test_frame_delta_encoding_dimension_mismatch() {
        // Test error handling for dimension mismatches in extra data
        let positions = Array2::<f32>::zeros((2, 3));
        let mut frame_data1 = FrameData::new(positions.clone());
        let mut frame_data2 = FrameData::new(positions);

        // Add different dimension extra data with same key
        let extra1d = Array1::<f32>::from(vec![1.0, 2.0]);
        let extra2d = Array2::<f32>::from_shape_vec((2, 1), vec![1.0, 2.0]).unwrap();

        frame_data1.set_extra_1d("conflicting", extra1d);
        frame_data2.set_extra_2d("conflicting", extra2d);

        // This should fail with dimension mismatch
        let result = frame_data2.delta_from(&frame_data1);
        assert!(result.is_err());
        if let Err(e) = result {
            let error_msg = format!("{e}");
            assert!(error_msg.contains("mismatch"));
        }
    }
}
