//! Frame structures and operations for TACO format
//!
//! This module contains structures for representing frames in trajectory data,
//! including positions, velocities, and forces.

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use ndarray::{Array1, Array2, Array3};
use std::collections::HashMap;
use std::io::{Read, Write};

use crate::compression::{
    CompressionSettings, PrecisionMode, compress_tensor, compress_tensor_f64, decompress_tensor,
    decompress_tensor_f64,
};
use crate::header::{ExtraCodec, ExtraDataType, ExtraDatasetSpec};
use crate::{Error, Result};
use fpzip_rs::{
    FpZipCompressor, decompress_f32 as fpzip_decompress_f32, decompress_f64 as fpzip_decompress_f64,
};

/// Enum for storing extra arrays of different dimensions
#[derive(Debug, Clone, PartialEq)]
pub enum ExtraArray {
    Array1(Array1<f32>),
    Array2(Array2<f32>),
    Array3(Array3<f32>),
    Array1F64(Array1<f64>),
    Array2F64(Array2<f64>),
    Array3F64(Array3<f64>),
}

impl ExtraArray {
    pub fn ndim(&self) -> usize {
        match self {
            ExtraArray::Array1(_) => 1,
            ExtraArray::Array2(_) => 2,
            ExtraArray::Array3(_) => 3,
            ExtraArray::Array1F64(_) => 1,
            ExtraArray::Array2F64(_) => 2,
            ExtraArray::Array3F64(_) => 3,
        }
    }

    pub fn dtype(&self) -> ExtraDataType {
        match self {
            ExtraArray::Array1(_) | ExtraArray::Array2(_) | ExtraArray::Array3(_) => {
                ExtraDataType::F32
            }
            ExtraArray::Array1F64(_) | ExtraArray::Array2F64(_) | ExtraArray::Array3F64(_) => {
                ExtraDataType::F64
            }
        }
    }

    pub fn shape(&self) -> Vec<u32> {
        match self {
            ExtraArray::Array1(arr) => vec![arr.len() as u32],
            ExtraArray::Array2(arr) => vec![arr.shape()[0] as u32, arr.shape()[1] as u32],
            ExtraArray::Array3(arr) => vec![
                arr.shape()[0] as u32,
                arr.shape()[1] as u32,
                arr.shape()[2] as u32,
            ],
            ExtraArray::Array1F64(arr) => vec![arr.len() as u32],
            ExtraArray::Array2F64(arr) => vec![arr.shape()[0] as u32, arr.shape()[1] as u32],
            ExtraArray::Array3F64(arr) => vec![
                arr.shape()[0] as u32,
                arr.shape()[1] as u32,
                arr.shape()[2] as u32,
            ],
        }
    }
}

/// A single frame in a molecular dynamics trajectory
#[derive(Debug, Clone)]
pub struct Frame {
    /// Frame number (sequential, starting from 0)
    pub frame_number: u64,

    /// Time in picoseconds
    pub time: f64,

    /// Coordinate data (positions, velocities, forces)
    pub data: FrameData,

    /// Box dimensions (x, y, z) in nanometers
    pub box_dimensions: Option<[f32; 3]>,

    /// Box angles (alpha, beta, gamma) in degrees
    pub box_angles: Option<[f32; 3]>,

    /// Optional full box vectors (rows are lattice vectors a, b, c) in
    /// nanometers. Set when the caller needs to preserve the cell's 3D
    /// orientation through a round-trip — otherwise lengths+angles are
    /// sufficient (but only define a canonical orientation).
    /// Stored at the end of the frame so old readers can ignore it.
    pub box_vectors: Option<[[f32; 3]; 3]>,

    /// Potential energy in kJ/mol
    pub potential_energy: Option<f32>,

    /// Kinetic energy in kJ/mol
    pub kinetic_energy: Option<f32>,

    /// Temperature in Kelvin
    pub temperature: Option<f32>,

    /// Pressure in bars
    pub pressure: Option<f32>,

    /// Flag indicating if this is a full frame or delta frame
    pub is_full_frame: bool,
}

impl Frame {
    /// Create a new frame with the given data
    pub fn new(frame_number: u64, time: f64, data: FrameData) -> Self {
        Frame {
            frame_number,
            time,
            data,
            box_dimensions: None,
            box_angles: None,
            box_vectors: None,
            potential_energy: None,
            kinetic_energy: None,
            temperature: None,
            pressure: None,
            is_full_frame: true, // Default to full frame
        }
    }

    /// Set the full box vectors (rows are lattice vectors a, b, c) in nanometers.
    /// Use this when the cell's 3D orientation must round-trip exactly;
    /// `box_dimensions` and `box_angles` only carry the cell's shape, not its
    /// orientation in space.
    pub fn with_box_vectors(mut self, vectors: [[f32; 3]; 3]) -> Self {
        self.box_vectors = Some(vectors);
        self
    }

    /// Set box dimensions for this frame
    pub fn with_box_dimensions(mut self, dimensions: [f32; 3], angles: [f32; 3]) -> Self {
        self.box_dimensions = Some(dimensions);
        self.box_angles = Some(angles);
        self
    }

    /// Set energy values for this frame
    pub fn with_energy(mut self, potential: f32, kinetic: f32) -> Self {
        self.potential_energy = Some(potential);
        self.kinetic_energy = Some(kinetic);
        self
    }

    /// Set temperature and pressure for this frame
    pub fn with_thermo(mut self, temperature: f32, pressure: f32) -> Self {
        self.temperature = Some(temperature);
        self.pressure = Some(pressure);
        self
    }

    /// Write the frame to a writer with the given compression settings
    pub fn write<W: Write + std::io::Seek>(
        &self,
        writer: &mut W,
        settings: &CompressionSettings,
        extra_specs: &HashMap<String, ExtraDatasetSpec>,
    ) -> Result<u64> {
        let start_pos = writer.stream_position()?;

        // Write frame metadata
        writer.write_u64::<LittleEndian>(self.frame_number)?;
        writer.write_f64::<LittleEndian>(self.time)?;

        // Write frame data
        self.data.write(writer, settings, extra_specs)?;

        // Write box dimensions if present
        if let Some(dimensions) = self.box_dimensions {
            writer.write_u8(1)?; // Box dimensions present
            for value in dimensions.iter() {
                writer.write_f32::<LittleEndian>(*value)?;
            }

            if let Some(angles) = self.box_angles {
                writer.write_u8(1)?; // Box angles present
                for value in angles.iter() {
                    writer.write_f32::<LittleEndian>(*value)?;
                }
            } else {
                writer.write_u8(0)?; // No box angles
            }
        } else {
            writer.write_u8(0)?; // No box dimensions
        }

        // Write energies if present
        if self.potential_energy.is_some() || self.kinetic_energy.is_some() {
            writer.write_u8(1)?; // Energies present
            let pot = self.potential_energy.unwrap_or(f32::NAN);
            let kin = self.kinetic_energy.unwrap_or(f32::NAN);
            writer.write_f32::<LittleEndian>(pot)?;
            writer.write_f32::<LittleEndian>(kin)?;
        } else {
            writer.write_u8(0)?; // No energies
        }

        // Write temperature and pressure independently
        // Temperature
        if let Some(temp) = self.temperature {
            writer.write_u8(1)?; // Temperature present
            writer.write_f32::<LittleEndian>(temp)?;
        } else {
            writer.write_u8(0)?; // No temperature
        }

        // Pressure
        if let Some(press) = self.pressure {
            writer.write_u8(1)?; // Pressure present
            writer.write_f32::<LittleEndian>(press)?;
        } else {
            writer.write_u8(0)?; // No pressure
        }

        // Optional full box vectors (9 floats). Appended after pressure for
        // backwards compatibility — older readers stop at pressure and ignore
        // the trailing block.
        if let Some(vectors) = self.box_vectors {
            writer.write_u8(1)?; // Box vectors present
            for row in vectors.iter() {
                for value in row.iter() {
                    writer.write_f32::<LittleEndian>(*value)?;
                }
            }
        } else {
            writer.write_u8(0)?; // No box vectors
        }

        Ok(writer.stream_position()? - start_pos)
    }

    /// Write the frame to a Vec<u8> with the given compression settings
    pub fn write_to_vec(
        &self,
        buffer: &mut Vec<u8>,
        settings: &CompressionSettings,
        extra_specs: &HashMap<String, ExtraDatasetSpec>,
    ) -> Result<()> {
        // Write frame metadata
        buffer.extend_from_slice(&self.frame_number.to_le_bytes());
        buffer.extend_from_slice(&self.time.to_le_bytes());

        // Write frame data
        self.data.write(buffer, settings, extra_specs)?;

        // Write box dimensions if present
        if let Some(dimensions) = self.box_dimensions {
            buffer.push(1); // Box dimensions present
            for value in dimensions.iter() {
                buffer.extend_from_slice(&value.to_le_bytes());
            }

            if let Some(angles) = self.box_angles {
                buffer.push(1); // Box angles present
                for value in angles.iter() {
                    buffer.extend_from_slice(&value.to_le_bytes());
                }
            } else {
                buffer.push(0); // No box angles
            }
        } else {
            buffer.push(0); // No box dimensions
        }

        // Write energies if present
        if self.potential_energy.is_some() || self.kinetic_energy.is_some() {
            buffer.push(1); // Energies present
            let pot = self.potential_energy.unwrap_or(f32::NAN);
            let kin = self.kinetic_energy.unwrap_or(f32::NAN);
            buffer.extend_from_slice(&pot.to_le_bytes());
            buffer.extend_from_slice(&kin.to_le_bytes());
        } else {
            buffer.push(0); // No energies
        }

        // Write temperature and pressure independently
        // Temperature
        if let Some(temp) = self.temperature {
            buffer.push(1); // Temperature present
            buffer.extend_from_slice(&temp.to_le_bytes());
        } else {
            buffer.push(0); // No temperature
        }

        // Pressure
        if let Some(press) = self.pressure {
            buffer.push(1); // Pressure present
            buffer.extend_from_slice(&press.to_le_bytes());
        } else {
            buffer.push(0); // No pressure
        }

        // Optional full box vectors (see `write` for rationale).
        if let Some(vectors) = self.box_vectors {
            buffer.push(1);
            for row in vectors.iter() {
                for value in row.iter() {
                    buffer.extend_from_slice(&value.to_le_bytes());
                }
            }
        } else {
            buffer.push(0);
        }

        Ok(())
    }

    /// Read a frame from a reader with the given compression settings
    pub fn read<R: Read>(
        reader: &mut R,
        settings: &CompressionSettings,
        extra_specs: &HashMap<String, ExtraDatasetSpec>,
    ) -> Result<Self> {
        // Read frame metadata
        let frame_number = reader.read_u64::<LittleEndian>()?;
        let time = reader.read_f64::<LittleEndian>()?;

        // Read frame data
        let data = FrameData::read(reader, settings, extra_specs)?;

        let mut frame = Frame::new(frame_number, time, data);

        // The flag is set later by the Reader based on the frame type byte read before this function is called

        // Read box dimensions if present
        let has_box = reader.read_u8()? != 0;
        if has_box {
            let mut dimensions = [0.0; 3];
            for value in dimensions.iter_mut() {
                *value = reader.read_f32::<LittleEndian>()?;
            }
            frame.box_dimensions = Some(dimensions);

            let has_angles = reader.read_u8()? != 0;
            if has_angles {
                let mut angles = [0.0; 3];
                for value in angles.iter_mut() {
                    *value = reader.read_f32::<LittleEndian>()?;
                }
                frame.box_angles = Some(angles);
            }
        }

        // Read energies if present
        let has_energy = reader.read_u8()? != 0;
        if has_energy {
            let potential = reader.read_f32::<LittleEndian>()?;
            let kinetic = reader.read_f32::<LittleEndian>()?;
            frame.potential_energy = if potential.is_nan() {
                None
            } else {
                Some(potential)
            };
            frame.kinetic_energy = if kinetic.is_nan() {
                None
            } else {
                Some(kinetic)
            };
        }

        // Read temperature and pressure independently
        // Temperature
        let has_temperature = reader.read_u8()? != 0;
        if has_temperature {
            let temperature = reader.read_f32::<LittleEndian>()?;
            frame.temperature = Some(temperature);
        }

        // Pressure
        let has_pressure = reader.read_u8()? != 0;
        if has_pressure {
            let pressure = reader.read_f32::<LittleEndian>()?;
            frame.pressure = Some(pressure);
        }

        // Optional full box vectors (added in a later TACO version). Old
        // frames stop here; the trailing byte may not exist, so EOF on the
        // initial flag read is interpreted as "no box_vectors".
        match reader.read_u8() {
            Ok(1) => {
                let mut vectors = [[0.0; 3]; 3];
                for row in vectors.iter_mut() {
                    for value in row.iter_mut() {
                        *value = reader.read_f32::<LittleEndian>()?;
                    }
                }
                frame.box_vectors = Some(vectors);
            }
            Ok(_) => {}
            Err(ref e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {}
            Err(e) => return Err(e.into()),
        }

        Ok(frame)
    }
}

/// Coordinate data (positions, velocities, forces) for a frame
#[derive(Debug, Clone, PartialEq)]
pub struct FrameData {
    /// Positions in nanometers. Shape: (n_atoms, 3)
    pub positions: Option<Array2<f32>>,

    /// Velocities in nm/ps. Shape: (n_atoms, 3)
    pub velocities: Option<Array2<f32>>,

    /// Forces in kJ/(mol*nm). Shape: (n_atoms, 3)
    pub forces: Option<Array2<f32>>,

    /// Stress tensor per atom or per frame (optional, shape: (n_atoms, 6) or (1, 6))
    pub stress: Option<Array2<f32>>,

    /// Arbitrary extra data (keyed by name, supports Array1, Array2, Array3)
    pub extra: HashMap<String, ExtraArray>,
}

impl FrameData {
    /// Create new frame data with the given positions
    pub fn new(positions: Array2<f32>) -> Self {
        FrameData {
            positions: Some(positions),
            velocities: None,
            forces: None,
            stress: None,
            extra: HashMap::new(),
        }
    }

    /// Create an empty frame data (used for reference/first frames)
    pub fn empty(n_atoms: usize) -> Self {
        FrameData {
            positions: Some(Array2::<f32>::zeros((n_atoms, 3))),
            velocities: None,
            forces: None,
            stress: None,
            extra: HashMap::new(),
        }
    }

    /// Add velocities to the frame data
    pub fn with_velocities(mut self, velocities: Array2<f32>) -> Self {
        self.velocities = Some(velocities);
        self
    }

    /// Add forces to the frame data
    pub fn with_forces(mut self, forces: Array2<f32>) -> Self {
        self.forces = Some(forces);
        self
    }

    pub fn with_stress(mut self, stress: Array2<f32>) -> Self {
        self.stress = Some(stress);
        self
    }

    pub fn set_stress(&mut self, stress: Array2<f32>) {
        self.stress = Some(stress);
    }

    pub fn with_extra_1d(mut self, key: &str, data: Array1<f32>) -> Self {
        self.extra.insert(key.to_string(), ExtraArray::Array1(data));
        self
    }

    pub fn with_extra_2d(mut self, key: &str, data: Array2<f32>) -> Self {
        self.extra.insert(key.to_string(), ExtraArray::Array2(data));
        self
    }

    pub fn with_extra_3d(mut self, key: &str, data: Array3<f32>) -> Self {
        self.extra.insert(key.to_string(), ExtraArray::Array3(data));
        self
    }

    pub fn with_extra_1d_f64(mut self, key: &str, data: Array1<f64>) -> Self {
        self.extra
            .insert(key.to_string(), ExtraArray::Array1F64(data));
        self
    }

    pub fn with_extra_2d_f64(mut self, key: &str, data: Array2<f64>) -> Self {
        self.extra
            .insert(key.to_string(), ExtraArray::Array2F64(data));
        self
    }

    pub fn with_extra_3d_f64(mut self, key: &str, data: Array3<f64>) -> Self {
        self.extra
            .insert(key.to_string(), ExtraArray::Array3F64(data));
        self
    }

    pub fn set_extra_1d(&mut self, key: &str, data: Array1<f32>) {
        self.extra.insert(key.to_string(), ExtraArray::Array1(data));
    }

    pub fn set_extra_2d(&mut self, key: &str, data: Array2<f32>) {
        self.extra.insert(key.to_string(), ExtraArray::Array2(data));
    }

    pub fn set_extra_3d(&mut self, key: &str, data: Array3<f32>) {
        self.extra.insert(key.to_string(), ExtraArray::Array3(data));
    }

    pub fn set_extra_1d_f64(&mut self, key: &str, data: Array1<f64>) {
        self.extra
            .insert(key.to_string(), ExtraArray::Array1F64(data));
    }

    pub fn set_extra_2d_f64(&mut self, key: &str, data: Array2<f64>) {
        self.extra
            .insert(key.to_string(), ExtraArray::Array2F64(data));
    }

    pub fn set_extra_3d_f64(&mut self, key: &str, data: Array3<f64>) {
        self.extra
            .insert(key.to_string(), ExtraArray::Array3F64(data));
    }

    pub fn get_extra(&self, key: &str) -> Option<&ExtraArray> {
        self.extra.get(key)
    }

    /// Calculate the delta (difference) from a reference frame
    pub fn delta_from(&self, reference: &FrameData) -> Result<FrameData> {
        let mut result = FrameData {
            positions: None,
            velocities: None,
            forces: None,
            stress: None,
            extra: HashMap::new(),
        };

        // Calculate delta positions if both frames have positions
        if let (Some(pos), Some(ref_pos)) = (&self.positions, &reference.positions) {
            if pos.shape() != ref_pos.shape() {
                return Err(Error::ShapeMismatch(format!(
                    "Position shapes don't match: {:?} vs {:?}",
                    pos.shape(),
                    ref_pos.shape()
                )));
            }

            result.positions = Some(pos - ref_pos);
        }

        // Calculate delta velocities if both frames have velocities
        if let (Some(vel), Some(ref_vel)) = (&self.velocities, &reference.velocities) {
            if vel.shape() != ref_vel.shape() {
                return Err(Error::ShapeMismatch(format!(
                    "Velocity shapes don't match: {:?} vs {:?}",
                    vel.shape(),
                    ref_vel.shape()
                )));
            }

            result.velocities = Some(vel - ref_vel);
        } else if let Some(vel) = &self.velocities {
            // If reference has no velocities, but current frame does, store full velocities
            result.velocities = Some(vel.clone());
        }

        // Calculate delta forces if both frames have forces
        if let (Some(force), Some(ref_force)) = (&self.forces, &reference.forces) {
            if force.shape() != ref_force.shape() {
                return Err(Error::ShapeMismatch(format!(
                    "Force shapes don't match: {:?} vs {:?}",
                    force.shape(),
                    ref_force.shape()
                )));
            }

            result.forces = Some(force - ref_force);
        } else if let Some(force) = &self.forces {
            // If reference has no forces, but current frame does, store full forces
            result.forces = Some(force.clone());
        }

        // Calculate delta stress if both frames have stress
        if let (Some(stress), Some(ref_stress)) = (&self.stress, &reference.stress) {
            if stress.shape() != ref_stress.shape() {
                return Err(Error::ShapeMismatch(format!(
                    "Stress shapes don't match: {:?} vs {:?}",
                    stress.shape(),
                    ref_stress.shape()
                )));
            }

            result.stress = Some(stress - ref_stress);
        } else if let Some(stress) = &self.stress {
            // If reference has no stress, but current frame does, store full stress
            result.stress = Some(stress.clone());
        }

        // Calculate delta extra data
        for (key, current_array) in &self.extra {
            if let Some(ref_array) = reference.extra.get(key) {
                // Both frames have this extra data, calculate delta
                match (current_array, ref_array) {
                    (ExtraArray::Array1(arr), ExtraArray::Array1(ref_arr)) => {
                        if arr.shape() != ref_arr.shape() {
                            return Err(Error::ShapeMismatch(format!(
                                "Extra array '{}' shapes don't match: {:?} vs {:?}",
                                key,
                                arr.shape(),
                                ref_arr.shape()
                            )));
                        }
                        result
                            .extra
                            .insert(key.clone(), ExtraArray::Array1(arr - ref_arr));
                    }
                    (ExtraArray::Array2(arr), ExtraArray::Array2(ref_arr)) => {
                        if arr.shape() != ref_arr.shape() {
                            return Err(Error::ShapeMismatch(format!(
                                "Extra array '{}' shapes don't match: {:?} vs {:?}",
                                key,
                                arr.shape(),
                                ref_arr.shape()
                            )));
                        }
                        result
                            .extra
                            .insert(key.clone(), ExtraArray::Array2(arr - ref_arr));
                    }
                    (ExtraArray::Array3(arr), ExtraArray::Array3(ref_arr)) => {
                        if arr.shape() != ref_arr.shape() {
                            return Err(Error::ShapeMismatch(format!(
                                "Extra array '{}' shapes don't match: {:?} vs {:?}",
                                key,
                                arr.shape(),
                                ref_arr.shape()
                            )));
                        }
                        result
                            .extra
                            .insert(key.clone(), ExtraArray::Array3(arr - ref_arr));
                    }
                    (ExtraArray::Array1F64(arr), ExtraArray::Array1F64(ref_arr)) => {
                        if arr.shape() != ref_arr.shape() {
                            return Err(Error::ShapeMismatch(format!(
                                "Extra array '{}' shapes don't match: {:?} vs {:?}",
                                key,
                                arr.shape(),
                                ref_arr.shape()
                            )));
                        }
                        result
                            .extra
                            .insert(key.clone(), ExtraArray::Array1F64(arr - ref_arr));
                    }
                    (ExtraArray::Array2F64(arr), ExtraArray::Array2F64(ref_arr)) => {
                        if arr.shape() != ref_arr.shape() {
                            return Err(Error::ShapeMismatch(format!(
                                "Extra array '{}' shapes don't match: {:?} vs {:?}",
                                key,
                                arr.shape(),
                                ref_arr.shape()
                            )));
                        }
                        result
                            .extra
                            .insert(key.clone(), ExtraArray::Array2F64(arr - ref_arr));
                    }
                    (ExtraArray::Array3F64(arr), ExtraArray::Array3F64(ref_arr)) => {
                        if arr.shape() != ref_arr.shape() {
                            return Err(Error::ShapeMismatch(format!(
                                "Extra array '{}' shapes don't match: {:?} vs {:?}",
                                key,
                                arr.shape(),
                                ref_arr.shape()
                            )));
                        }
                        result
                            .extra
                            .insert(key.clone(), ExtraArray::Array3F64(arr - ref_arr));
                    }
                    _ => {
                        return Err(Error::ShapeMismatch(format!(
                            "Extra array '{}' mismatch: ndim current={}, reference={}, dtype current={:?}, reference={:?}",
                            key,
                            current_array.ndim(),
                            ref_array.ndim(),
                            current_array.dtype(),
                            ref_array.dtype()
                        )));
                    }
                }
            } else {
                // Reference has no such extra data, store full current data
                result.extra.insert(key.clone(), current_array.clone());
            }
        }

        Ok(result)
    }

    /// Apply a delta to this frame data
    pub fn apply_delta(&self, delta: &FrameData) -> Result<FrameData> {
        let mut result = FrameData {
            positions: None,
            velocities: None,
            forces: None,
            stress: None,
            extra: HashMap::new(),
        };

        // Apply delta positions if both have positions
        if let (Some(pos), Some(delta_pos)) = (&self.positions, &delta.positions) {
            if pos.shape() != delta_pos.shape() {
                return Err(Error::ShapeMismatch(format!(
                    "Position shapes don't match: {:?} vs {:?}",
                    pos.shape(),
                    delta_pos.shape()
                )));
            }

            // Element-wise addition to ensure precision
            let mut new_pos = pos.clone();
            for ((i, j), &delta_val) in delta_pos.indexed_iter() {
                new_pos[[i, j]] += delta_val;
            }
            result.positions = Some(new_pos);
        } else if let Some(delta_pos) = &delta.positions {
            // If we only have delta positions, copy them as is (this shouldn't normally happen)
            result.positions = Some(delta_pos.clone());
            // Warning: Applying delta positions without reference positions
        } else if let Some(pos) = &self.positions {
            // If we have positions but no delta, keep the original positions
            result.positions = Some(pos.clone());
        }

        // Apply delta velocities if both have velocities
        if let (Some(vel), Some(delta_vel)) = (&self.velocities, &delta.velocities) {
            if vel.shape() != delta_vel.shape() {
                return Err(Error::ShapeMismatch(format!(
                    "Velocity shapes don't match: {:?} vs {:?}",
                    vel.shape(),
                    delta_vel.shape()
                )));
            }

            // Element-wise addition to ensure precision
            let mut new_vel = vel.clone();
            for ((i, j), &delta_val) in delta_vel.indexed_iter() {
                new_vel[[i, j]] += delta_val;
            }
            result.velocities = Some(new_vel);
        } else if let Some(delta_vel) = &delta.velocities {
            // If we only have delta velocities, use them directly (new feature)
            result.velocities = Some(delta_vel.clone());
        } else if let Some(vel) = &self.velocities {
            // If we have velocities but no delta, keep the original velocities
            result.velocities = Some(vel.clone());
        }

        // Apply delta forces if both have forces
        if let (Some(force), Some(delta_force)) = (&self.forces, &delta.forces) {
            if force.shape() != delta_force.shape() {
                return Err(Error::ShapeMismatch(format!(
                    "Force shapes don't match: {:?} vs {:?}",
                    force.shape(),
                    delta_force.shape()
                )));
            }

            // Element-wise addition to ensure precision
            let mut new_force = force.clone();
            for ((i, j), &delta_val) in delta_force.indexed_iter() {
                new_force[[i, j]] += delta_val;
            }
            result.forces = Some(new_force);
        } else if let Some(delta_force) = &delta.forces {
            // If we only have delta forces, use them directly (new feature)
            result.forces = Some(delta_force.clone());
        } else if let Some(force) = &self.forces {
            // If we have forces but no delta, keep the original forces
            result.forces = Some(force.clone());
        }

        // Apply delta stress if both have stress
        if let (Some(stress), Some(delta_stress)) = (&self.stress, &delta.stress) {
            if stress.shape() != delta_stress.shape() {
                return Err(Error::ShapeMismatch(format!(
                    "Stress shapes don't match: {:?} vs {:?}",
                    stress.shape(),
                    delta_stress.shape()
                )));
            }

            // Element-wise addition to ensure precision
            let mut new_stress = stress.clone();
            for ((i, j), &delta_val) in delta_stress.indexed_iter() {
                new_stress[[i, j]] += delta_val;
            }
            result.stress = Some(new_stress);
        } else if let Some(delta_stress) = &delta.stress {
            // If we only have delta stress, use it directly (new feature)
            result.stress = Some(delta_stress.clone());
        } else if let Some(stress) = &self.stress {
            // If we have stress but no delta, keep the original stress
            result.stress = Some(stress.clone());
        }

        // Apply delta extra data
        // Only process extra data that has deltas in the delta frame
        for (key, delta_array) in &delta.extra {
            if let Some(current_array) = self.extra.get(key) {
                // Both have this extra data, apply delta
                match (current_array, delta_array) {
                    (ExtraArray::Array1(arr), ExtraArray::Array1(delta_arr)) => {
                        if arr.shape() != delta_arr.shape() {
                            return Err(Error::ShapeMismatch(format!(
                                "Extra array '{}' shapes don't match: {:?} vs {:?}",
                                key,
                                arr.shape(),
                                delta_arr.shape()
                            )));
                        }
                        // Element-wise addition
                        let mut new_arr = arr.clone();
                        for (i, &delta_val) in delta_arr.indexed_iter() {
                            new_arr[i] += delta_val;
                        }
                        result
                            .extra
                            .insert(key.clone(), ExtraArray::Array1(new_arr));
                    }
                    (ExtraArray::Array2(arr), ExtraArray::Array2(delta_arr)) => {
                        if arr.shape() != delta_arr.shape() {
                            return Err(Error::ShapeMismatch(format!(
                                "Extra array '{}' shapes don't match: {:?} vs {:?}",
                                key,
                                arr.shape(),
                                delta_arr.shape()
                            )));
                        }
                        // Element-wise addition
                        let mut new_arr = arr.clone();
                        for ((i, j), &delta_val) in delta_arr.indexed_iter() {
                            new_arr[[i, j]] += delta_val;
                        }
                        result
                            .extra
                            .insert(key.clone(), ExtraArray::Array2(new_arr));
                    }
                    (ExtraArray::Array3(arr), ExtraArray::Array3(delta_arr)) => {
                        if arr.shape() != delta_arr.shape() {
                            return Err(Error::ShapeMismatch(format!(
                                "Extra array '{}' shapes don't match: {:?} vs {:?}",
                                key,
                                arr.shape(),
                                delta_arr.shape()
                            )));
                        }
                        // Element-wise addition
                        let mut new_arr = arr.clone();
                        for ((i, j, k), &delta_val) in delta_arr.indexed_iter() {
                            new_arr[[i, j, k]] += delta_val;
                        }
                        result
                            .extra
                            .insert(key.clone(), ExtraArray::Array3(new_arr));
                    }
                    (ExtraArray::Array1F64(arr), ExtraArray::Array1F64(delta_arr)) => {
                        if arr.shape() != delta_arr.shape() {
                            return Err(Error::ShapeMismatch(format!(
                                "Extra array '{}' shapes don't match: {:?} vs {:?}",
                                key,
                                arr.shape(),
                                delta_arr.shape()
                            )));
                        }
                        let mut new_arr = arr.clone();
                        for (i, &delta_val) in delta_arr.indexed_iter() {
                            new_arr[i] += delta_val;
                        }
                        result
                            .extra
                            .insert(key.clone(), ExtraArray::Array1F64(new_arr));
                    }
                    (ExtraArray::Array2F64(arr), ExtraArray::Array2F64(delta_arr)) => {
                        if arr.shape() != delta_arr.shape() {
                            return Err(Error::ShapeMismatch(format!(
                                "Extra array '{}' shapes don't match: {:?} vs {:?}",
                                key,
                                arr.shape(),
                                delta_arr.shape()
                            )));
                        }
                        let mut new_arr = arr.clone();
                        for ((i, j), &delta_val) in delta_arr.indexed_iter() {
                            new_arr[[i, j]] += delta_val;
                        }
                        result
                            .extra
                            .insert(key.clone(), ExtraArray::Array2F64(new_arr));
                    }
                    (ExtraArray::Array3F64(arr), ExtraArray::Array3F64(delta_arr)) => {
                        if arr.shape() != delta_arr.shape() {
                            return Err(Error::ShapeMismatch(format!(
                                "Extra array '{}' shapes don't match: {:?} vs {:?}",
                                key,
                                arr.shape(),
                                delta_arr.shape()
                            )));
                        }
                        let mut new_arr = arr.clone();
                        for ((i, j, k), &delta_val) in delta_arr.indexed_iter() {
                            new_arr[[i, j, k]] += delta_val;
                        }
                        result
                            .extra
                            .insert(key.clone(), ExtraArray::Array3F64(new_arr));
                    }
                    _ => {
                        return Err(Error::ShapeMismatch(format!(
                            "Extra array '{}' mismatch: ndim current={}, delta={}, dtype current={:?}, delta={:?}",
                            key,
                            current_array.ndim(),
                            delta_array.ndim(),
                            current_array.dtype(),
                            delta_array.dtype()
                        )));
                    }
                }
            } else {
                // Only delta has this extra data, use it directly (new feature)
                result.extra.insert(key.clone(), delta_array.clone());
            }
        }

        Ok(result)
    }

    /// Write the frame data to a writer with the given compression settings
    pub fn write<W: Write>(
        &self,
        writer: &mut W,
        settings: &CompressionSettings,
        extra_specs: &HashMap<String, ExtraDatasetSpec>,
    ) -> Result<()> {
        // Write positions if present
        if let Some(ref positions) = self.positions {
            writer.write_u8(1)?; // Positions present

            // Compress positions
            let compressed = compress_tensor(positions, settings.precision, settings.zstd_level)?;

            // Write compressed data length and data
            writer.write_u32::<LittleEndian>(compressed.len() as u32)?;
            writer.write_all(&compressed)?;
        } else {
            writer.write_u8(0)?; // No positions
        }

        // Write velocities if present
        if let Some(ref velocities) = self.velocities {
            writer.write_u8(1)?; // Velocities present

            // Compress velocities
            let compressed = compress_tensor(velocities, settings.precision, settings.zstd_level)?;

            // Write compressed data length and data
            writer.write_u32::<LittleEndian>(compressed.len() as u32)?;
            writer.write_all(&compressed)?;
        } else {
            writer.write_u8(0)?; // No velocities
        }

        // Write forces if present
        if let Some(ref forces) = self.forces {
            writer.write_u8(1)?; // Forces present

            // Compress forces
            let compressed = compress_tensor(forces, settings.precision, settings.zstd_level)?;

            // Write compressed data length and data
            writer.write_u32::<LittleEndian>(compressed.len() as u32)?;
            writer.write_all(&compressed)?;
        } else {
            writer.write_u8(0)?; // No forces
        }

        // Write stress if present
        if let Some(ref stress) = self.stress {
            writer.write_u8(1)?;
            let compressed = compress_tensor(stress, settings.precision, settings.zstd_level)?;
            writer.write_u32::<LittleEndian>(compressed.len() as u32)?;
            writer.write_all(&compressed)?;
        } else {
            writer.write_u8(0)?;
        }

        // Write extra data
        writer.write_u32::<LittleEndian>(self.extra.len() as u32)?;
        for (key, array) in &self.extra {
            let ndim = array.ndim();
            let shape = array.shape();

            let compressed = compress_extra_array(key, array, ndim, &shape, settings, extra_specs)?;

            // Write key length and key bytes
            let key_bytes = key.as_bytes();
            writer.write_u16::<LittleEndian>(key_bytes.len() as u16)?;
            writer.write_all(key_bytes)?;

            // Write ndim and shape
            writer.write_u8(ndim as u8)?;
            match ndim {
                1 => {
                    writer.write_u32::<LittleEndian>(shape[0])?;
                }
                2 => {
                    writer.write_u32::<LittleEndian>(shape[0])?;
                    writer.write_u32::<LittleEndian>(shape[1])?;
                }
                3 => {
                    writer.write_u32::<LittleEndian>(shape[0])?;
                    writer.write_u32::<LittleEndian>(shape[1])?;
                    writer.write_u32::<LittleEndian>(shape[2])?;
                }
                _ => {
                    return Err(Error::UnsupportedFeature(format!(
                        "Unsupported extra array ndim: {ndim}"
                    )));
                }
            }

            writer.write_u32::<LittleEndian>(compressed.len() as u32)?;
            writer.write_all(&compressed)?;
        }

        Ok(())
    }

    /// Read frame data from a reader with the given compression settings
    pub fn read<R: Read>(
        reader: &mut R,
        settings: &CompressionSettings,
        extra_specs: &HashMap<String, ExtraDatasetSpec>,
    ) -> Result<Self> {
        let mut data = FrameData {
            positions: None,
            velocities: None,
            forces: None,
            stress: None,
            extra: HashMap::new(),
        };

        // Read positions if present
        let has_positions = reader.read_u8()? != 0;
        if has_positions {
            // Read compressed data length
            let data_length = reader.read_u32::<LittleEndian>()? as usize;

            // Read compressed data
            let mut compressed = vec![0u8; data_length];
            reader.read_exact(&mut compressed)?;

            // Decompress positions
            let positions = decompress_tensor(&compressed, settings.precision)?;

            data.positions = Some(positions);
        }

        // Read velocities if present
        let has_velocities = reader.read_u8()? != 0;
        if has_velocities {
            // Read compressed data length
            let data_length = reader.read_u32::<LittleEndian>()? as usize;

            // Read compressed data
            let mut compressed = vec![0u8; data_length];
            reader.read_exact(&mut compressed)?;

            // Decompress velocities
            let velocities = decompress_tensor(&compressed, settings.precision)?;

            data.velocities = Some(velocities);
        }

        // Read forces if present
        let has_forces = reader.read_u8()? != 0;
        if has_forces {
            // Read compressed data length
            let data_length = reader.read_u32::<LittleEndian>()? as usize;

            // Read compressed data
            let mut compressed = vec![0u8; data_length];
            reader.read_exact(&mut compressed)?;

            // Decompress forces
            let forces = decompress_tensor(&compressed, settings.precision)?;

            data.forces = Some(forces);
        }

        // Read stress if present
        let has_stress = reader.read_u8()? != 0;
        if has_stress {
            let data_length = reader.read_u32::<LittleEndian>()? as usize;
            let mut compressed = vec![0u8; data_length];
            reader.read_exact(&mut compressed)?;
            let stress = decompress_tensor(&compressed, settings.precision)?;
            data.stress = Some(stress);
        }

        // Read extra data
        let extra_count = reader.read_u32::<LittleEndian>()?;
        for _ in 0..extra_count {
            let key_len = reader.read_u16::<LittleEndian>()? as usize;
            let mut key_bytes = vec![0u8; key_len];
            reader.read_exact(&mut key_bytes)?;
            let key =
                String::from_utf8(key_bytes).map_err(|e| Error::Deserialization(e.to_string()))?;
            let ndim = reader.read_u8()? as usize;

            let mut shape = Vec::with_capacity(ndim);
            match ndim {
                1 => {
                    let len = reader.read_u32::<LittleEndian>()?;
                    shape.push(len);
                }
                2 => {
                    let dim0 = reader.read_u32::<LittleEndian>()?;
                    let dim1 = reader.read_u32::<LittleEndian>()?;
                    shape.push(dim0);
                    shape.push(dim1);
                }
                3 => {
                    let dim0 = reader.read_u32::<LittleEndian>()?;
                    let dim1 = reader.read_u32::<LittleEndian>()?;
                    let dim2 = reader.read_u32::<LittleEndian>()?;
                    shape.push(dim0);
                    shape.push(dim1);
                    shape.push(dim2);
                }
                _ => {
                    return Err(Error::Deserialization(
                        "Unsupported extra array ndim".to_string(),
                    ));
                }
            }

            let data_length = reader.read_u32::<LittleEndian>()? as usize;
            let mut compressed = vec![0u8; data_length];
            reader.read_exact(&mut compressed)?;

            let array = decode_extra_array(&key, ndim, &shape, &compressed, settings, extra_specs)?;
            data.extra.insert(key, array);
        }

        Ok(data)
    }

    pub fn set_positions(&mut self, positions: Array2<f32>) {
        self.positions = Some(positions);
    }
}

fn validate_extra_spec(
    key: &str,
    ndim: usize,
    shape: &[u32],
    dtype: ExtraDataType,
    spec: Option<&ExtraDatasetSpec>,
) -> Result<ExtraCodec> {
    if shape.len() != ndim {
        return Err(Error::ShapeMismatch(format!(
            "Extra array '{}' shape length {} does not match ndim {}",
            key,
            shape.len(),
            ndim
        )));
    }
    if ndim == 0 || ndim > 3 {
        return Err(Error::UnsupportedFeature(format!(
            "Unsupported extra array ndim: {ndim}"
        )));
    }
    if dtype == ExtraDataType::F64 && spec.is_none() {
        return Err(Error::InvalidMetadata(format!(
            "Extra array '{}' uses f64 but no dataset spec is present in the header",
            key
        )));
    }

    if let Some(spec) = spec {
        if spec.dtype != dtype {
            return Err(Error::InvalidMetadata(format!(
                "Extra array '{}' dtype mismatch: header={:?}, data={:?}",
                key, spec.dtype, dtype
            )));
        }
        if let Some(spec_ndim) = spec.ndim {
            if spec_ndim as usize != ndim {
                return Err(Error::ShapeMismatch(format!(
                    "Extra array '{}' ndim mismatch: header={}, data={}",
                    key, spec_ndim, ndim
                )));
            }
        }
        if let Some(spec_shape) = &spec.shape {
            if spec_shape.as_slice() != shape {
                return Err(Error::ShapeMismatch(format!(
                    "Extra array '{}' shape mismatch: header={:?}, data={:?}",
                    key, spec_shape, shape
                )));
            }
        }
        Ok(spec.codec.clone())
    } else {
        Ok(ExtraCodec::default())
    }
}

fn shape_element_count(shape: &[u32]) -> Result<usize> {
    let mut total: u64 = 1;
    for &dim in shape {
        if dim == 0 {
            return Err(Error::ShapeMismatch(
                "Extra array shape contains a zero dimension".to_string(),
            ));
        }
        total = total
            .checked_mul(dim as u64)
            .ok_or_else(|| Error::ShapeMismatch("Extra array size overflow".to_string()))?;
    }
    Ok(total as usize)
}

fn fpzip_dims_from_shape(shape: &[u32]) -> Result<(u32, u32, u32, u32)> {
    for &dim in shape {
        if dim == 0 {
            return Err(Error::ShapeMismatch(
                "Extra array shape contains a zero dimension".to_string(),
            ));
        }
    }
    match shape.len() {
        1 => Ok((shape[0], 1, 1, 1)),
        2 => Ok((shape[0], shape[1], 1, 1)),
        3 => Ok((shape[0], shape[1], shape[2], 1)),
        _ => Err(Error::UnsupportedFeature(format!(
            "Unsupported fpzip ndim: {}",
            shape.len()
        ))),
    }
}

fn resolve_fpzip_prec(prec: Option<u8>, dtype: ExtraDataType, key: &str) -> Result<u32> {
    let bits = prec.unwrap_or(0) as u32;
    if bits == 0 {
        return Ok(bits);
    }
    match dtype {
        ExtraDataType::F32 => {
            if !(2..=32).contains(&bits) {
                return Err(Error::Compression(format!(
                    "fpzip prec out of range for '{}' (f32): {}",
                    key, bits
                )));
            }
        }
        ExtraDataType::F64 => {
            if !(4..=64).contains(&bits) || bits % 2 != 0 {
                return Err(Error::Compression(format!(
                    "fpzip prec out of range for '{}' (f64): {}",
                    key, bits
                )));
            }
        }
    }
    Ok(bits)
}

fn compress_extra_array(
    key: &str,
    array: &ExtraArray,
    ndim: usize,
    shape: &[u32],
    settings: &CompressionSettings,
    extra_specs: &HashMap<String, ExtraDatasetSpec>,
) -> Result<Vec<u8>> {
    let dtype = array.dtype();
    let spec = extra_specs.get(key);
    let codec = validate_extra_spec(key, ndim, shape, dtype, spec)?;

    match codec {
        ExtraCodec::Zstd { precision } => {
            let precision = precision.unwrap_or(settings.precision);
            match array {
                ExtraArray::Array1(arr) => {
                    let flat: Vec<f32> = arr.iter().copied().collect();
                    let arr2 = Array2::from_shape_vec((arr.len(), 1), flat)
                        .map_err(|e| Error::Deserialization(e.to_string()))?;
                    compress_tensor(&arr2, precision, settings.zstd_level)
                }
                ExtraArray::Array2(arr) => compress_tensor(arr, precision, settings.zstd_level),
                ExtraArray::Array3(arr) => {
                    let shape = arr.shape();
                    let flat: Vec<f32> = arr.iter().copied().collect();
                    let arr2 = Array2::from_shape_vec((shape[0], shape[1] * shape[2]), flat)
                        .map_err(|e| Error::Deserialization(e.to_string()))?;
                    compress_tensor(&arr2, precision, settings.zstd_level)
                }
                ExtraArray::Array1F64(arr) => {
                    if precision == PrecisionMode::Half {
                        return Err(Error::UnsupportedFeature(format!(
                            "Half precision zstd not supported for f64 extra array '{key}'"
                        )));
                    }
                    let flat: Vec<f64> = arr.iter().copied().collect();
                    let arr2 = Array2::from_shape_vec((arr.len(), 1), flat)
                        .map_err(|e| Error::Deserialization(e.to_string()))?;
                    compress_tensor_f64(&arr2, settings.zstd_level)
                }
                ExtraArray::Array2F64(arr) => {
                    if precision == PrecisionMode::Half {
                        return Err(Error::UnsupportedFeature(format!(
                            "Half precision zstd not supported for f64 extra array '{key}'"
                        )));
                    }
                    compress_tensor_f64(arr, settings.zstd_level)
                }
                ExtraArray::Array3F64(arr) => {
                    if precision == PrecisionMode::Half {
                        return Err(Error::UnsupportedFeature(format!(
                            "Half precision zstd not supported for f64 extra array '{key}'"
                        )));
                    }
                    let shape = arr.shape();
                    let flat: Vec<f64> = arr.iter().copied().collect();
                    let arr2 = Array2::from_shape_vec((shape[0], shape[1] * shape[2]), flat)
                        .map_err(|e| Error::Deserialization(e.to_string()))?;
                    compress_tensor_f64(&arr2, settings.zstd_level)
                }
            }
        }
        ExtraCodec::Fpzip { prec } => {
            let prec_bits = resolve_fpzip_prec(prec, dtype, key)?;
            let (nx, ny, nz, nf) = fpzip_dims_from_shape(shape)?;
            match array {
                ExtraArray::Array1(arr) => {
                    let flat: Vec<f32> = arr.iter().copied().collect();
                    FpZipCompressor::new(nx)
                        .ny(ny)
                        .nz(nz)
                        .nf(nf)
                        .prec(prec_bits)
                        .compress_f32(&flat)
                        .map_err(|e| Error::Compression(format!("fpzip: {e}")))
                }
                ExtraArray::Array2(arr) => {
                    let flat: Vec<f32> = arr.iter().copied().collect();
                    FpZipCompressor::new(nx)
                        .ny(ny)
                        .nz(nz)
                        .nf(nf)
                        .prec(prec_bits)
                        .compress_f32(&flat)
                        .map_err(|e| Error::Compression(format!("fpzip: {e}")))
                }
                ExtraArray::Array3(arr) => {
                    let flat: Vec<f32> = arr.iter().copied().collect();
                    FpZipCompressor::new(nx)
                        .ny(ny)
                        .nz(nz)
                        .nf(nf)
                        .prec(prec_bits)
                        .compress_f32(&flat)
                        .map_err(|e| Error::Compression(format!("fpzip: {e}")))
                }
                ExtraArray::Array1F64(arr) => {
                    let flat: Vec<f64> = arr.iter().copied().collect();
                    FpZipCompressor::new(nx)
                        .ny(ny)
                        .nz(nz)
                        .nf(nf)
                        .prec(prec_bits)
                        .compress_f64(&flat)
                        .map_err(|e| Error::Compression(format!("fpzip: {e}")))
                }
                ExtraArray::Array2F64(arr) => {
                    let flat: Vec<f64> = arr.iter().copied().collect();
                    FpZipCompressor::new(nx)
                        .ny(ny)
                        .nz(nz)
                        .nf(nf)
                        .prec(prec_bits)
                        .compress_f64(&flat)
                        .map_err(|e| Error::Compression(format!("fpzip: {e}")))
                }
                ExtraArray::Array3F64(arr) => {
                    let flat: Vec<f64> = arr.iter().copied().collect();
                    FpZipCompressor::new(nx)
                        .ny(ny)
                        .nz(nz)
                        .nf(nf)
                        .prec(prec_bits)
                        .compress_f64(&flat)
                        .map_err(|e| Error::Compression(format!("fpzip: {e}")))
                }
            }
        }
    }
}

fn decode_extra_array(
    key: &str,
    ndim: usize,
    shape: &[u32],
    compressed: &[u8],
    settings: &CompressionSettings,
    extra_specs: &HashMap<String, ExtraDatasetSpec>,
) -> Result<ExtraArray> {
    let spec = extra_specs.get(key);
    let dtype = spec.map(|s| s.dtype).unwrap_or(ExtraDataType::F32);
    let codec = validate_extra_spec(key, ndim, shape, dtype, spec)?;

    match (dtype, codec) {
        (ExtraDataType::F32, ExtraCodec::Zstd { precision }) => {
            let precision = precision.unwrap_or(settings.precision);
            let arr2 = decompress_tensor(compressed, precision)?;
            let expected = shape_element_count(shape)?;
            if arr2.len() != expected {
                return Err(Error::Decompression(format!(
                    "Extra array '{}' size mismatch after zstd decode",
                    key
                )));
            }
            match ndim {
                1 => {
                    let flat: Vec<f32> = arr2.iter().copied().collect();
                    Ok(ExtraArray::Array1(Array1::from(flat)))
                }
                2 => {
                    let dim0 = shape[0] as usize;
                    let dim1 = shape[1] as usize;
                    let arr2 = arr2
                        .to_shape((dim0, dim1))
                        .map_err(|e| Error::Deserialization(e.to_string()))?
                        .to_owned();
                    Ok(ExtraArray::Array2(arr2))
                }
                3 => {
                    let dim0 = shape[0] as usize;
                    let dim1 = shape[1] as usize;
                    let dim2 = shape[2] as usize;
                    let flat: Vec<f32> = arr2.iter().copied().collect();
                    let arr3 = Array3::from_shape_vec((dim0, dim1, dim2), flat)
                        .map_err(|e| Error::Deserialization(e.to_string()))?;
                    Ok(ExtraArray::Array3(arr3))
                }
                _ => Err(Error::UnsupportedFeature(format!(
                    "Unsupported extra array ndim: {ndim}"
                ))),
            }
        }
        (ExtraDataType::F32, ExtraCodec::Fpzip { prec: _ }) => {
            let flat = fpzip_decompress_f32(compressed)
                .map_err(|e| Error::Decompression(format!("fpzip: {e}")))?;
            let expected = shape_element_count(shape)?;
            if flat.len() != expected {
                return Err(Error::Decompression(format!(
                    "Extra array '{}' size mismatch after fpzip decode",
                    key
                )));
            }
            match ndim {
                1 => Ok(ExtraArray::Array1(Array1::from(flat))),
                2 => {
                    let dim0 = shape[0] as usize;
                    let dim1 = shape[1] as usize;
                    let arr2 = Array2::from_shape_vec((dim0, dim1), flat)
                        .map_err(|e| Error::Deserialization(e.to_string()))?;
                    Ok(ExtraArray::Array2(arr2))
                }
                3 => {
                    let dim0 = shape[0] as usize;
                    let dim1 = shape[1] as usize;
                    let dim2 = shape[2] as usize;
                    let arr3 = Array3::from_shape_vec((dim0, dim1, dim2), flat)
                        .map_err(|e| Error::Deserialization(e.to_string()))?;
                    Ok(ExtraArray::Array3(arr3))
                }
                _ => Err(Error::UnsupportedFeature(format!(
                    "Unsupported extra array ndim: {ndim}"
                ))),
            }
        }
        (ExtraDataType::F64, ExtraCodec::Zstd { precision }) => {
            let precision = precision.unwrap_or(settings.precision);
            if precision == PrecisionMode::Half {
                return Err(Error::UnsupportedFeature(format!(
                    "Half precision zstd not supported for f64 extra array '{key}'"
                )));
            }
            let arr2 = decompress_tensor_f64(compressed)?;
            let expected = shape_element_count(shape)?;
            if arr2.len() != expected {
                return Err(Error::Decompression(format!(
                    "Extra array '{}' size mismatch after zstd decode",
                    key
                )));
            }
            match ndim {
                1 => {
                    let flat: Vec<f64> = arr2.iter().copied().collect();
                    Ok(ExtraArray::Array1F64(Array1::from(flat)))
                }
                2 => {
                    let dim0 = shape[0] as usize;
                    let dim1 = shape[1] as usize;
                    let arr2 = arr2
                        .to_shape((dim0, dim1))
                        .map_err(|e| Error::Deserialization(e.to_string()))?
                        .to_owned();
                    Ok(ExtraArray::Array2F64(arr2))
                }
                3 => {
                    let dim0 = shape[0] as usize;
                    let dim1 = shape[1] as usize;
                    let dim2 = shape[2] as usize;
                    let flat: Vec<f64> = arr2.iter().copied().collect();
                    let arr3 = Array3::from_shape_vec((dim0, dim1, dim2), flat)
                        .map_err(|e| Error::Deserialization(e.to_string()))?;
                    Ok(ExtraArray::Array3F64(arr3))
                }
                _ => Err(Error::UnsupportedFeature(format!(
                    "Unsupported extra array ndim: {ndim}"
                ))),
            }
        }
        (ExtraDataType::F64, ExtraCodec::Fpzip { prec: _ }) => {
            let flat = fpzip_decompress_f64(compressed)
                .map_err(|e| Error::Decompression(format!("fpzip: {e}")))?;
            let expected = shape_element_count(shape)?;
            if flat.len() != expected {
                return Err(Error::Decompression(format!(
                    "Extra array '{}' size mismatch after fpzip decode",
                    key
                )));
            }
            match ndim {
                1 => Ok(ExtraArray::Array1F64(Array1::from(flat))),
                2 => {
                    let dim0 = shape[0] as usize;
                    let dim1 = shape[1] as usize;
                    let arr2 = Array2::from_shape_vec((dim0, dim1), flat)
                        .map_err(|e| Error::Deserialization(e.to_string()))?;
                    Ok(ExtraArray::Array2F64(arr2))
                }
                3 => {
                    let dim0 = shape[0] as usize;
                    let dim1 = shape[1] as usize;
                    let dim2 = shape[2] as usize;
                    let arr3 = Array3::from_shape_vec((dim0, dim1, dim2), flat)
                        .map_err(|e| Error::Deserialization(e.to_string()))?;
                    Ok(ExtraArray::Array3F64(arr3))
                }
                _ => Err(Error::UnsupportedFeature(format!(
                    "Unsupported extra array ndim: {ndim}"
                ))),
            }
        }
    }
}
