tekhsi_rs 0.1.1

use num_complex::Complex64;
use smol_str::SmolStr;

/// Interleaved IQ samples in native integer formats.
///
/// ```rust
/// use tekhsi_rs::data::IqSamples;
///
/// let samples = IqSamples::I8(vec![1, 2, 3, 4]);
/// assert_eq!(samples.len(), 4);
/// ```
#[derive(Debug, Clone)]
pub enum IqSamples {
    I8(Vec<i8>),
    I16(Vec<i16>),
    I32(Vec<i32>),
}

impl IqSamples {
    /// Number of raw samples (I and Q interleaved).
    pub fn len(&self) -> usize {
        match self {
            IqSamples::I8(values) => values.len(),
            IqSamples::I16(values) => values.len(),
            IqSamples::I32(values) => values.len(),
        }
    }

    /// Returns true when the payload is empty.
    pub fn is_empty(&self) -> bool {
        match self {
            IqSamples::I8(values) => values.is_empty(),
            IqSamples::I16(values) => values.is_empty(),
            IqSamples::I32(values) => values.is_empty(),
        }
    }

    /// Returns an iterator over the raw interleaved IQ samples as f64 values.
    pub fn iter_raw(&self) -> Box<dyn ExactSizeIterator<Item = f64> + '_> {
        match self {
            Self::I8(v) => Box::new(v.iter().map(|&x| x as f64)),
            Self::I16(v) => Box::new(v.iter().map(|&x| x as f64)),
            Self::I32(v) => Box::new(v.iter().map(|&x| x as f64)),
        }
    }

    /// Returns complex samples scaled and offset using axis calibration.
    #[inline]
    pub fn iter_normalized(
        &self,
        spacing: f64,
        offset: f64,
    ) -> Box<dyn ExactSizeIterator<Item = Complex64> + '_> {
        fn scaled_complex_iter<'a, T>(
            values: &'a [T],
            spacing: f64,
            offset: f64,
        ) -> Box<dyn ExactSizeIterator<Item = Complex64> + 'a>
        where
            T: Copy + Into<f64>,
        {
            Box::new(values.chunks_exact(2).map(move |c| {
                Complex64::new(
                    2.0f64 * (c[0].into() * spacing + offset),
                    2.0f64 * (c[1].into() * spacing + offset),
                )
            }))
        }

        match self {
            Self::I8(v) => scaled_complex_iter(v, spacing, offset),
            Self::I16(v) => scaled_complex_iter(v, spacing, offset),
            Self::I32(v) => scaled_complex_iter(v, spacing, offset),
        }
    }
}

/// IQ metadata describing frequency and windowing configuration.
///
/// ```rust
/// use tekhsi_rs::data::IqMetaInfo;
/// use smol_str::SmolStr;
///
/// let info = IqMetaInfo {
///     iq_center_frequency: 1.0,
///     iq_fft_length: 1024.0,
///     iq_resolution_bandwidth: 1.0,
///     iq_span: 2.0,
///     iq_window_type: SmolStr::new("Rectangle"),
///     iq_sample_rate: 1.0,
/// };
/// assert_eq!(info.iq_fft_length, 1024.0);
/// ```
#[derive(Debug, Clone)]
pub struct IqMetaInfo {
    /// Center frequency of IQ acquisition in Hz.
    pub iq_center_frequency: f64,
    /// FFT length used for IQ computation.
    pub iq_fft_length: f64,
    /// Resolution bandwidth (RBW) of spectrum in Hz.
    pub iq_resolution_bandwidth: f64,
    /// Frequency span of IQ acquisition in Hz.
    pub iq_span: f64,
    /// Window type applied to IQ data (e.g., "Rectangle", "Hann").
    pub iq_window_type: SmolStr,
    /// Computed sample rate based on FFT length and window type.
    pub iq_sample_rate: f64,
}

/// Decoded IQ waveform with axis metadata.
///
/// ```rust
/// use tekhsi_rs::data::{IqMetaInfo, IqSamples, IqWaveform};
/// use smol_str::SmolStr;
///
/// let waveform = IqWaveform {
///     source_name: SmolStr::new("ch1_iq"),
///     interleaved_iq: IqSamples::I8(vec![1, 2]),
///     iq_axis_spacing: 1.0,
///     iq_axis_offset: 0.0,
///     iq_axis_units: SmolStr::new("V"),
///     x_axis_spacing: 1.0,
///     x_axis_units: SmolStr::new("s"),
///     trigger_index: 0.0,
///     meta_info: IqMetaInfo {
///         iq_center_frequency: 0.0,
///         iq_fft_length: 0.0,
///         iq_resolution_bandwidth: 0.0,
///         iq_span: 0.0,
///         iq_window_type: SmolStr::new("Rectangle"),
///         iq_sample_rate: 0.0,
///     },
/// };
/// assert_eq!(waveform.record_length(), 1);
/// ```
#[derive(Debug, Clone)]
pub struct IqWaveform {
    /// Channel or waveform identifier (e.g., "ch1_iq", "spectrum1").
    pub source_name: SmolStr,
    /// Interleaved I/Q samples (even indices = I, odd indices = Q).
    pub interleaved_iq: IqSamples,
    /// Scale factor applied to raw IQ samples.
    pub iq_axis_spacing: f64,
    /// Offset applied to scaled IQ samples.
    pub iq_axis_offset: f64,
    /// Unit string for IQ axis (typically "V").
    pub iq_axis_units: SmolStr,
    /// Sample period in seconds (time per sample).
    pub x_axis_spacing: f64,
    /// Unit string for horizontal axis (typically "s").
    pub x_axis_units: SmolStr,
    /// Trigger location expressed as sample index from trigger.
    pub trigger_index: f64,
    /// IQ-specific metadata including frequency, span, RBW, and window info.
    pub meta_info: IqMetaInfo,
}

impl IqWaveform {
    /// Returns the number of complex samples.
    pub fn record_length(&self) -> usize {
        self.interleaved_iq.len() / 2
    }

    /// Returns normalized I values.
    #[inline]
    pub fn iter_normalized_i_values(&self) -> impl ExactSizeIterator<Item = f64> {
        self.interleaved_iq
            .iter_raw()
            .step_by(2)
            .map(|v| 2.0f64 * v * self.iq_axis_spacing + self.iq_axis_offset)
    }

    /// Returns normalized Q values.
    #[inline]
    pub fn iter_normalized_q_values(&self) -> impl ExactSizeIterator<Item = f64> {
        self.interleaved_iq
            .iter_raw()
            .skip(1)
            .step_by(2)
            .map(|v| 2.0f64 * v * self.iq_axis_spacing + self.iq_axis_offset)
    }

    /// Returns normalized complex samples as an iterator.
    ///
    /// ```rust
    /// use smol_str::SmolStr;
    /// use num_complex::Complex64;
    /// use tekhsi_rs::data::{IqMetaInfo, IqSamples, IqWaveform};
    ///
    /// let waveform = IqWaveform {
    ///     source_name: SmolStr::new("ch1_iq"),
    ///     interleaved_iq: IqSamples::I16(vec![0, 2, 4, 6]),
    ///     iq_axis_spacing: 0.5,
    ///     iq_axis_offset: 1.0,
    ///     iq_axis_units: SmolStr::new("V"),
    ///     x_axis_spacing: 1.0,
    ///     x_axis_units: SmolStr::new("s"),
    ///     trigger_index: 0.0,
    ///     meta_info: IqMetaInfo {
    ///         iq_center_frequency: 0.0,
    ///         iq_fft_length: 0.0,
    ///         iq_resolution_bandwidth: 0.0,
    ///         iq_span: 0.0,
    ///         iq_window_type: SmolStr::new("Rectangle"),
    ///         iq_sample_rate: 0.0,
    ///     },
    /// };
    /// let values: Vec<Complex64> = waveform.iter_normalized_values().collect();
    /// assert_eq!(values.len(), 2);
    /// ```
    #[inline]
    pub fn iter_normalized_values(&self) -> impl ExactSizeIterator<Item = Complex64> {
        self.interleaved_iq
            .iter_normalized(self.iq_axis_spacing, self.iq_axis_offset)
    }
}

#[cfg(test)]
#[cfg_attr(coverage, coverage(off))]
mod tests {
    use super::{IqMetaInfo, IqSamples, IqWaveform};
    use crate::data::Waveform;
    use num_complex::Complex64;
    use smol_str::SmolStr;

    #[test]
    fn record_length_iq_is_half_interleaved() {
        let waveform = IqWaveform {
            source_name: SmolStr::new("ch1_iq"),
            interleaved_iq: IqSamples::I16(vec![1, 2, 3, 4]),
            iq_axis_spacing: 1.0,
            iq_axis_offset: 0.0,
            iq_axis_units: SmolStr::new("V"),
            x_axis_spacing: 1.0,
            x_axis_units: SmolStr::new("s"),
            trigger_index: 0.0,
            meta_info: IqMetaInfo {
                iq_center_frequency: 0.0,
                iq_fft_length: 0.0,
                iq_resolution_bandwidth: 0.0,
                iq_span: 0.0,
                iq_window_type: SmolStr::new("Rectangle"),
                iq_sample_rate: 0.0,
            },
        };

        assert_eq!(waveform.record_length(), 2);
        assert_eq!(Waveform::Iq(waveform).record_length(), 2);
    }

    #[test]
    fn iq_axis_values_split_interleaved() {
        let waveform = IqWaveform {
            source_name: SmolStr::new("ch1_iq"),
            interleaved_iq: IqSamples::I16(vec![1, 2, 3, 4]),
            iq_axis_spacing: 1.0,
            iq_axis_offset: 0.0,
            iq_axis_units: SmolStr::new("V"),
            x_axis_spacing: 1.0,
            x_axis_units: SmolStr::new("s"),
            trigger_index: 0.0,
            meta_info: IqMetaInfo {
                iq_center_frequency: 0.0,
                iq_fft_length: 0.0,
                iq_resolution_bandwidth: 0.0,
                iq_span: 0.0,
                iq_window_type: SmolStr::new("Rectangle"),
                iq_sample_rate: 0.0,
            },
        };
        let i_values: Vec<f64> = waveform.iter_normalized_i_values().collect();
        let q_values: Vec<f64> = waveform.iter_normalized_q_values().collect();
        assert_eq!(i_values, vec![2.0, 6.0]);
        assert_eq!(q_values, vec![4.0, 8.0]);
    }

    #[test]
    fn iq_normalized_vertical_values_applies_scale_offset() {
        let waveform = IqWaveform {
            source_name: SmolStr::new("ch1_iq"),
            interleaved_iq: IqSamples::I16(vec![0, 2, 4, 6]),
            iq_axis_spacing: 0.5,
            iq_axis_offset: 1.0,
            iq_axis_units: SmolStr::new("V"),
            x_axis_spacing: 1.0,
            x_axis_units: SmolStr::new("s"),
            trigger_index: 0.0,
            meta_info: IqMetaInfo {
                iq_center_frequency: 0.0,
                iq_fft_length: 0.0,
                iq_resolution_bandwidth: 0.0,
                iq_span: 0.0,
                iq_window_type: SmolStr::new("Rectangle"),
                iq_sample_rate: 0.0,
            },
        };
        let values: Vec<Complex64> = waveform.iter_normalized_values().collect();
        assert_eq!(
            values,
            vec![Complex64::new(2.0, 4.0), Complex64::new(6.0, 8.0)]
        );
    }

    #[test]
    fn iq_samples_len_and_empty_for_all_variants() {
        let samples_i8 = IqSamples::I8(vec![1, 2]);
        let samples_i16 = IqSamples::I16(vec![1, 2]);
        let samples_i32 = IqSamples::I32(vec![1, 2]);
        assert_eq!(samples_i8.len(), 2);
        assert_eq!(samples_i16.len(), 2);
        assert_eq!(samples_i32.len(), 2);
        assert!(!samples_i8.is_empty());
        assert!(!samples_i16.is_empty());
        assert!(!samples_i32.is_empty());

        assert!(IqSamples::I8(Vec::new()).is_empty());
        assert!(IqSamples::I16(Vec::new()).is_empty());
        assert!(IqSamples::I32(Vec::new()).is_empty());
    }

    #[test]
    fn iq_samples_iter_raw_and_normalized_i32() {
        let samples = IqSamples::I32(vec![1, 2, 3, 4]);
        let raw: Vec<f64> = samples.iter_raw().collect();
        assert_eq!(raw, vec![1.0, 2.0, 3.0, 4.0]);

        let normalized: Vec<Complex64> = samples.iter_normalized(0.5, 1.0).collect();
        assert_eq!(
            normalized,
            vec![Complex64::new(3.0, 4.0), Complex64::new(5.0, 6.0)]
        );
    }

    #[test]
    fn iq_samples_iter_normalized_i8() {
        let samples = IqSamples::I8(vec![1, 2, 3, 4]);
        let raw: Vec<f64> = samples.iter_raw().collect();
        assert_eq!(raw, vec![1.0, 2.0, 3.0, 4.0]);
        let normalized: Vec<Complex64> = samples.iter_normalized(0.5, 1.0).collect();
        assert_eq!(
            normalized,
            vec![Complex64::new(3.0, 4.0), Complex64::new(5.0, 6.0)]
        );
    }
}