//! Tests for FrameData extra and stress fields, including Array1, Array2, Array3
//! This file tests FrameData's support for stress and extra arrays (1D, 2D, 3D)
//! and roundtrip serialization/deserialization.

use ndarray::{Array1, Array2, Array3};
use std::collections::HashMap;
use std::io::Cursor;
use taco_format::{CompressionSettings, ExtraDatasetSpec, Frame, FrameData};

fn array2_allclose(a: &ndarray::Array2<f32>, b: &ndarray::Array2<f32>, tol: f32) -> bool {
    if a.shape() != b.shape() {
        return false;
    }
    a.iter().zip(b.iter()).all(|(x, y)| (*x - *y).abs() <= tol)
}

#[test]
fn test_framedata_stress_and_extra_roundtrip() {
    let compression_settings = CompressionSettings::default();

    // Create test positions
    let positions =
        Array2::<f32>::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    let mut frame_data = FrameData::new(positions.clone());

    // Add stress (2x6)
    let stress = Array2::<f32>::from_shape_vec(
        (2, 6),
        vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2],
    )
    .unwrap();
    frame_data.set_stress(stress.clone());

    // Add extra Array1
    let arr1 = Array1::<f32>::from(vec![9.0, 8.0]);
    frame_data.set_extra_1d("arr1", arr1.clone());
    // Add extra Array2
    let arr2 = Array2::<f32>::from_shape_vec((2, 2), vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    frame_data.set_extra_2d("arr2", arr2.clone());
    // Add extra Array3
    let arr3 =
        Array3::<f32>::from_shape_vec((2, 1, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    frame_data.set_extra_3d("arr3", arr3.clone());

    // Serialize and deserialize
    let mut buffer = Cursor::new(Vec::new());
    let extra_specs: HashMap<String, ExtraDatasetSpec> = HashMap::new();
    frame_data
        .write(&mut buffer, &compression_settings, &extra_specs)
        .unwrap();
    buffer.set_position(0);
    let frame_data2 = FrameData::read(&mut buffer, &compression_settings, &extra_specs).unwrap();

    // Check stress (allow for small floating point error due to compression)
    assert!(array2_allclose(
        frame_data2.stress.as_ref().unwrap(),
        &stress,
        1e-3
    ));
    // Check extra Array1
    match frame_data2.get_extra("arr1").unwrap() {
        taco_format::ExtraArray::Array1(a) => assert_eq!(a, &arr1),
        _ => panic!("arr1 not Array1"),
    }
    // Check extra Array2
    match frame_data2.get_extra("arr2").unwrap() {
        taco_format::ExtraArray::Array2(a) => assert_eq!(a, &arr2),
        _ => panic!("arr2 not Array2"),
    }
    // Check extra Array3
    match frame_data2.get_extra("arr3").unwrap() {
        taco_format::ExtraArray::Array3(a) => assert_eq!(a, &arr3),
        _ => panic!("arr3 not Array3"),
    }
}

#[test]
fn test_frame_with_stress_and_extra_roundtrip() {
    let compression_settings = CompressionSettings::default();
    let positions = Array2::<f32>::from_shape_vec((1, 3), vec![1.0, 2.0, 3.0]).unwrap();
    let mut frame_data = FrameData::new(positions);
    let stress = Array2::<f32>::from_shape_vec((1, 6), vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6]).unwrap();
    frame_data.set_stress(stress.clone());
    let arr1 = Array1::<f32>::from(vec![42.0]);
    frame_data.set_extra_1d("foo", arr1.clone());
    let frame = Frame::new(0, 0.0, frame_data);
    let mut buffer = Cursor::new(Vec::new());
    let extra_specs: HashMap<String, ExtraDatasetSpec> = HashMap::new();
    frame
        .write(&mut buffer, &compression_settings, &extra_specs)
        .unwrap();
    buffer.set_position(0);
    let frame2 = Frame::read(&mut buffer, &compression_settings, &extra_specs).unwrap();
    assert!(array2_allclose(
        frame2.data.stress.as_ref().unwrap(),
        &stress,
        1e-3
    ));
    match frame2.data.get_extra("foo").unwrap() {
        taco_format::ExtraArray::Array1(a) => assert_eq!(a, &arr1),
        _ => panic!("foo not Array1"),
    }
}
