lasprs 0.14.1

//! Some measurement helper tools, specifally chunk-based reading of the data
//! from HDF5 datasets, according to our specification of a measurmenent file,
//! where the raw data is stored in blocks, that correspond to the number of
//! frames per block that comes from a DAQ device.
use super::*;
use crate::file_attr_names::{
    INPUT_SIGNAL_CLIPPED_ATTR, MEASURED_INPUT_DATASET_NAME, MEASUREMENT_STILL_RUNNING,
};
use crate::{Flt, config::*, daq::DataType};
use dasp_sample::{Sample, ToSample};
use hdf5_metno::Dataset;
use ndarray::Array2;
use num::traits::SaturatingSub;
use reinterpret::reinterpret_vec;
use smallvec::{SmallVec, ToSmallVec};
use snafu::prelude::*;
use std::iter::Iterator;
use std::marker::PhantomData;
type Result<T> = std::result::Result<T, MeasurementError>;

/// Default size hint for each chunk in MiB.
pub const DEFAULT_MIB_PER_CHUNK: usize = 5;

/// De-interleaved chunk of data as stored in the HDF5 file. Data is stored in
/// channel-by-channel, and after that sample-by-sample for each channel.
pub enum RawChunk {
    Datai8(Array2<i8>),
    Datai16(Array2<i16>),
    Datai24(Array2<dasp_sample::I24>),
    Datai32(Array2<i32>),
    Dataf32(Array2<f32>),
    Dataf64(Array2<f64>),
}
impl RawChunk {
    pub(crate) fn toFloat(&self) -> Array2<Flt> {
        match self {
            RawChunk::Datai8(arr) => arr.map(|x| Flt::from_sample(*x)),
            RawChunk::Datai16(arr) => arr.map(|x| Flt::from_sample(*x)),
            RawChunk::Datai24(arr) => arr.map(|x| Flt::from_sample(*x)),
            RawChunk::Datai32(arr) => arr.map(|x| Flt::from_sample(*x)),
            RawChunk::Dataf32(arr) => arr.map(|x| Flt::from_sample(*x)),
            RawChunk::Dataf64(arr) => arr.map(|x| Flt::from_sample(*x)),
        }
    }
}

/// Read a single (full) block from the HDF5 file. Reads all channels.
///
/// # Args
///
/// - `dataset`: The HDF5 dataset to read from.
/// - `channels`: Optional list of channels to read. If None, reads data from all channels.
/// - `tot_nchannels`: Total number of channels in the dataset.
/// - `block_idx`: Index of the block to read.
/// - `dtype`: Data type of the dataset.
/// - `startframe`: Start frame in the block.
/// - `nframes`: Number of frames to read in the block.
#[inline]
fn read_block(
    dataset: &Dataset,
    channels: Option<&[usize]>,
    dtype: DataType,
    startframe: usize,
    nframes: usize,
) -> RawChunk {
    let tot_nchannels = dataset.shape()[1];
    let nchannels = channels.map(|ch| ch.len()).unwrap_or(tot_nchannels);
    macro_rules! get_chunk {
        ($t:ty, $u:expr) => {{
            let dat: Array2<$t> = dataset
                .read_slice((startframe..startframe + nframes, ..))
                .expect("Cannot read slice");
            if let Some(channels) = channels {
                // TODO: Can we avoid the copy here?
                let mut block_data = Array2::<$t>::zeros((nframes, nchannels));

                for (i, ch) in channels.iter().copied().enumerate() {
                    let ax_in = dat.column(ch);
                    let mut ax_out = block_data.column_mut(i);
                    ax_out.assign(&ax_in);
                }
                $u(block_data)
            } else {
                $u(dat)
            }
        }};
    }
    match dtype {
        DataType::F32 => {
            get_chunk!(f32, RawChunk::Dataf32)
        }
        DataType::F64 => {
            get_chunk!(f64, RawChunk::Dataf64)
        }
        DataType::I8 => {
            get_chunk!(i8, RawChunk::Datai8)
        }
        DataType::I16 => {
            get_chunk!(i16, RawChunk::Datai16)
        }
        DataType::I24 => {
            let RawChunk::Datai32(v) = get_chunk!(i32, RawChunk::Datai32) else {
                panic!("Expected Datai32")
            };
            let (v, _) = v.into_raw_vec_and_offset();
            let vi24 = unsafe { reinterpret_vec(v) };
            let vi24 = Array2::from_shape_vec((nchannels, nframes), vi24).unwrap();
            RawChunk::Datai24(vi24)
        }
        DataType::I32 => {
            get_chunk!(i32, RawChunk::Datai32)
        }
    }
}

pub struct RawBlockIter<'a> {
    // We take a handle to the measurement, just to make sure a RawBlockIter
    // does not outlive the measurement.
    meas: PhantomData<&'a Measurement>,
    dataset: Dataset,
    dtype: DataType,
    istart: usize,
    iend: usize,
    tot_nchannels: usize,
    channels: Option<SmallVec<[usize; TYPICAL_CHANNELS]>>,
    max_frames_per_chunk: usize,
}

impl<'a> RawBlockIter<'a> {
    /// # Create a new iterator over raw blocks of data.
    ///
    /// # Args
    ///
    /// * `meas` - The measurement to read from.
    /// * `dataset` - The dataset to read from (owned; the dataset handle is
    ///   refcounted, so the underlying HDF5 file stays open while iterating).
    /// * `istart_and_iend` - Optional start and end indices for the blocks to read.
    /// * `channels` - Optional slice of channel indices to include in the chunks.
    /// * `MiB_per_chunk` - Size hint for the each chunk in MiB. Defaults to 5 MiB.
    pub fn new(
        meas: &'a Measurement,
        dataset: Dataset,
        istart_and_iend: Option<(usize, usize)>,
        channels: Option<&[usize]>,
        MiB_per_chunk: Option<usize>,
    ) -> Result<Self> {
        let shape = dataset.shape();
        let nframes = shape[0];
        let MiB_per_chunk = MiB_per_chunk.unwrap_or(DEFAULT_MIB_PER_CHUNK);
        let dtype = meas.dataType();
        let bytes_per_sample = dtype.bytes_per_sample_in_memory();
        ensure!(
            MiB_per_chunk > 0,
            LogicSnafu {
                name: meas.name(),
                message: "MiB_per_chunk must be greater than 0".to_string()
            }
        );
        let tot_nchannels = shape[1];
        let nchannels = if let Some(channels) = &channels {
            ensure! {!channels.is_empty(), LogicSnafu {name: meas.name(),
            message: "No channels specified".to_string()}}
            let max_ch = *channels.iter().max().unwrap();
            ensure!(
                max_ch < tot_nchannels,
                ChannelIdxOutOfBoundsSnafu {
                    channel_idx: max_ch,
                    max_channels: tot_nchannels
                }
            );
            channels.len()
        } else {
            tot_nchannels
        };
        let max_frames_per_chunk = MiB_per_chunk * 1024 * 1024 / nchannels / bytes_per_sample;
        let (istart, iend) = istart_and_iend.unwrap_or((0, nframes));
        ensure!(
            iend <= nframes,
            SampleIdxOutOfBoundsSnafu {
                sample_idx: iend,
                max_samples: nframes
            }
        );
        ensure!(
            istart <= iend,
            LogicSnafu {
                name: meas.name(),
                message: "istart must not exceed iend".to_string()
            }
        );
        let channels = channels.map(|channels| channels.to_smallvec());
        Ok(Self {
            meas: PhantomData,
            dataset,
            dtype,
            istart,
            iend,
            tot_nchannels,
            channels,
            max_frames_per_chunk,
        })
    }
}
impl<'a> Iterator for RawBlockIter<'a> {
    type Item = RawChunk;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let frames_remaining = self.iend.saturating_sub(self.istart);
        if frames_remaining > 0 {
            // How many frames are we going to read this time? All, or up to the
            // chunk size.
            let nframes = frames_remaining.min(self.max_frames_per_chunk);

            let block = read_block(
                &self.dataset,
                self.channels.as_deref(),
                self.dtype,
                self.istart,
                nframes,
            );

            self.istart += nframes;
            Some(block)
        } else {
            None
        }
    }
}

pub struct ConvertedBlockIter<'a> {
    raw_iter: RawBlockIter<'a>,
    sensitivities: Option<SmallVec<[f64; TYPICAL_CHANNELS]>>,
}

impl<'a> ConvertedBlockIter<'a> {
    /// Create a new iterator over converted blocks, reading from a dataset
    /// of the given measurement. The dataset is opened from the measurement
    /// itself, so it is always consistent with the measurement's data type
    /// and channel configuration.
    ///
    /// # Arguments
    /// * `meas` - The measurement to read from. Borrowed for the
    ///   lifetime of the iterator (lifetime guard).
    /// * `dataset_name` - Name of the dataset to read (e.g.
    ///   `MEASURED_INPUT_DATASET_NAME` or `MEASURED_PRECAPTURE_DATASET_NAME`).
    /// * `istart_and_iend` - The start and end indices of the dataset to read. If not specified,
    ///   the entire dataset will be read.
    ///
    ///  * `channels` - The channels to read.
    ///  * `sensitivities` - The sensitivities to apply. If not specified,
    ///    no sensitivity correction will be applied.
    ///  * `MiB_per_chunk` - The size of each chunk in MiB. If not specified,
    ///    the default value will be used.
    ///
    /// # Returns
    /// A new iterator over measurement data as floating-point values.
    pub fn new(
        meas: &'a Measurement,
        dataset_name: &str,
        istart_and_iend: Option<(usize, usize)>,
        channels: Option<&[usize]>,
        sensitivities: Option<&[f64]>,
        MiB_per_chunk: Option<usize>,
    ) -> Result<Self> {
        let name = meas.name();
        let file = meas.open_file(false)?;
        let dataset =
            file.dataset(dataset_name)
                .map_err(H5Error::from)
                .context(H5FileProblemSnafu {
                    operation: format!("opening dataset '{dataset_name}'"),
                })?;
        let raw_iter = RawBlockIter::new(meas, dataset, istart_and_iend, channels, MiB_per_chunk)?;
        if let Some(sensitivities) = &sensitivities {
            let errmsg = "Number of channels and sensitivities must match";
            match &raw_iter.channels {
                Some(ch) => {
                    ensure!(
                        ch.len() == sensitivities.len(),
                        LogicSnafu {
                            name,
                            message: errmsg
                        }
                    );
                }
                None => {
                    ensure!(
                        sensitivities.len() == raw_iter.tot_nchannels,
                        LogicSnafu {
                            name,
                            message: errmsg
                        }
                    );
                }
            }
        }
        let sensitivities = sensitivities.map(|sens| sens.to_smallvec());

        Ok(Self {
            raw_iter,
            sensitivities,
        })
    }
}
impl<'a> Iterator for ConvertedBlockIter<'a> {
    type Item = Array2<Flt>;

    fn next(&mut self) -> Option<Self::Item> {
        self.raw_iter.next().map(|rawblock| {
            let floatblock = rawblock.toFloat();

            match &self.sensitivities {
                Some(sens) => {
                    debug_assert_eq!(
                        sens.len(),
                        floatblock.ncols(),
                        "Number of sensitivities must match number of channels"
                    );
                    let mut converted = floatblock;
                    for (i, &s) in sens.iter().enumerate() {
                        // Divide each element by the sensitivity
                        converted.column_mut(i).map_inplace(|v| *v /= s);
                    }
                    converted
                }
                None => floatblock,
            }
        })
    }
}