mzdata_bindata/

map.rs

use std::borrow::Cow;
use std::cmp::Ordering;
use std::collections::HashMap;
use std::collections::hash_map::{Iter, IterMut};
use std::convert::TryFrom;
use std::ffi::c_void;
use std::hash::BuildHasherDefault;

use identity_hash::IdentityHasher;
#[cfg(feature = "parallelism")]
use rayon::prelude::*;

use mzpeaks::Tolerance;

use mzdata_param::{Param, ParamDescribed, Unit};

use super::BinaryDataArrayType;
use super::array::DataArray;
use super::encodings::{ArrayRetrievalError, ArrayType, BinaryCompressionType};
use super::traits::{ByteArrayView, ByteArrayViewMut};

/// A collection of [`DataArray`]s that are identified by name.
///
/// # Examples
/// ```
/// use mzdata_bindata::{ArrayType, BinaryArrayMap, BinaryDataArrayType, DataArray};
///
/// let mut map = BinaryArrayMap::new();
/// assert!(!map.has_array(&ArrayType::MZArray));
///
/// let array = DataArray::from_name_and_type(&ArrayType::MZArray, BinaryDataArrayType::Float64);
/// map.add(array);
///
/// assert!(map.has_array(&ArrayType::MZArray));
/// assert_eq!(map.len(), 1);
/// ```
#[derive(Debug, Default, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct BinaryArrayMap {
    pub byte_buffer_map: HashMap<ArrayType, DataArray>,
}

impl BinaryArrayMap {
    pub fn new() -> BinaryArrayMap {
        BinaryArrayMap {
            ..Default::default()
        }
    }

    /// Get the number of arrays in the map
    pub fn len(&self) -> usize {
        self.byte_buffer_map.len()
    }

    pub fn is_empty(&self) -> bool {
        self.byte_buffer_map.is_empty()
    }

    /// Check if there is an ion mobility array present
    pub fn has_ion_mobility(&self) -> bool {
        self.byte_buffer_map.keys().any(|a| a.is_ion_mobility())
    }

    /// Iterate over references to the key-value pairs of this map
    pub fn iter(&self) -> Iter<'_, ArrayType, DataArray> {
        self.byte_buffer_map.iter()
    }

    #[cfg(feature = "parallelism")]
    pub fn par_iter(&self) -> rayon::collections::hash_map::Iter<'_, ArrayType, DataArray> {
        self.byte_buffer_map.par_iter()
    }

    /// Iterate over mutable references to the key-value pairs of this map
    pub fn iter_mut(&mut self) -> IterMut<'_, ArrayType, DataArray> {
        self.byte_buffer_map.iter_mut()
    }

    #[cfg(feature = "parallelism")]
    pub fn par_iter_mut(
        &mut self,
    ) -> rayon::collections::hash_map::IterMut<'_, ArrayType, DataArray> {
        self.byte_buffer_map.par_iter_mut()
    }

    /// Compress a specific [`DataArray`] with the `compression` scheme provided, if it is present.
    ///
    /// This method may fail if the compression fails or if the array type is missing.
    ///
    /// ## See Also
    /// [`DataArray::store_compressed`]
    pub fn encode_array(
        &mut self,
        array_type: &ArrayType,
        compression: BinaryCompressionType,
    ) -> Result<(), ArrayRetrievalError> {
        if let Some(arr) = self.get_mut(array_type) {
            arr.store_compressed(compression)?;
            Ok(())
        } else {
            Err(ArrayRetrievalError::NotFound(array_type.clone()))
        }
    }

    /// Decode all [`DataArray`] in this map if they have not been decoded already so
    /// that they are ready for use. If there are many arrays and the `parallelism` feature
    /// is enabled, arrays may be decoded in parallel.
    ///
    /// # Examples
    /// ```
    /// use mzdata_bindata::{ArrayType, BinaryArrayMap, BinaryCompressionType, BinaryDataArrayType, DataArray};
    ///
    /// let mut mz_array = DataArray::from_name_and_type(&ArrayType::MZArray, BinaryDataArrayType::Float64);
    /// mz_array.extend(&[1.0f64, 2.0, 3.0]).unwrap();
    /// mz_array.store_compressed(BinaryCompressionType::Zlib).unwrap();
    ///
    /// let mut map = BinaryArrayMap::new();
    /// map.add(mz_array);
    /// assert_eq!(map.get(&ArrayType::MZArray).unwrap().compression, BinaryCompressionType::Zlib);
    ///
    /// map.decode_all_arrays().unwrap();
    /// assert_eq!(map.get(&ArrayType::MZArray).unwrap().compression, BinaryCompressionType::Decoded);
    /// ```
    pub fn decode_all_arrays(&mut self) -> Result<(), ArrayRetrievalError> {
        #[cfg(not(feature = "parallelism"))]
        {
            self._decode_all_arrays()
        }
        #[cfg(feature = "parallelism")]
        {
            if self.len() > 2 {
                self._decode_all_arrays_parallel()
            } else {
                self._decode_all_arrays()
            }
        }
    }

    fn _decode_all_arrays(&mut self) -> Result<(), ArrayRetrievalError> {
        for (_key, value) in self.iter_mut() {
            match value.compression {
                BinaryCompressionType::Decoded => {}
                _ => {
                    value.decode_and_store()?;
                }
            }
        }
        Ok(())
    }

    #[cfg(feature = "parallelism")]
    fn _decode_all_arrays_parallel(&mut self) -> Result<(), ArrayRetrievalError> {
        let res: Result<(), ArrayRetrievalError> = self
            .iter_mut()
            .par_bridge()
            .map(|(_key, value)| {
                match value.compression {
                    BinaryCompressionType::Decoded => {}
                    _ => {
                        value.decode_and_store()?;
                    }
                }
                Ok(())
            })
            .collect::<Result<(), ArrayRetrievalError>>();
        res
    }

    /// Decode a specific [`DataArray`] if it is present.
    ///
    /// This method may fail if decoding fails or if the array type is missing.
    ///
    /// See [`DataArray::decode_and_store`]
    pub fn decode_array(&mut self, array_type: &ArrayType) -> Result<(), ArrayRetrievalError> {
        if let Some(array) = self.get_mut(array_type) {
            array.decode_and_store()?;
            Ok(())
        } else {
            Err(ArrayRetrievalError::NotFound(array_type.clone()))
        }
    }

    /// Add a [`DataArray`] to the map by its [`ArrayType`] name
    pub fn add(&mut self, array: DataArray) {
        self.byte_buffer_map.insert(array.name.clone(), array);
    }

    /// Get a reference to a specific [`DataArray`] if present
    pub fn get(&self, array_type: &ArrayType) -> Option<&DataArray> {
        self.byte_buffer_map.get(array_type)
    }

    /// Get a mutable reference to a specific [`DataArray`] if present
    pub fn get_mut(&mut self, array_type: &ArrayType) -> Option<&mut DataArray> {
        self.byte_buffer_map.get_mut(array_type)
    }

    /// Check whether a specific [`ArrayType`] is present
    pub fn has_array(&self, array_type: &ArrayType) -> bool {
        self.byte_buffer_map.contains_key(array_type)
    }

    /// Clear the map, discarding any array data
    pub fn clear(&mut self) {
        self.byte_buffer_map.clear();
    }

    /// Search for a specific m/z within a specified [`Tolerance`].
    ///
    /// This effectively calls [`BinaryArrayMap::mzs`] to get the m/z array and then
    /// performs a binary search to find the best index within range, if one exists.
    pub fn search(&self, query: f64, error_tolerance: Tolerance) -> Option<usize> {
        if let Ok(mzs) = self.mzs() {
            let (lower, _upper) = error_tolerance.bounds(query);
            match mzs[..].binary_search_by(|m| m.partial_cmp(&lower).unwrap()) {
                Ok(i) => {
                    let mut best_error = error_tolerance.call(query, mzs[i]).abs();
                    let mut best_index = i;
                    let mut index = i + 1;
                    while index < mzs.len() {
                        let error = error_tolerance.call(query, mzs[index]).abs();
                        if error < best_error {
                            best_index = index;
                            best_error = error;
                        }
                        index += 1;
                    }
                    if best_error < error_tolerance.tol() {
                        return Some(best_index);
                    }
                    None
                }
                Err(_err) => None,
            }
        } else {
            None
        }
    }

    /// Get a reference to the m/z array if it is present
    ///
    /// # Examples
    /// ```
    /// use mzdata_bindata::{ArrayType, BinaryArrayMap, BinaryDataArrayType, DataArray};
    ///
    /// let mut mz_array = DataArray::from_name_and_type(&ArrayType::MZArray, BinaryDataArrayType::Float64);
    /// mz_array.extend(&[100.0f64, 200.0, 300.0]).unwrap();
    ///
    /// let mut map = BinaryArrayMap::new();
    /// map.add(mz_array);
    ///
    /// assert_eq!(&*map.mzs().unwrap(), &[100.0f64, 200.0, 300.0][..]);
    /// ```
    pub fn mzs(&'_ self) -> Result<Cow<'_, [f64]>, ArrayRetrievalError> {
        let mz_array = self
            .get(&ArrayType::MZArray)
            .ok_or(ArrayRetrievalError::NotFound(ArrayType::MZArray))?
            .to_f64()
            .inspect_err(|e| log::error!("Failed to decode m/z array: {e}"))?;
        Ok(mz_array)
    }

    /// Re-order all arrays of matching length with respect to the provided index `mask` operating
    /// on the decoded native value states.
    ///
    /// The mask would be created using an indirect sort from another array:
    /// ```notest
    /// let vals = self.mzs().unwrap();
    /// let n = vals.len();
    /// let mut mask: Vec<usize> = (0..n).into_iter().collect();
    /// mask.sort_by(|i, j| {
    ///     let a = vals[*i];
    ///     let b = vals[*j];
    ///     a.total_cmp(&b)
    /// });
    /// drop(vals);
    /// self.sort_from_indices(mask)
    /// ```
    pub fn sort_from_indices(&mut self, mut mask: Vec<usize>) -> Result<(), ArrayRetrievalError> {
        let n = mask.len();
        const TOMBSTONE: usize = usize::MAX;
        let mut tracks: Vec<(*mut c_void, BinaryDataArrayType)> = Vec::new();
        for (_, v) in self.iter_mut() {
            if v.data_len()? != n {
                continue;
            }
            match v.dtype {
                BinaryDataArrayType::Float64 => {
                    let view = v.coerce_mut::<f64>()?;
                    tracks.push((view.as_mut_ptr() as *mut c_void, v.dtype))
                }
                BinaryDataArrayType::Float32 => {
                    let view = v.coerce_mut::<f32>()?;
                    tracks.push((view.as_mut_ptr() as *mut c_void, v.dtype))
                }
                BinaryDataArrayType::Int64 => {
                    let view = v.coerce_mut::<i64>()?;
                    tracks.push((view.as_mut_ptr() as *mut c_void, v.dtype))
                }
                BinaryDataArrayType::Int32 => {
                    let view = v.coerce_mut::<i32>()?;
                    tracks.push((view.as_mut_ptr() as *mut c_void, v.dtype))
                }
                BinaryDataArrayType::ASCII => todo!(),
                BinaryDataArrayType::Unknown => todo!(),
            }
        }

        for idx in 0..n {
            if mask[idx] != TOMBSTONE {
                let mut current_idx = idx;
                loop {
                    let next_idx = mask[current_idx];
                    mask[current_idx] = TOMBSTONE;
                    if mask[next_idx] == TOMBSTONE {
                        break;
                    }
                    for (dat, dtype) in tracks.iter_mut () {
                        match dtype {
                            BinaryDataArrayType::Float64 => {
                                unsafe {
                                    let v = core::slice::from_raw_parts_mut(*dat as *mut f64, n);
                                    v.swap(current_idx, next_idx);
                                }
                            },
                            BinaryDataArrayType::Float32 => {
                                unsafe {
                                    let v = core::slice::from_raw_parts_mut(*dat as *mut f32, n);
                                    v.swap(current_idx, next_idx);
                                }
                            },
                            BinaryDataArrayType::Int64 => {
                                unsafe {
                                    let v = core::slice::from_raw_parts_mut(*dat as *mut i64, n);
                                    v.swap(current_idx, next_idx);
                                }
                            },
                            BinaryDataArrayType::Int32 => {
                                unsafe {
                                    let v = core::slice::from_raw_parts_mut(*dat as *mut i32, n);
                                    v.swap(current_idx, next_idx);
                                }
                            },
                            BinaryDataArrayType::Unknown => todo!(),
                            BinaryDataArrayType::ASCII => todo!(),
                        }
                    }
                    current_idx = next_idx;
                }
            }
        }
        /*
        for idx in 0..n {
            if mask[idx] != TOMBSTONE {
                let mut current_idx = idx;
                loop {
                    let next_idx = mask[current_idx];
                    mask[current_idx] = TOMBSTONE;
                    if mask[next_idx] == TOMBSTONE {
                        break;
                    }
                    for (_, v) in self.iter_mut() {
                        if v.data_len()? != n {
                            continue;
                        }
                        match v.dtype {
                            BinaryDataArrayType::Float64 => {
                                let view = v.coerce_mut::<f64>()?;
                                view.swap(current_idx, next_idx);
                            }
                            BinaryDataArrayType::Float32 => {
                                let view = v.coerce_mut::<f32>()?;
                                view.swap(current_idx, next_idx);
                            }
                            BinaryDataArrayType::Int64 => {
                                let view = v.coerce_mut::<i64>()?;
                                view.swap(current_idx, next_idx);
                            }
                            BinaryDataArrayType::Int32 => {
                                let view = v.coerce_mut::<i32>()?;
                                view.swap(current_idx, next_idx);
                            }
                            BinaryDataArrayType::ASCII => todo!(),
                            BinaryDataArrayType::Unknown => todo!(),
                        }
                    }
                    current_idx = next_idx;
                }
            }
        }
        */
        Ok(())
    }

    /// Sort all the arrays in the map by `name` if they are the same length. This operates
    /// on the decoded state directly.
    ///
    /// This builds the sorting mask and then leaves [`Self::sort_from_indices`] to perform
    /// the re-ordering.
    ///
    /// Not supported for [`BinaryDataArrayType::ASCII`] and [`BinaryDataArrayType::Unknown`].
    pub fn sort_by_array(&mut self, name: &ArrayType) -> Result<(), ArrayRetrievalError> {
        let query_axis = self
            .get(name)
            .ok_or_else(|| ArrayRetrievalError::NotFound(name.clone()))?;
        macro_rules! sort_mask {
            ($conv:expr, $cmp:expr) => {{
                let vals = $conv?;
                if vals.is_sorted() {
                    return Ok(());
                }
                let n = vals.len();
                let mut mask: Vec<usize> = (0..n).into_iter().collect();
                mask.sort_by(|i, j| {
                    let a = vals[*i];
                    let b = vals[*j];
                    $cmp(&a, &b)
                });
                (mask, n)
            }};
        }
        let (mask, _) = match query_axis.dtype() {
            BinaryDataArrayType::Float64 => {
                sort_mask!(query_axis.to_f64(), f64::total_cmp)
            }
            BinaryDataArrayType::Float32 => {
                sort_mask!(query_axis.to_f32(), f32::total_cmp)
            }
            BinaryDataArrayType::Int64 => {
                sort_mask!(query_axis.to_i64(), i64::cmp)
            }
            BinaryDataArrayType::Int32 => {
                sort_mask!(query_axis.to_i32(), i32::cmp)
            }
            BinaryDataArrayType::ASCII => todo!(),
            BinaryDataArrayType::Unknown => todo!(),
        };

        self.sort_from_indices(mask)
    }

    /// Get a mutable reference to the m/z array if it is present.
    ///
    /// This method obtains the m/z array, calls [`DataArray::decode_and_store`] to force
    /// the data to be stored natively in memory, and forces the storage to be
    /// [`BinaryDataArrayType::Float64`] if not already. Finally it returns a mutable
    /// view of that decoded dataas `&mut [f64]`
    pub fn mzs_mut(&mut self) -> Result<&mut [f64], ArrayRetrievalError> {
        if let Some(mz_array) = self.get_mut(&ArrayType::MZArray) {
            mz_array
                .decode_and_store()
                .inspect_err(|e| log::error!("Failed to decode m/z array: {e}"))?;
            mz_array.store_as(BinaryDataArrayType::Float64)?;
            mz_array.coerce_mut()
        } else {
            Err(ArrayRetrievalError::NotFound(ArrayType::MZArray))
        }
    }

    /// Get a reference to the intensity array if it is present
    pub fn intensities(&'_ self) -> Result<Cow<'_, [f32]>, ArrayRetrievalError> {
        let intensities = self
            .get(&ArrayType::IntensityArray)
            .ok_or(ArrayRetrievalError::NotFound(ArrayType::IntensityArray))?
            .to_f32()
            .inspect_err(|e| log::error!("Failed to decode intensity array: {e:?}"))?;
        Ok(intensities)
    }

    /// Get a mutable reference to the intensity array if it is present.
    ///
    /// See [`BinaryArrayMap::mzs_mut`] for an explanation of the intermediate steps,
    /// save that the forced storage type is [`BinaryArrayDataType::Float32`].
    pub fn intensities_mut(&mut self) -> Result<&mut [f32], ArrayRetrievalError> {
        if let Some(arr) = self.get_mut(&ArrayType::IntensityArray) {
            arr
                .decode_and_store()
                .inspect_err(|e| log::error!("Failed to decode intensity array: {e}"))?;
            arr
                .store_as(BinaryDataArrayType::Float32)
                .inspect_err(|e| log::error!("Failed to decode intensity array: {e}"))?;
            arr.coerce_mut()
        } else {
            Err(ArrayRetrievalError::NotFound(ArrayType::IntensityArray))
        }
    }

    /// Get a reference to the charge array if it is present
    pub fn charges(&'_ self) -> Result<Cow<'_, [i32]>, ArrayRetrievalError> {
        match self.get(&ArrayType::ChargeArray) {
            Some(data_array) => data_array.to_i32(),
            None => Err(ArrayRetrievalError::NotFound(ArrayType::ChargeArray)),
        }
    }

    /// Get a mutable reference to the charge array if it is present.
    ///
    /// See [`BinaryArrayMap::mzs_mut`] for an explanation of the intermediate steps,
    /// save that the forced storage type is [`BinaryArrayDataType::Int32`].
    pub fn charge_mut(&mut self) -> Result<&mut [i32], ArrayRetrievalError> {
        if let Some(arr) = self.get_mut(&ArrayType::ChargeArray) {
            arr.decode_and_store()?;
            arr
                .store_as(BinaryDataArrayType::Int32)
                .inspect_err(|e| log::error!("Failed to decode charge array: {e}"))?;
            arr.coerce_mut()
        } else {
            Err(ArrayRetrievalError::NotFound(ArrayType::ChargeArray))
        }
    }

    /// Get a reference to the ion mobility array if it is present
    pub fn ion_mobility(&self) -> Result<(Cow<'_, [f64]>, ArrayType), ArrayRetrievalError> {
        if let Some((array_type, data_array)) = self
            .byte_buffer_map
            .iter()
            .find(|(a, _)| a.is_ion_mobility())
        {
            Ok((
                data_array
                    .to_f64()
                    .inspect_err(|e| log::error!("Failed to decode ion mobility array: {e}"))?,
                array_type.clone(),
            ))
        } else {
            Err(ArrayRetrievalError::NotFound(ArrayType::IonMobilityArray))
        }
    }

    /// Get a mutable reference to the ion mobility array if it is present.
    ///
    /// See [`BinaryArrayMap::mzs_mut`] for an explanation of the intermediate steps.
    pub fn ion_mobility_mut(&mut self) -> Result<(&mut [f64], ArrayType), ArrayRetrievalError> {
        if let Some((array_type, data_array)) = self
            .byte_buffer_map
            .iter_mut()
            .find(|(a, _)| a.is_ion_mobility())
        {
            data_array
                .decode_and_store()
                .inspect_err(|e| log::error!("Failed to decode ion mobility array: {e}"))?;
            data_array.store_as(BinaryDataArrayType::Float32)?;
            Ok((data_array.coerce_mut()?, array_type.clone()))
        } else {
            Err(ArrayRetrievalError::NotFound(ArrayType::IonMobilityArray))
        }
    }

    /// Split an array map along the ion mobility dimension if it is present.
    ///
    /// This function will fail if there is no ion mobility dimension.
    pub fn stack_ion_mobility(self) -> Result<BinaryArrayMap3D, ArrayRetrievalError> {
        BinaryArrayMap3D::try_from(self)
    }
}

impl IntoIterator for BinaryArrayMap {
    type Item = (ArrayType, DataArray);

    type IntoIter = <HashMap<ArrayType, DataArray> as IntoIterator>::IntoIter;

    fn into_iter(self) -> Self::IntoIter {
        self.byte_buffer_map.into_iter()
    }
}

#[cfg(feature = "mzsignal")]
mod mzsignal_impl {
    use super::*;

    use crate::to_bytes;
    use mzsignal::ArrayPair;

    impl From<ArrayPair<'_>> for BinaryArrayMap {
        fn from(value: ArrayPair<'_>) -> Self {
            let mz_array = DataArray::wrap(
                &ArrayType::MZArray,
                BinaryDataArrayType::Float64,
                to_bytes(&value.mz_array),
            );

            let intensity_array = DataArray::wrap(
                &ArrayType::IntensityArray,
                BinaryDataArrayType::Float32,
                to_bytes(&value.intensity_array),
            );

            let mut array_map = BinaryArrayMap::new();
            array_map.add(mz_array);
            array_map.add(intensity_array);
            array_map
        }
    }
}

#[derive(Debug, Default, Clone, Copy, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
struct NonNaNF64(f64);

impl std::hash::Hash for NonNaNF64 {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        state.write_u64(self.0.to_bits());
    }
}

impl identity_hash::IdentityHashable for NonNaNF64 {}

impl From<f64> for NonNaNF64 {
    fn from(value: f64) -> Self {
        Self::wrap(value)
            .unwrap_or_else(|| panic!("Expected an order-able f64 value, but found {}", value))
    }
}

impl NonNaNF64 {
    fn wrap(value: f64) -> Option<Self> {
        if value.is_nan() {
            None
        } else {
            Some(Self(value))
        }
    }
}

impl PartialOrd for NonNaNF64 {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Eq for NonNaNF64 {}

impl Ord for NonNaNF64 {
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.total_cmp(&other.0)
    }
}

macro_rules! _populate_stacked_array_from {
    ($im_dim:ident, $view:ident, $index_map:ident, $array_bins:ident, $array_type:ident, $array:ident) => {
        for (i_im, im) in $im_dim.iter() {
            let v = $view[*i_im];
            let i_axis = $index_map[&NonNaNF64(*im)];
            let bin = &mut $array_bins[i_axis];
            if let Some(bin_array) = bin.get_mut($array_type) {
                bin_array.push(v)?;
            } else {
                let mut bin_array = DataArray::from_name_and_type($array_type, $array.dtype());
                *bin_array.unit_mut() = $array.unit();
                bin_array.push(v)?;
                bin.add(bin_array);
            }
        }
    };
}

/// Represent a set of [`BinaryArrayMap`] that has been split across an
/// ion mobility dimension.
///
/// This effectively removes an ion mobility [`DataArray`] and for each unique
/// value in that array, creates a new [`BinaryArrayMap`] with subsets of each
/// other [`DataArray`] that co-occur at the same ion mobility coordinate.
///
/// Because [`BinaryArrayMap3D`] spreads [`DataArray`] over multiple ion mobility
/// bins, it cannot share a single [`ParamList`] for each of the dimensions. Instead,
/// all split [`ArrayType`] have their [`ParamList`] moved to separate storage accessed
/// using [`params_of`] and its siblins. These [`ParamList`] will be propagated by [`Self::stack`]
/// and [`BinaryArrayMap3D::unstack`] as appropriate.
#[derive(Debug, Default, Clone)]
#[cfg_attr(feature = "serde", serde_with::serde_as)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct BinaryArrayMap3D {
    pub ion_mobility_dimension: Vec<f64>,
    pub ion_mobility_type: ArrayType,
    pub ion_mobility_unit: Unit,
    pub arrays: Vec<BinaryArrayMap>,
    pub additional_arrays: BinaryArrayMap,
    #[cfg_attr(feature = "serde", serde_as(as = "Vec<(_, _)>"))]
    ion_mobility_index: HashMap<NonNaNF64, usize, BuildHasherDefault<IdentityHasher<NonNaNF64>>>,
    parameters_of: HashMap<ArrayType, mzdata_param::ParamList>
}

impl BinaryArrayMap3D {
    /// Create an empty [`BinaryArrayMap3D`] with a defined ion mobility dimension
    /// but no content at each ion mobility point.
    ///
    /// # Panics
    /// If the ion mobility dimension contains a NaN value
    pub fn from_ion_mobility_dimension(
        ion_mobility_dimension: Vec<f64>,
        ion_mobility_type: ArrayType,
        ion_mobility_unit: Unit,
    ) -> BinaryArrayMap3D {
        let ion_mobility_index: HashMap<_, _, BuildHasherDefault<IdentityHasher<NonNaNF64>>> = ion_mobility_dimension
            .iter()
            .copied()
            .enumerate()
            .map(|(i, f)| {
                (
                    NonNaNF64::wrap(f)
                        .unwrap_or_else(|| panic!("Expected non-NaN value for ion mobility")),
                    i,
                )
            })
            .collect();
        let mut arrays = Vec::new();
        arrays.resize_with(ion_mobility_dimension.len(), BinaryArrayMap::default);
        Self {
            ion_mobility_dimension,
            ion_mobility_type,
            ion_mobility_unit,
            arrays,
            ion_mobility_index,
            additional_arrays: Default::default(),
            parameters_of: Default::default(),
        }
    }

    /// Construct a [`BinaryArrayMap3D`] from a defined ion mobility dimension and
    /// a [`Vec`] of [`BinaryArrayMap`] at each point.
    ///
    /// # Panics
    /// If the ion mobility dimension contains a NaN value or `ion_mobility_dimension.len() != arrays.len()`
    pub fn from_ion_mobility_dimension_and_arrays(
        ion_mobility_dimension: Vec<f64>,
        ion_mobility_type: ArrayType,
        ion_mobility_unit: Unit,
        arrays: Vec<BinaryArrayMap>,
    ) -> BinaryArrayMap3D {
        let ion_mobility_index: HashMap<_, _, _> = ion_mobility_dimension
            .iter()
            .copied()
            .enumerate()
            .map(|(i, f)| {
                (
                    NonNaNF64::wrap(f)
                        .unwrap_or_else(|| panic!("Expected non-NaN value for ion mobility")),
                    i,
                )
            })
            .collect();

        assert_eq!(
            ion_mobility_dimension.len(),
            arrays.len(),
            "The number of ion mobility points does not match the number of point-level maps!"
        );

        Self {
            ion_mobility_dimension,
            ion_mobility_type,
            ion_mobility_unit,
            arrays,
            ion_mobility_index,
            additional_arrays: Default::default(),
            parameters_of: Default::default(),
        }
    }

    /// Get a reference to the [`ParamList`] for the specified [`ArrayType`], if one exists.
    ///
    /// # See also
    /// - [`Self::params_of_mut`] to mutate existing [`ParamList`] for extant [`ArrayType`]
    /// - [`Self::insert_params_of`] to create or overwrite the [`ParamList`] for an [`ArrayType`]
    pub fn params_of(&self, array_type: &ArrayType) -> Option<&Vec<Param>> {
        self.parameters_of.get(array_type)
    }

    /// Get a mutable reference to the [`ParamList`] for the specified [`ArrayType`], if one exists.
    ///
    /// # See also
    /// - [`Self::params_of`] to read existing [`ParamList`] for extant [`ArrayType`]
    /// - [`Self::insert_params_of`] to create or overwrite the [`ParamList`] for an [`ArrayType`]
    /// - [`Self::params_of_entry`] to create or modify the [`ParamList`] using a lower
    /// level Entry API
    pub fn params_of_mut(&mut self, array_type: &ArrayType) -> Option<&mut Vec<Param>> {
        self.parameters_of.get_mut(array_type)
    }

    /// A lower level create or update version of [`Self::insert_params_of`] that uses the
    /// underlying [`HashMap`]'s Entry API
    pub fn params_of_entry(&mut self, array_type: ArrayType) -> std::collections::hash_map::Entry<'_, ArrayType, Vec<Param>> {
        self.parameters_of.entry(array_type)
    }

    /// Set the [`ParamList`] for the specified [`ArrayType`].
    ///
    /// Returns any previous [`ParamList`] associated with that [`ArrayType`].
    ///
    /// # See also
    /// - [`Self::params_of`] to read existing [`ParamList`] for extant [`ArrayType`]
    /// - [`Self::params_of_mut`] to mutate existing [`ParamList`] for extant [`ArrayType`]
    /// - [`Self::params_of_entry`] to create or modify the [`ParamList`] using a lower
    /// level Entry API
    pub fn insert_params_of(&mut self, array_type: ArrayType, params: Vec<Param>) -> Option<Vec<Param>> {
        self.parameters_of.insert(array_type, params)
    }

    /// Get the associated arrays at the requested ion mobility, if they exist
    pub fn get_ion_mobility(&self, ion_mobility: f64) -> Option<&BinaryArrayMap> {
        if let Some(i) = NonNaNF64::wrap(ion_mobility) {
            if let Some(i) = self.ion_mobility_index.get(&i) {
                self.arrays.get(*i)
            } else {
                None
            }
        } else {
            None
        }
    }

    /// Perform a binary search on the ion mobility dimension and find the nearest
    /// point to `ion_mobility` within `error_tolerance` absolute distance.
    ///
    /// Returns the [`BinaryArrayMap`] at that point and the distance from `ion_mobility`
    /// matched.
    pub fn search_ion_mobility(
        &self,
        ion_mobility: f64,
        error_tolerance: f64,
    ) -> Option<(&BinaryArrayMap, f64)> {
        match self
            .ion_mobility_dimension
            .binary_search_by(|x: &f64| x.total_cmp(&ion_mobility))
        {
            Ok(i) => {
                let delta = ion_mobility - self.ion_mobility_dimension[i];
                if delta.abs() <= error_tolerance {
                    self.arrays.get(i).map(|a| (a, delta))
                } else {
                    None
                }
            }
            Err(i) => {
                if self.arrays.is_empty() {
                    return None;
                }
                let delta = ion_mobility - self.ion_mobility_dimension[i];
                if delta.abs() <= error_tolerance {
                    self.arrays.get(i).map(|a| (a, delta))
                } else {
                    None
                }
            }
        }
    }

    /// Perform a binary search on the ion mobility dimension and find the nearest
    /// point to `ion_mobility` within `error_tolerance` absolute distance.
    ///
    /// Unlike [`Self::search_ion_mobility`], this returns the index in the ion mobility
    /// dimension that matched and the ion mobility value at that index.
    pub fn search_ion_mobility_index(&self, ion_mobility: f64, error_tolerance: f64) -> Option<(usize, f64)> {
        match self
            .ion_mobility_dimension
            .binary_search_by(|x: &f64| x.total_cmp(&ion_mobility))
        {
            Ok(i) => {
                let delta = ion_mobility - self.ion_mobility_dimension[i];
                if delta.abs() <= error_tolerance {
                    Some((i, self.ion_mobility_dimension[i]))
                } else {
                    None
                }
            }
            Err(i) => {
                if self.arrays.is_empty() {
                    return None;
                }
                let delta = ion_mobility - self.ion_mobility_dimension[i];
                if delta.abs() <= error_tolerance {
                    Some((i, self.ion_mobility_dimension[i]))
                } else {
                    None
                }
            }
        }
    }

    /// Get the a mutable reference to the associated arrays at the requested ion mobility, if they exist.
    ///
    /// The ion mobility value must be an *exact* match
    pub fn get_ion_mobility_mut(&mut self, ion_mobility: f64) -> Option<&mut BinaryArrayMap> {
        if let Some(i) = NonNaNF64::wrap(ion_mobility) {
            if let Some(i) = self.ion_mobility_index.get(&i) {
                self.arrays.get_mut(*i)
            } else {
                None
            }
        } else {
            None
        }
    }

    /// Iterate over the ion mobility dimension and associated arrays
    /// at each point.
    pub fn iter(&self) -> impl Iterator<Item = (f64, &BinaryArrayMap)> {
        self.ion_mobility_dimension
            .iter()
            .copied()
            .zip(self.arrays.iter())
    }

    /// Iterate over the ion mobility dimension and a mutable reference to the associated arrays
    /// at each point.
    pub fn iter_mut(&mut self) -> impl Iterator<Item = (f64, &mut BinaryArrayMap)> {
        self.ion_mobility_dimension
            .iter()
            .copied()
            .zip(self.arrays.iter_mut())
    }

    /// Flatten this array into a single [`BinaryArrayMap`].
    ///
    /// The flattened map is sorted by m/z and ion mobility. If no m/z array is found,
    /// *no* sorting is done, and the flattened order is left in roughly ion mobility sorted
    /// order.
    ///
    /// # Errors
    /// [`ArrayRetrievalError`] errors related to array decoding occur if
    /// any [`DataArray`] cannot be decoded, or if an expected array is absent.
    pub fn unstack(&self) -> Result<BinaryArrayMap, ArrayRetrievalError> {
        let mut destination = self.additional_arrays.clone();

        let mut im_dim =
            DataArray::from_name_and_type(&self.ion_mobility_type, BinaryDataArrayType::Float64);
        im_dim.unit = self.ion_mobility_unit;
        let mut sizes = Vec::new();
        let mut mz_size = None;
        for (layer, im) in self.arrays.iter().zip(self.ion_mobility_dimension.iter()) {
            sizes.clear();
            for (key, array) in layer.iter() {
                match destination.get_mut(key) {
                    Some(sink) => {
                        sink.extend_raw(&array.data).inspect_err(|e| {
                            log::error!("Failed to extend {key:?}: {e}");
                        })?;
                    }
                    None => {
                        destination.add(array.clone());
                    }
                }
                if matches!(key, ArrayType::MZArray) {
                    mz_size = Some(array.data_len()?);
                } else {
                    sizes.push((key, array.data_len()?));
                }
            }
            if let Some(mz_size) = mz_size {
                im_dim.extend_iter(std::iter::repeat_n(*im, mz_size))?;
            } else if let Some((_, size)) = sizes.first() {
                im_dim.extend_iter(std::iter::repeat_n(*im, *size))?;
            }
        }

        let sorter = if let Some(mz_array) = destination.get(&ArrayType::MZArray) {
            let mzs = mz_array.to_f64()?;
            let ims = im_dim.to_f64()?;

            let mut indices: Vec<usize> = Vec::with_capacity(mzs.len());
            indices.extend(0..mzs.len());
            indices.sort_by(|i, j| {
                mzs[*i]
                    .total_cmp(&mzs[*j])
                    .then_with(|| ims[*i].total_cmp(&ims[*j]))
            });
            Some(indices)
        } else {
            None
        };

        let final_size = im_dim.data_len()?;

        destination.add(im_dim);

        if let Some(sorter) = sorter {
            destination.sort_from_indices(sorter)?;
        } else if final_size > 0 {
            log::debug!("Unsorted unstack");
        }

        for (k, v) in destination.iter_mut(){
            if let Some(pars) = self.parameters_of.get(k) {
                v.params_mut().extend(pars.iter().cloned());
            }
        }
        Ok(destination)
    }

    /// Convert a [`BinaryArrayMap`] into a [`BinaryArrayMap3D`] if it has an ion mobility dimension.
    ///
    /// Any arrays that aren't the same length as the ion mobility dimension will be in
    /// [`BinaryArrayMap3D::additional_arrays`].
    ///
    /// This uses *exact* ion mobility values as keys, bit-level equality. If ion mobility
    /// values are computed, they may vary at trivial bits and be equivalent, but not **equal**.
    ///
    /// # Errors
    /// [`ArrayRetrievalError`] errors related to array decoding occur if any [`DataArray`]
    /// cannot be decoded, or if an expected array is absent.
    pub fn stack(source: &BinaryArrayMap) -> Result<Self, ArrayRetrievalError> {
        let mut this = Self::default();
        if !source.has_ion_mobility() {
            return Err(ArrayRetrievalError::NotFound(ArrayType::IonMobilityArray));
        }
        let (im_dim, im_type) = source.ion_mobility()?;
        this.ion_mobility_unit = source.get(&im_type).unwrap().unit;
        this.ion_mobility_type = im_type;
        if im_dim.is_empty() {
            return Ok(this);
        }
        let mut im_dim: Vec<(usize, f64)> = im_dim.iter().copied().enumerate().collect();
        im_dim.sort_by(|(_, va), (_, vb)| va.total_cmp(vb));

        let mut im_axis = Vec::with_capacity(200);
        let mut last_v = im_dim.first().unwrap().1 - 1.0;
        let mut index_map = HashMap::default();
        for (_, v) in im_dim.iter() {
            if v.total_cmp(&last_v).is_gt() {
                last_v = *v;
                index_map.insert(NonNaNF64::from(*v), im_axis.len());
                im_axis.push(*v);
            }
        }

        let mut array_bins: Vec<BinaryArrayMap> = Vec::with_capacity(im_axis.len());
        array_bins.resize(im_axis.len(), BinaryArrayMap::default());
        for (array_type, array) in source.iter() {
            if array_type.is_ion_mobility() {
                continue;
            }
            if array.data_len()? != im_dim.len() {
                this.additional_arrays.add(array.clone());
                continue;
            }
            if !array.params().is_empty() {
                this.parameters_of.insert(array_type.clone(), array.params().to_vec());
            }
            match array.dtype() {
                BinaryDataArrayType::Unknown => {
                    panic!("Cannot re-sort opaque or unknown dimension data types")
                }
                BinaryDataArrayType::Float64 => {
                    let view = array.to_f64()?;
                    _populate_stacked_array_from!(
                        im_dim, view, index_map, array_bins, array_type, array
                    );
                }
                BinaryDataArrayType::Float32 => {
                    let view = array.to_f32()?;
                    _populate_stacked_array_from!(
                        im_dim, view, index_map, array_bins, array_type, array
                    );
                }
                BinaryDataArrayType::Int64 => {
                    let view = array.to_i64()?;
                    _populate_stacked_array_from!(
                        im_dim, view, index_map, array_bins, array_type, array
                    );
                }
                BinaryDataArrayType::Int32 => {
                    let view = array.to_i32()?;
                    _populate_stacked_array_from!(
                        im_dim, view, index_map, array_bins, array_type, array
                    );
                }
                BinaryDataArrayType::ASCII => {
                    let view = array.decode()?;
                    _populate_stacked_array_from!(
                        im_dim, view, index_map, array_bins, array_type, array
                    );
                }
            }
        }
        this.ion_mobility_dimension = im_axis;
        this.ion_mobility_index = index_map;
        this.arrays = array_bins;
        Ok(this)
    }
}

impl TryFrom<BinaryArrayMap> for BinaryArrayMap3D {
    type Error = ArrayRetrievalError;

    /// A wrapper around [`BinaryArrayMap3D::stack`]
    fn try_from(value: BinaryArrayMap) -> Result<Self, Self::Error> {
        Self::stack(&value)
    }
}

/// A [`BinaryArrayMap3D`] aliases [`TryFrom`] on [`BinaryArrayMap`] with [`BinaryArrayMap3D::stack`]
impl TryFrom<&BinaryArrayMap> for BinaryArrayMap3D {
    type Error = ArrayRetrievalError;

    fn try_from(value: &BinaryArrayMap) -> Result<Self, Self::Error> {
        Self::stack(value)
    }
}

impl TryFrom<BinaryArrayMap3D> for BinaryArrayMap {
    type Error = ArrayRetrievalError;

    fn try_from(value: BinaryArrayMap3D) -> Result<Self, Self::Error> {
        value.unstack()
    }
}

/// A [`BinaryArrayMap`] aliases [`TryFrom`] on [`BinaryArrayMap3D`] with [`BinaryArrayMap3D::unstack`]
impl TryFrom<&BinaryArrayMap3D> for BinaryArrayMap {
    type Error = ArrayRetrievalError;

    fn try_from(value: &BinaryArrayMap3D) -> Result<Self, Self::Error> {
        value.unstack()
    }
}



#[cfg(test)]
mod test {
    use crate::BinaryDataArrayType;

    use super::*;
    use std::fs;
    use std::io::{self, prelude::*};

    fn make_array_from_file() -> io::Result<DataArray> {
        let mut fh = fs::File::open("../../test/data/mz_f64_zlib_bas64.txt")?;
        let mut buf = String::new();
        fh.read_to_string(&mut buf)?;
        let bytes: Vec<u8> = buf.into();
        let mut da = DataArray::wrap(&ArrayType::MZArray, BinaryDataArrayType::Float64, bytes);
        da.compression = BinaryCompressionType::Zlib;
        Ok(da)
    }

    #[test]
    fn test_construction() -> io::Result<()> {
        let da = make_array_from_file()?;
        let mut map = BinaryArrayMap::new();
        assert!(!map.has_array(&ArrayType::MZArray));
        map.add(da);
        assert!(map.has_array(&ArrayType::MZArray));
        Ok(())
    }

    #[test]
    fn test_decode() -> io::Result<()> {
        let da = make_array_from_file()?;
        let mut map = BinaryArrayMap::new();
        map.add(da);
        assert_eq!(
            map.get(&ArrayType::MZArray).unwrap().compression,
            BinaryCompressionType::Zlib
        );
        map.decode_all_arrays()?;
        assert_eq!(
            map.get(&ArrayType::MZArray).unwrap().compression,
            BinaryCompressionType::Decoded
        );
        Ok(())
    }

    #[test]
    fn test_sort_from() -> io::Result<()> {
        let mut arr = make_array_from_file()?;
        let dat = arr.coerce_mut::<f64>()?;
        dat.reverse();

        let mut arr2 = make_array_from_file()?;
        arr2.name = ArrayType::TimeArray;
        arr2.store_as(BinaryDataArrayType::Float32)?;
        let dat = arr2.coerce_mut::<f32>()?;
        dat.reverse();

        let mut map = BinaryArrayMap::new();
        map.add(arr);
        map.add(arr2);
        map.sort_by_array(&ArrayType::MZArray)?;
        let arr = map.mzs()?;
        let (_, state) = arr.iter().copied().fold((-f64::INFINITY, true), |(last, state), next|{
            if !state {
                (last, state)
            } else {
                (next, last <= next)
            }
        });
        assert!(state, "Failed to sort m/z array {arr:?}");

        let arr = map.get(&ArrayType::TimeArray).unwrap().to_f64()?;
        let (_, state) = arr.iter().copied().fold((-f64::INFINITY, true), |(last, state), next|{
            if !state {
                (last, state)
            } else {
                (next, last <= next)
            }
        });
        assert!(state, "Failed to sort parallel array {arr:?}");
        Ok(())
    }
}