rustdf/data/

handle.rs

use crate::data::meta::{read_global_meta_sql, read_meta_data_sql, FrameMeta, GlobalMetaData};
use crate::data::raw::BrukerTimsDataLibrary;
use crate::data::utility::{
    flatten_scan_values, parse_decompressed_bruker_binary_data, zstd_decompress,
};
use byteorder::{LittleEndian, ReadBytesExt};
use mscore::data::spectrum::MsType;
use mscore::timstof::frame::{ImsFrame, RawTimsFrame, TimsFrame};
use mscore::timstof::slice::TimsSlice;
use std::fs::File;
use std::io::{Cursor, Read, Seek, SeekFrom};
use std::path::PathBuf;

use crate::data::acquisition::AcquisitionMode;
use rayon::prelude::*;
use rayon::ThreadPoolBuilder;

use std::error::Error;

/// Derive m/z calibration coefficients by using the Bruker SDK.
///
/// This function temporarily loads the SDK to convert a range of TOF indices to m/z,
/// then fits a linear regression to derive accurate calibration coefficients.
///
/// Returns `Some((intercept, slope))` for the formula: `sqrt(mz) = intercept + slope * tof`
/// Returns `None` if SDK is not available (e.g., on macOS).
fn derive_mz_calibration(
    bruker_lib_path: &str,
    data_path: &str,
    tof_max_index: u32,
) -> Option<(f64, f64)> {
    // Check if SDK path is valid before trying to load
    // Common invalid values: "NO_SDK", empty string, or non-existent paths
    if bruker_lib_path.is_empty()
        || bruker_lib_path == "NO_SDK"
        || bruker_lib_path == "CALIBRATED"
        || !std::path::Path::new(bruker_lib_path).exists()
    {
        return None;
    }

    // Try to create a BrukerLib converter. Use the FALLIBLE `try_new`, not `catch_unwind` around the
    // panicking `new`: the workspace builds release with `panic = "abort"`, so catch_unwind can't
    // recover — a present-but-unloadable SDK (wrong arch / mismatched version) would abort instead of
    // returning None. try_new returns Err on load failure with no panic.
    let sdk_converter = match BrukerLibTimsDataConverter::try_new(bruker_lib_path, data_path) {
        Ok(converter) => converter,
        Err(_) => return None,
    };

    // Generate a range of TOF indices across the spectrum
    let n_points = 1000;
    let step = tof_max_index / n_points;
    let tof_indices: Vec<u32> = (0..n_points).map(|i| i * step + step / 2).collect();

    // Convert to m/z using SDK
    let mz_values = sdk_converter.tof_to_mz(1, &tof_indices);

    // Filter out any invalid values (zero or negative)
    let valid_pairs: Vec<(f64, f64)> = tof_indices
        .iter()
        .zip(mz_values.iter())
        .filter(|(_, mz)| **mz > 0.0)
        .map(|(tof, mz)| (*tof as f64, mz.sqrt()))
        .collect();

    if valid_pairs.len() < 10 {
        return None;
    }

    // Linear regression: sqrt(mz) = intercept + slope * tof
    // Using simple least squares: slope = Cov(x,y) / Var(x), intercept = mean_y - slope * mean_x
    let n = valid_pairs.len() as f64;
    let sum_x: f64 = valid_pairs.iter().map(|(x, _)| x).sum();
    let sum_y: f64 = valid_pairs.iter().map(|(_, y)| y).sum();
    let sum_xy: f64 = valid_pairs.iter().map(|(x, y)| x * y).sum();
    let sum_xx: f64 = valid_pairs.iter().map(|(x, _)| x * x).sum();

    let mean_x = sum_x / n;
    let mean_y = sum_y / n;

    let slope = (sum_xy - n * mean_x * mean_y) / (sum_xx - n * mean_x * mean_x);
    let intercept = mean_y - slope * mean_x;

    Some((intercept, slope))
}

/// Build a `LookupIndexConverter`, preferring an SDK-derived m/z calibration.
///
/// The `LookupIndexConverter` always carries the pre-computed scan→1/K0 lookup
/// for ion mobility. For m/z it normally uses a 2-point boundary model
/// (`LookupIndexConverter::new`), which can have a large m/z error on some
/// datasets. When a valid `bruker_lib_path` is supplied, m/z is instead
/// calibrated with the same regression fit used by the `Calibrated` converter
/// (`derive_mz_calibration`). Falls back to the boundary model when the SDK is
/// unavailable (e.g. macOS, or `bruker_lib_path` is empty / "NO_SDK").
fn build_lookup_converter(
    bruker_lib_path: &str,
    data_path: &str,
    tof_max_index: u32,
    mz_lower: f64,
    mz_upper: f64,
    im_lookup: Vec<f64>,
) -> LookupIndexConverter {
    match derive_mz_calibration(bruker_lib_path, data_path, tof_max_index) {
        Some((intercept, slope)) => {
            LookupIndexConverter::with_mz_fit(intercept, slope, im_lookup)
        }
        None => {
            if !bruker_lib_path.is_empty() && bruker_lib_path != "NO_SDK" {
                eprintln!(
                    "Warning: Could not derive m/z calibration from SDK at '{}'. \
                    Falling back to the 2-point boundary model, which may have a \
                    large m/z error on some datasets.",
                    bruker_lib_path
                );
            }
            LookupIndexConverter::new(mz_lower, mz_upper, tof_max_index, im_lookup)
        }
    }
}

fn lzf_decompress(data: &[u8], max_output_size: usize) -> Result<Vec<u8>, Box<dyn Error>> {
    let decompressed_data = lzf::decompress(data, max_output_size)
        .map_err(|e| format!("LZF decompression failed: {}", e))?;
    Ok(decompressed_data)
}

fn parse_decompressed_bruker_binary_type1(
    decompressed_bytes: &[u8],
    scan_indices: &mut [i64],
    tof_indices: &mut [u32],
    intensities: &mut [u16],
    scan_start: usize,
    scan_index: usize,
) -> usize {
    // Interpret decompressed_bytes as a slice of i32
    let int_count = decompressed_bytes.len() / 4;
    let buffer =
        unsafe { std::slice::from_raw_parts(decompressed_bytes.as_ptr() as *const i32, int_count) };

    let mut tof_index = 0i32;
    let mut previous_was_intensity = true;
    let mut current_index = scan_start;

    for &value in buffer {
        if value >= 0 {
            // positive value => intensity
            if previous_was_intensity {
                tof_index += 1;
            }
            tof_indices[current_index] = tof_index as u32;
            intensities[current_index] = value as u16;
            previous_was_intensity = true;
            current_index += 1;
        } else {
            // negative value => indicates a jump in tof_index
            tof_index -= value; // value is negative, so this adds |value| to tof_index
            previous_was_intensity = false;
        }
    }

    let scan_size = current_index - scan_start;
    scan_indices[scan_index] = scan_size as i64;
    scan_size
}

pub struct TimsRawDataLayout {
    pub raw_data_path: String,
    pub global_meta_data: GlobalMetaData,
    pub frame_meta_data: Vec<FrameMeta>,
    pub max_scan_count: i64,
    pub frame_id_ptr: Vec<i64>,
    pub tims_offset_values: Vec<i64>,
    pub acquisition_mode: AcquisitionMode,
}

impl TimsRawDataLayout {
    pub fn new(data_path: &str) -> Self {
        // get the global and frame meta data
        let global_meta_data = read_global_meta_sql(data_path).unwrap();
        let frame_meta_data = read_meta_data_sql(data_path).unwrap();

        // get the max scan count
        let max_scan_count = frame_meta_data.iter().map(|x| x.num_scans).max().unwrap();

        let mut frame_id_ptr: Vec<i64> = Vec::new();
        frame_id_ptr.resize(frame_meta_data.len() + 1, 0);

        // get the frame id_ptr values
        for (i, row) in frame_meta_data.iter().enumerate() {
            frame_id_ptr[i + 1] = row.num_peaks + frame_id_ptr[i];
        }

        // get the tims offset values
        let tims_offset_values = frame_meta_data
            .iter()
            .map(|x| x.tims_id)
            .collect::<Vec<i64>>();

        // get the acquisition mode
        let acquisition_mode = match frame_meta_data[0].scan_mode {
            8 => AcquisitionMode::DDA,
            9 => AcquisitionMode::DIA,
            _ => AcquisitionMode::Unknown,
        };

        TimsRawDataLayout {
            raw_data_path: data_path.to_string(),
            global_meta_data,
            frame_meta_data,
            max_scan_count,
            frame_id_ptr,
            tims_offset_values,
            acquisition_mode,
        }
    }
}

pub trait TimsData {
    fn get_frame(&self, frame_id: u32) -> TimsFrame;
    fn get_raw_frame(&self, frame_id: u32) -> RawTimsFrame;
    fn get_slice(&self, frame_ids: Vec<u32>, num_threads: usize) -> TimsSlice;
    fn get_acquisition_mode(&self) -> AcquisitionMode;
    fn get_frame_count(&self) -> i32;
    fn get_data_path(&self) -> &str;
}

pub trait IndexConverter {
    fn tof_to_mz(&self, frame_id: u32, tof_values: &Vec<u32>) -> Vec<f64>;
    fn mz_to_tof(&self, frame_id: u32, mz_values: &Vec<f64>) -> Vec<u32>;
    fn scan_to_inverse_mobility(&self, frame_id: u32, scan_values: &Vec<u32>) -> Vec<f64>;
    fn inverse_mobility_to_scan(
        &self,
        frame_id: u32,
        inverse_mobility_values: &Vec<f64>,
    ) -> Vec<u32>;
}

pub struct BrukerLibTimsDataConverter {
    pub bruker_lib: BrukerTimsDataLibrary,
}

impl BrukerLibTimsDataConverter {
    pub fn new(bruker_lib_path: &str, data_path: &str) -> Self {
        Self::try_new(bruker_lib_path, data_path).unwrap()
    }
    /// Fallible constructor — `Err` if the Bruker SDK shared library can't be loaded (missing,
    /// wrong architecture, mismatched version). Lets callers degrade to the simple/derived
    /// calibration instead of panicking.
    pub fn try_new(
        bruker_lib_path: &str,
        data_path: &str,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let bruker_lib = BrukerTimsDataLibrary::new(bruker_lib_path, data_path)?;
        Ok(BrukerLibTimsDataConverter { bruker_lib })
    }
}
impl IndexConverter for BrukerLibTimsDataConverter {
    /// translate tof to mz values calling the bruker library
    ///
    /// # Arguments
    ///
    /// * `frame_id` - A u32 that holds the frame id
    /// * `tof` - A vector of u32 that holds the tof values
    ///
    /// # Returns
    ///
    /// * `mz_values` - A vector of f64 that holds the mz values
    ///
    fn tof_to_mz(&self, frame_id: u32, tof: &Vec<u32>) -> Vec<f64> {
        let mut dbl_tofs: Vec<f64> = Vec::new();
        dbl_tofs.resize(tof.len(), 0.0);

        for (i, &val) in tof.iter().enumerate() {
            dbl_tofs[i] = val as f64;
        }

        let mut mz_values: Vec<f64> = Vec::new();
        mz_values.resize(tof.len(), 0.0);

        self.bruker_lib
            .tims_index_to_mz(frame_id, &dbl_tofs, &mut mz_values)
            .expect("Bruker binary call failed at: tims_index_to_mz;");

        mz_values
    }

    fn mz_to_tof(&self, frame_id: u32, mz: &Vec<f64>) -> Vec<u32> {
        let mut dbl_mz: Vec<f64> = Vec::new();
        dbl_mz.resize(mz.len(), 0.0);

        for (i, &val) in mz.iter().enumerate() {
            dbl_mz[i] = val;
        }

        let mut tof_values: Vec<f64> = Vec::new();
        tof_values.resize(mz.len(), 0.0);

        self.bruker_lib
            .tims_mz_to_index(frame_id, &dbl_mz, &mut tof_values)
            .expect("Bruker binary call failed at: tims_mz_to_index;");

        tof_values.iter().map(|&x| x.round() as u32).collect()
    }

    /// translate scan to inverse mobility values calling the bruker library
    ///
    /// # Arguments
    ///
    /// * `frame_id` - A u32 that holds the frame id
    /// * `scan` - A vector of i32 that holds the scan values
    ///
    /// # Returns
    ///
    /// * `inv_mob` - A vector of f64 that holds the inverse mobility values
    ///
    fn scan_to_inverse_mobility(&self, frame_id: u32, scan: &Vec<u32>) -> Vec<f64> {
        let mut dbl_scans: Vec<f64> = Vec::new();
        dbl_scans.resize(scan.len(), 0.0);

        for (i, &val) in scan.iter().enumerate() {
            dbl_scans[i] = val as f64;
        }

        let mut inv_mob: Vec<f64> = Vec::new();
        inv_mob.resize(scan.len(), 0.0);

        self.bruker_lib
            .tims_scan_to_inv_mob(frame_id, &dbl_scans, &mut inv_mob)
            .expect("Bruker binary call failed at: tims_scannum_to_oneoverk0;");

        inv_mob
    }

    /// translate inverse mobility to scan values calling the bruker library
    ///
    /// # Arguments
    ///
    /// * `frame_id` - A u32 that holds the frame id
    /// * `inv_mob` - A vector of f64 that holds the inverse mobility values
    ///
    /// # Returns
    ///
    /// * `scan_values` - A vector of i32 that holds the scan values
    ///
    fn inverse_mobility_to_scan(&self, frame_id: u32, inv_mob: &Vec<f64>) -> Vec<u32> {
        let mut dbl_inv_mob: Vec<f64> = Vec::new();
        dbl_inv_mob.resize(inv_mob.len(), 0.0);

        for (i, &val) in inv_mob.iter().enumerate() {
            dbl_inv_mob[i] = val;
        }

        let mut scan_values: Vec<f64> = Vec::new();
        scan_values.resize(inv_mob.len(), 0.0);

        self.bruker_lib
            .inv_mob_to_tims_scan(frame_id, &dbl_inv_mob, &mut scan_values)
            .expect("Bruker binary call failed at: tims_oneoverk0_to_scannum;");

        scan_values.iter().map(|&x| x.round() as u32).collect()
    }
}

/// SDK-free converter using the exact Bruker calibration formulas.
///
/// Reads the `MzCalibration`, `TimsCalibration` and `Frames` tables from
/// analysis.tdf and evaluates the published calibration curves directly (see
/// [`crate::data::calibration`]). Verified against the Bruker SDK:
///   * scan <-> 1/K0 : machine-precision exact (ModelType 2).
///   * TOF  <-> m/z  : bit-exact for MzCalibration ModelType 1; ModelType 2 uses
///     the same C0/C1/C2 quadratic-in-sqrt(m) curve and reproduces the SDK to
///     ~1 ppm (the proprietary ModelType-2 C8..C14 fine correction is not
///     modelled).
///
/// The calibration is resolved **per frame**: every frame's `Frames.T1`/`T2` and
/// its `MzCalibration`/`TimsCalibration` foreign keys are read, so a run that
/// carries more than one calibration row is handled, and the temperature
/// compensation is evaluated for the frame the caller asks about. That matters
/// because `dC1` is ~20 ppm/degC: pinning one frame's temperature for a whole run
/// costs ~20 ppm of mass accuracy per degC the digitizer drifts.
/// [`Self::temperature_spread`] reports how far the temperatures actually moved,
/// so callers can tell a quiet run from a drifting one.
///
/// Frame ids that are not in the file, and frames whose calibration rows use an
/// unsupported ModelType, fall back to the calibrators built from
/// `fallback_frame_id`.
pub struct BrukerFormulaConverter {
    /// Calibrator built from `fallback_frame_id`; also used for unknown frames.
    pub mz: crate::data::calibration::MzCalibrator,
    /// Mobility calibrator for `fallback_frame_id`; also used for unknown frames.
    pub im: crate::data::calibration::MobilityCalibrator,
    /// `MzCalibration.ModelType` of the fallback frame's row.
    pub mz_model_type: i64,
    /// Supported `MzCalibration` rows, by `Id`.
    mz_rows: std::collections::HashMap<i64, crate::data::meta::MzCalibration>,
    /// Ready-built mobility calibrators, by `TimsCalibration.Id`.
    im_rows: std::collections::HashMap<i64, crate::data::calibration::MobilityCalibrator>,
    /// Per-frame calibration linkage and digitizer temperatures.
    frames: std::collections::HashMap<u32, FrameCalibrationRef>,
    /// (T1, T2) spread (max - min) over all frames of the run, in degC.
    temp_spread: (f64, f64),
}

/// One frame's link into the calibration tables, plus its digitizer temperatures.
struct FrameCalibrationRef {
    mz_id: i64,
    tims_id: i64,
    t1: f64,
    t2: f64,
}

impl BrukerFormulaConverter {
    /// Build the converter from a `.d` folder.
    ///
    /// `fallback_frame_id` (conventionally 1) selects the calibration used for
    /// frames that are missing from the `Frames` table or whose calibration rows
    /// this converter cannot evaluate; every frame that *is* in the file uses its
    /// own rows and temperatures.
    pub fn from_d_folder(
        data_path: &str,
        fallback_frame_id: u32,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        use crate::data::calibration::{MobilityCalibrator, MzCalibrator};
        use crate::data::meta::{read_mz_calibration, read_tims_calibration};
        use std::collections::HashMap;

        let tdf = PathBuf::from(data_path).join("analysis.tdf");
        let con = rusqlite::Connection::open(&tdf)?;

        // Per-frame calibration linkage + digitizer temperatures for the whole run.
        // T1/T2 are nullable in Bruker's schema; frames without them cannot be
        // temperature compensated and are left to the fallback calibration.
        let mut stmt =
            con.prepare("SELECT Id, T1, T2, MzCalibration, TimsCalibration FROM Frames")?;
        let rows = stmt.query_map([], |r| {
            Ok((
                r.get::<_, i64>(0)?,
                r.get::<_, Option<f64>>(1)?,
                r.get::<_, Option<f64>>(2)?,
                r.get::<_, i64>(3)?,
                r.get::<_, i64>(4)?,
            ))
        })?;

        let mut frames: HashMap<u32, FrameCalibrationRef> = HashMap::new();
        let (mut t1_min, mut t1_max) = (f64::INFINITY, f64::NEG_INFINITY);
        let (mut t2_min, mut t2_max) = (f64::INFINITY, f64::NEG_INFINITY);
        for row in rows {
            let (id, t1, t2, mz_id, tims_id) = row?;
            let (t1, t2) = match (t1, t2) {
                (Some(t1), Some(t2)) => (t1, t2),
                _ => continue,
            };
            t1_min = t1_min.min(t1);
            t1_max = t1_max.max(t1);
            t2_min = t2_min.min(t2);
            t2_max = t2_max.max(t2);
            frames.insert(
                id as u32,
                FrameCalibrationRef {
                    mz_id,
                    tims_id,
                    t1,
                    t2,
                },
            );
        }
        drop(stmt);

        let fallback = frames
            .get(&fallback_frame_id)
            .ok_or("fallback frame has no calibration/temperature row in Frames")?;

        let mz_all = read_mz_calibration(data_path)?;
        let tims_all = read_tims_calibration(data_path)?;

        let mzc = mz_all
            .iter()
            .find(|c| c.id == fallback.mz_id)
            .ok_or("MzCalibration row not found")?;
        let tc = tims_all
            .iter()
            .find(|c| c.id == fallback.tims_id)
            .ok_or("TimsCalibration row not found")?;

        // Reject calibration models this converter does not implement, rather
        // than silently mis-computing (m/z model 2 is the default branch).
        if mzc.model_type != 1 && mzc.model_type != 2 {
            return Err(format!("unsupported MzCalibration ModelType {}", mzc.model_type).into());
        }
        if tc.model_type != 2 {
            return Err(format!("unsupported TimsCalibration ModelType {}", tc.model_type).into());
        }

        let mz = MzCalibrator::from_calibration(mzc, fallback.t1, fallback.t2);
        let im = MobilityCalibrator::from_calibration(tc);
        let mz_model_type = mzc.model_type;

        // Keep only the rows this converter can evaluate; frames pointing at any
        // other row fall back to the calibration above.
        let mz_rows: HashMap<i64, crate::data::meta::MzCalibration> = mz_all
            .into_iter()
            .filter(|c| c.model_type == 1 || c.model_type == 2)
            .map(|c| (c.id, c))
            .collect();
        let im_rows: HashMap<i64, MobilityCalibrator> = tims_all
            .iter()
            .filter(|c| c.model_type == 2)
            .map(|c| (c.id, MobilityCalibrator::from_calibration(c)))
            .collect();

        let temp_spread = if frames.is_empty() {
            (0.0, 0.0)
        } else {
            (t1_max - t1_min, t2_max - t2_min)
        };

        Ok(Self {
            mz,
            im,
            mz_model_type,
            mz_rows,
            im_rows,
            frames,
            temp_spread,
        })
    }

    /// Spread (max - min, in degC) of `Frames.T1` and `Frames.T2` across the run.
    ///
    /// A near-zero spread means per-frame and fixed-frame calibration agree; a
    /// large one means the digitizer temperature moved and per-frame
    /// compensation is doing real work (~20 ppm of m/z per degC).
    pub fn temperature_spread(&self) -> (f64, f64) {
        self.temp_spread
    }

    /// m/z calibrator for `frame_id`, falling back to the default calibration.
    fn mz_for(&self, frame_id: u32) -> crate::data::calibration::MzCalibrator {
        self.frames
            .get(&frame_id)
            .and_then(|f| {
                self.mz_rows.get(&f.mz_id).map(|row| {
                    crate::data::calibration::MzCalibrator::from_calibration(row, f.t1, f.t2)
                })
            })
            .unwrap_or_else(|| self.mz.clone())
    }

    /// Mobility calibrator for `frame_id`, falling back to the default one.
    fn im_for(&self, frame_id: u32) -> &crate::data::calibration::MobilityCalibrator {
        self.frames
            .get(&frame_id)
            .and_then(|f| self.im_rows.get(&f.tims_id))
            .unwrap_or(&self.im)
    }
}

impl IndexConverter for BrukerFormulaConverter {
    fn tof_to_mz(&self, frame_id: u32, tof_values: &Vec<u32>) -> Vec<f64> {
        let mz = self.mz_for(frame_id);
        tof_values.iter().map(|&t| mz.tof_to_mz(t)).collect()
    }
    fn mz_to_tof(&self, frame_id: u32, mz_values: &Vec<f64>) -> Vec<u32> {
        let mz = self.mz_for(frame_id);
        mz_values.iter().map(|&m| mz.mz_to_tof(m)).collect()
    }
    fn scan_to_inverse_mobility(&self, frame_id: u32, scan_values: &Vec<u32>) -> Vec<f64> {
        let im = self.im_for(frame_id);
        scan_values
            .iter()
            .map(|&s| im.scan_to_one_over_k0(s))
            .collect()
    }
    fn inverse_mobility_to_scan(
        &self,
        frame_id: u32,
        inverse_mobility_values: &Vec<f64>,
    ) -> Vec<u32> {
        let im = self.im_for(frame_id);
        inverse_mobility_values
            .iter()
            .map(|&v| im.one_over_k0_to_scan(v))
            .collect()
    }
}

pub enum TimsIndexConverter {
    Simple(SimpleIndexConverter),
    Calibrated(CalibratedIndexConverter),
    BrukerLib(BrukerLibTimsDataConverter),
    Lookup(LookupIndexConverter),
    BrukerFormula(BrukerFormulaConverter),
}

impl IndexConverter for TimsIndexConverter {
    fn tof_to_mz(&self, frame_id: u32, tof_values: &Vec<u32>) -> Vec<f64> {
        match self {
            TimsIndexConverter::Simple(converter) => converter.tof_to_mz(frame_id, tof_values),
            TimsIndexConverter::Calibrated(converter) => converter.tof_to_mz(frame_id, tof_values),
            TimsIndexConverter::BrukerLib(converter) => converter.tof_to_mz(frame_id, tof_values),
            TimsIndexConverter::Lookup(converter) => converter.tof_to_mz(frame_id, tof_values),
            TimsIndexConverter::BrukerFormula(converter) => converter.tof_to_mz(frame_id, tof_values),
        }
    }

    fn mz_to_tof(&self, frame_id: u32, mz_values: &Vec<f64>) -> Vec<u32> {
        match self {
            TimsIndexConverter::Simple(converter) => converter.mz_to_tof(frame_id, mz_values),
            TimsIndexConverter::Calibrated(converter) => converter.mz_to_tof(frame_id, mz_values),
            TimsIndexConverter::BrukerLib(converter) => converter.mz_to_tof(frame_id, mz_values),
            TimsIndexConverter::Lookup(converter) => converter.mz_to_tof(frame_id, mz_values),
            TimsIndexConverter::BrukerFormula(converter) => converter.mz_to_tof(frame_id, mz_values),
        }
    }

    fn scan_to_inverse_mobility(&self, frame_id: u32, scan_values: &Vec<u32>) -> Vec<f64> {
        match self {
            TimsIndexConverter::Simple(converter) => {
                converter.scan_to_inverse_mobility(frame_id, scan_values)
            }
            TimsIndexConverter::Calibrated(converter) => {
                converter.scan_to_inverse_mobility(frame_id, scan_values)
            }
            TimsIndexConverter::BrukerLib(converter) => {
                converter.scan_to_inverse_mobility(frame_id, scan_values)
            }
            TimsIndexConverter::Lookup(converter) => {
                converter.scan_to_inverse_mobility(frame_id, scan_values)
            }
            TimsIndexConverter::BrukerFormula(converter) => {
                converter.scan_to_inverse_mobility(frame_id, scan_values)
            }
        }
    }

    fn inverse_mobility_to_scan(
        &self,
        frame_id: u32,
        inverse_mobility_values: &Vec<f64>,
    ) -> Vec<u32> {
        match self {
            TimsIndexConverter::Simple(converter) => {
                converter.inverse_mobility_to_scan(frame_id, inverse_mobility_values)
            }
            TimsIndexConverter::Calibrated(converter) => {
                converter.inverse_mobility_to_scan(frame_id, inverse_mobility_values)
            }
            TimsIndexConverter::BrukerLib(converter) => {
                converter.inverse_mobility_to_scan(frame_id, inverse_mobility_values)
            }
            TimsIndexConverter::Lookup(converter) => {
                converter.inverse_mobility_to_scan(frame_id, inverse_mobility_values)
            }
            TimsIndexConverter::BrukerFormula(converter) => {
                converter.inverse_mobility_to_scan(frame_id, inverse_mobility_values)
            }
        }
    }
}

pub struct TimsLazyLoder {
    pub raw_data_layout: TimsRawDataLayout,
    pub index_converter: TimsIndexConverter,
}

impl TimsData for TimsLazyLoder {
    fn get_frame(&self, frame_id: u32) -> TimsFrame {
        let frame_index = (frame_id - 1) as usize;

        // turns out, there can be empty frames in the data, check for that, if so, return an empty frame
        let num_peaks = self.raw_data_layout.frame_meta_data[frame_index].num_peaks;

        if num_peaks == 0 {

            let ms_type_raw = self.raw_data_layout.frame_meta_data[frame_index].ms_ms_type;

            let ms_type = match ms_type_raw {
                0 => MsType::Precursor,
                8 => MsType::FragmentDda,
                9 => MsType::FragmentDia,
                _ => MsType::Unknown,
            };

            return TimsFrame {
                frame_id: frame_id as i32,
                ms_type,
                scan: Vec::new(),
                tof: Vec::new(),
                ims_frame: ImsFrame::new(
                    self.raw_data_layout.frame_meta_data[(frame_id - 1) as usize].time,
                    Vec::new(),
                    Vec::new(),
                    Vec::new(),
                ),
            };
        }

        let offset = self.raw_data_layout.tims_offset_values[frame_index] as u64;

        let mut file_path = PathBuf::from(&self.raw_data_layout.raw_data_path);
        file_path.push("analysis.tdf_bin");
        let mut infile = File::open(&file_path).unwrap();

        infile.seek(SeekFrom::Start(offset)).unwrap();

        let mut bin_buffer = [0u8; 4];
        infile.read_exact(&mut bin_buffer).unwrap();
        let bin_size = Cursor::new(bin_buffer).read_i32::<LittleEndian>().unwrap();

        infile.read_exact(&mut bin_buffer).unwrap();

        match self.raw_data_layout.global_meta_data.tims_compression_type {
            1 => {
                let scan_count =
                    self.raw_data_layout.frame_meta_data[frame_index].num_scans as usize;
                let num_peaks = num_peaks as usize;
                let compression_offset = 8 + (scan_count + 1) * 4;

                let mut scan_offsets_buffer = vec![0u8; (scan_count + 1) * 4];
                infile.read_exact(&mut scan_offsets_buffer).unwrap();

                let mut scan_offsets = Vec::with_capacity(scan_count + 1);
                {
                    let mut rdr = Cursor::new(&scan_offsets_buffer);
                    for _ in 0..(scan_count + 1) {
                        scan_offsets.push(rdr.read_i32::<LittleEndian>().unwrap());
                    }
                }

                for offs in &mut scan_offsets {
                    *offs -= compression_offset as i32;
                }

                let remaining_size = (bin_size as usize - compression_offset) as usize;
                let mut compressed_data = vec![0u8; remaining_size];
                infile.read_exact(&mut compressed_data).unwrap();

                let mut scan_indices_ = vec![0i64; scan_count];
                let mut tof_indices_ = vec![0u32; num_peaks];
                let mut intensities_ = vec![0u16; num_peaks];

                let mut scan_start = 0usize;

                for scan_index in 0..scan_count {
                    let start = scan_offsets[scan_index] as usize;
                    let end = scan_offsets[scan_index + 1] as usize;

                    if start == end {
                        continue;
                    }

                    let max_output_size = num_peaks * 8;
                    let decompressed_bytes =
                        lzf_decompress(&compressed_data[start..end], max_output_size)
                            .expect("LZF decompression failed.");

                    scan_start += parse_decompressed_bruker_binary_type1(
                        &decompressed_bytes,
                        &mut scan_indices_,
                        &mut tof_indices_,
                        &mut intensities_,
                        scan_start,
                        scan_index,
                    );
                }

                // Create a flat scan vector to match what flatten_scan_values expects
                let mut scan = Vec::with_capacity(num_peaks);
                {
                    let mut current_scan_index = 0u32;
                    for &size in &scan_indices_ {
                        let sz = size as usize;
                        for _ in 0..sz {
                            scan.push(current_scan_index);
                        }
                        current_scan_index += 1;
                    }
                }

                let intensity_dbl = intensities_.iter().map(|&x| x as f64).collect::<Vec<f64>>();
                let tof_i32 = tof_indices_.iter().map(|&x| x as i32).collect::<Vec<i32>>();

                let mz = self.index_converter.tof_to_mz(frame_id, &tof_indices_);
                let inv_mobility = self
                    .index_converter
                    .scan_to_inverse_mobility(frame_id, &scan);

                let ms_type_raw = self.raw_data_layout.frame_meta_data[frame_index].ms_ms_type;
                let ms_type = match ms_type_raw {
                    0 => MsType::Precursor,
                    8 => MsType::FragmentDda,
                    9 => MsType::FragmentDia,
                    _ => MsType::Unknown,
                };

                TimsFrame {
                    frame_id: frame_id as i32,
                    ms_type,
                    scan: scan.iter().map(|&x| x as i32).collect(),
                    tof: tof_i32,
                    ims_frame: ImsFrame::new(
                        self.raw_data_layout.frame_meta_data[frame_index].time,
                        inv_mobility,
                        mz,
                        intensity_dbl,
                    ),
                }
            }

            // Existing handling of Type 2
            2 => {
                let mut compressed_data = vec![0u8; bin_size as usize - 8];
                infile.read_exact(&mut compressed_data).unwrap();

                let decompressed_bytes = zstd_decompress(&compressed_data).unwrap();

                let (scan, tof, intensity) =
                    parse_decompressed_bruker_binary_data(&decompressed_bytes).unwrap();
                let intensity_dbl = intensity.iter().map(|&x| x as f64).collect();
                let tof_i32 = tof.iter().map(|&x| x as i32).collect();
                let scan = flatten_scan_values(&scan, true);

                let mz = self.index_converter.tof_to_mz(frame_id, &tof);
                let inv_mobility = self
                    .index_converter
                    .scan_to_inverse_mobility(frame_id, &scan);

                let ms_type_raw = self.raw_data_layout.frame_meta_data[frame_index].ms_ms_type;

                let ms_type = match ms_type_raw {
                    0 => MsType::Precursor,
                    8 => MsType::FragmentDda,
                    9 => MsType::FragmentDia,
                    _ => MsType::Unknown,
                };

                TimsFrame {
                    frame_id: frame_id as i32,
                    ms_type,
                    scan: scan.iter().map(|&x| x as i32).collect(),
                    tof: tof_i32,
                    ims_frame: ImsFrame::new(
                        self.raw_data_layout.frame_meta_data[frame_index].time,
                        inv_mobility,
                        mz,
                        intensity_dbl,
                    ),
                }
            }

            _ => {
                panic!("TimsCompressionType is not 1 or 2.")
            }
        }
    }

    fn get_raw_frame(&self, frame_id: u32) -> RawTimsFrame {
        let frame_index = (frame_id - 1) as usize;
        let offset = self.raw_data_layout.tims_offset_values[frame_index] as u64;

        // turns out, there can be empty frames in the data, check for that, if so, return an empty frame
        let num_peaks = self.raw_data_layout.frame_meta_data[frame_index].num_peaks;

        if num_peaks == 0 {
            return RawTimsFrame {
                frame_id: frame_id as i32,
                retention_time: self.raw_data_layout.frame_meta_data[(frame_id - 1) as usize].time,
                ms_type: MsType::Unknown,
                scan: Vec::new(),
                tof: Vec::new(),
                intensity: Vec::new(),
            };
        }

        let mut file_path = PathBuf::from(&self.raw_data_layout.raw_data_path);
        file_path.push("analysis.tdf_bin");
        let mut infile = File::open(&file_path).unwrap();

        infile.seek(SeekFrom::Start(offset)).unwrap();

        let mut bin_buffer = [0u8; 4];
        infile.read_exact(&mut bin_buffer).unwrap();
        let bin_size = Cursor::new(bin_buffer).read_i32::<LittleEndian>().unwrap();

        infile.read_exact(&mut bin_buffer).unwrap();

        match self.raw_data_layout.global_meta_data.tims_compression_type {
            _ if self.raw_data_layout.global_meta_data.tims_compression_type == 1 => {
                panic!("Decompression Type1 not implemented.");
            }

            // Extract from ZSTD compressed binary
            _ if self.raw_data_layout.global_meta_data.tims_compression_type == 2 => {
                let mut compressed_data = vec![0u8; bin_size as usize - 8];
                infile.read_exact(&mut compressed_data).unwrap();

                let decompressed_bytes = zstd_decompress(&compressed_data).unwrap();

                let (scan, tof, intensity) =
                    parse_decompressed_bruker_binary_data(&decompressed_bytes).unwrap();

                let ms_type_raw = self.raw_data_layout.frame_meta_data[frame_index].ms_ms_type;

                let ms_type = match ms_type_raw {
                    0 => MsType::Precursor,
                    8 => MsType::FragmentDda,
                    9 => MsType::FragmentDia,
                    _ => MsType::Unknown,
                };

                let frame = RawTimsFrame {
                    frame_id: frame_id as i32,
                    retention_time: self.raw_data_layout.frame_meta_data[(frame_id - 1) as usize]
                        .time,
                    ms_type,
                    scan,
                    tof,
                    intensity: intensity.iter().map(|&x| x as f64).collect(),
                };

                return frame;
            }

            // Error on unknown compression algorithm
            _ => {
                panic!("TimsCompressionType is not 1 or 2.")
            }
        }
    }

    fn get_slice(&self, frame_ids: Vec<u32>, _num_threads: usize) -> TimsSlice {
        let result: Vec<TimsFrame> = frame_ids.into_iter().map(|f| self.get_frame(f)).collect();

        TimsSlice { frames: result }
    }

    fn get_acquisition_mode(&self) -> AcquisitionMode {
        self.raw_data_layout.acquisition_mode.clone()
    }

    fn get_frame_count(&self) -> i32 {
        self.raw_data_layout.frame_meta_data.len() as i32
    }

    fn get_data_path(&self) -> &str {
        &self.raw_data_layout.raw_data_path
    }
}

pub struct TimsInMemoryLoader {
    pub raw_data_layout: TimsRawDataLayout,
    pub index_converter: TimsIndexConverter,
    compressed_data: Vec<u8>,
}

impl TimsData for TimsInMemoryLoader {
    fn get_frame(&self, frame_id: u32) -> TimsFrame {
        let raw_frame = self.get_raw_frame(frame_id);

        let raw_frame = match raw_frame.ms_type {
            MsType::FragmentDda => raw_frame.smooth(1).centroid(1),
            _ => raw_frame,
        };

        // if raw frame is empty, return an empty frame
        if raw_frame.scan.is_empty() {
            return TimsFrame::default();
        }

        let tof_i32 = raw_frame.tof.iter().map(|&x| x as i32).collect();
        let scan = flatten_scan_values(&raw_frame.scan, true);

        let mz = self.index_converter.tof_to_mz(frame_id, &raw_frame.tof);
        let inverse_mobility = self
            .index_converter
            .scan_to_inverse_mobility(frame_id, &scan);

        let ims_frame = ImsFrame::new(
            raw_frame.retention_time,
            inverse_mobility,
            mz,
            raw_frame.intensity,
        );

        TimsFrame {
            frame_id: frame_id as i32,
            ms_type: raw_frame.ms_type,
            scan: scan.iter().map(|&x| x as i32).collect(),
            tof: tof_i32,
            ims_frame,
        }
    }

    fn get_raw_frame(&self, frame_id: u32) -> RawTimsFrame {
        let frame_index = (frame_id - 1) as usize;
        let offset = self.raw_data_layout.tims_offset_values[frame_index] as usize;

        let bin_size_offset = offset + 4; // Assuming the size is stored immediately before the frame data
        let bin_size = Cursor::new(&self.compressed_data[offset..bin_size_offset])
            .read_i32::<LittleEndian>()
            .unwrap();

        let data_offset = bin_size_offset + 4; // Adjust based on actual structure
        let frame_data = &self.compressed_data[data_offset..data_offset + bin_size as usize - 8];

        let decompressed_bytes = zstd_decompress(&frame_data).unwrap();

        let (scan, tof, intensity) =
            parse_decompressed_bruker_binary_data(&decompressed_bytes).unwrap();

        let ms_type_raw = self.raw_data_layout.frame_meta_data[frame_index].ms_ms_type;

        let ms_type = match ms_type_raw {
            0 => MsType::Precursor,
            8 => MsType::FragmentDda,
            9 => MsType::FragmentDia,
            _ => MsType::Unknown,
        };

        let raw_frame = RawTimsFrame {
            frame_id: frame_id as i32,
            retention_time: self.raw_data_layout.frame_meta_data[(frame_id - 1) as usize].time,
            ms_type,
            scan,
            tof,
            intensity: intensity.iter().map(|&x| x as f64).collect(),
        };

        raw_frame
    }

    fn get_slice(&self, frame_ids: Vec<u32>, num_threads: usize) -> TimsSlice {
        let pool = ThreadPoolBuilder::new()
            .num_threads(num_threads)
            .build()
            .unwrap();
        let frames = pool.install(|| {
            frame_ids
                .par_iter()
                .map(|&frame_id| self.get_frame(frame_id))
                .collect()
        });

        TimsSlice { frames }
    }

    fn get_acquisition_mode(&self) -> AcquisitionMode {
        self.raw_data_layout.acquisition_mode.clone()
    }

    fn get_frame_count(&self) -> i32 {
        self.raw_data_layout.frame_meta_data.len() as i32
    }

    fn get_data_path(&self) -> &str {
        &self.raw_data_layout.raw_data_path
    }
}

pub enum TimsDataLoader {
    InMemory(TimsInMemoryLoader),
    Lazy(TimsLazyLoder),
}

/// Pick the m/z/mobility index converter, shared by the lazy and in-memory loaders.
///
/// - `use_bruker_sdk` → try the live Bruker SDK converter; if the SDK shared library can't be
///   loaded (missing / wrong arch / mismatched version) this now **falls back** (with a warning)
///   instead of panicking.
/// - Otherwise (or after that fallback): evaluate the calibration curves stored in the file
///   itself ([`BrukerFormulaConverter`]) — no SDK needed, 1/K0 machine-exact and m/z within
///   ~1 ppm of the SDK. Only if those tables cannot be read does it fall back to the
///   SDK-derived regression, and finally to the simple boundary model (which is ~5 Da off on
///   some datasets for m/z and thousands of ppm off on the mobility axis).
fn build_index_converter(
    bruker_lib_path: &str,
    data_path: &str,
    use_bruker_sdk: bool,
    scan_max_index: u32,
    im_lower: f64,
    im_upper: f64,
    tof_max_index: u32,
    mz_lower: f64,
    mz_upper: f64,
) -> TimsIndexConverter {
    if use_bruker_sdk {
        match BrukerLibTimsDataConverter::try_new(bruker_lib_path, data_path) {
            Ok(converter) => return TimsIndexConverter::BrukerLib(converter),
            Err(e) => eprintln!(
                "Warning: Bruker SDK requested but failed to load ({e}); \
                falling back to derived/simple m/z calibration."
            ),
        }
    }
    // SDK-free, and better than anything derived: evaluate the calibration curves the tdf
    // carries for this very run (per-frame temperature compensated, exact 1/K0).
    match BrukerFormulaConverter::from_d_folder(data_path, 1) {
        Ok(converter) => return TimsIndexConverter::BrukerFormula(converter),
        Err(e) => eprintln!(
            "Warning: could not build the SDK-free Bruker formula converter ({e}); \
            falling back to derived/simple m/z calibration."
        ),
    }

    // Derive an accurate calibration via the SDK if possible, else the simple boundary model.
    match derive_mz_calibration(bruker_lib_path, data_path, tof_max_index) {
        Some((intercept, slope)) => TimsIndexConverter::Calibrated(CalibratedIndexConverter::new(
            intercept,
            slope,
            im_lower,
            im_upper,
            scan_max_index,
        )),
        None => {
            eprintln!(
                "Warning: Could not derive m/z calibration from SDK. \
                Using simple boundary model which may have ~5 Da error on some datasets. \
                This typically happens on macOS where Bruker SDK is not available."
            );
            TimsIndexConverter::Simple(SimpleIndexConverter::from_boundaries(
                mz_lower,
                mz_upper,
                tof_max_index,
                im_lower,
                im_upper,
                scan_max_index,
            ))
        }
    }
}

impl TimsDataLoader {
    pub fn new_lazy(
        bruker_lib_path: &str,
        data_path: &str,
        use_bruker_sdk: bool,
        scan_max_index: u32,
        im_lower: f64,
        im_upper: f64,
        tof_max_index: u32,
        mz_lower: f64,
        mz_upper: f64,
    ) -> Self {
        let raw_data_layout = TimsRawDataLayout::new(data_path);
        let index_converter = build_index_converter(
            bruker_lib_path,
            data_path,
            use_bruker_sdk,
            scan_max_index,
            im_lower,
            im_upper,
            tof_max_index,
            mz_lower,
            mz_upper,
        );

        TimsDataLoader::Lazy(TimsLazyLoder {
            raw_data_layout,
            index_converter,
        })
    }

    pub fn new_in_memory(
        bruker_lib_path: &str,
        data_path: &str,
        use_bruker_sdk: bool,
        scan_max_index: u32,
        im_lower: f64,
        im_upper: f64,
        tof_max_index: u32,
        mz_lower: f64,
        mz_upper: f64,
    ) -> Self {
        let raw_data_layout = TimsRawDataLayout::new(data_path);
        let index_converter = build_index_converter(
            bruker_lib_path,
            data_path,
            use_bruker_sdk,
            scan_max_index,
            im_lower,
            im_upper,
            tof_max_index,
            mz_lower,
            mz_upper,
        );

        let mut file_path = PathBuf::from(data_path);
        file_path.push("analysis.tdf_bin");
        let mut infile = File::open(file_path).unwrap();
        let mut data = Vec::new();
        infile.read_to_end(&mut data).unwrap();

        TimsDataLoader::InMemory(TimsInMemoryLoader {
            raw_data_layout,
            index_converter,
            compressed_data: data,
        })
    }

    /// Create a lazy loader with pre-computed ion mobility calibration lookup table.
    ///
    /// This method enables accurate ion mobility calibration with fast parallel extraction.
    /// The im_lookup table should be pre-computed using the Bruker SDK.
    ///
    /// # Arguments
    /// * `data_path` - Path to the .d folder
    /// * `bruker_lib_path` - Path to the Bruker SDK shared library; used to
    ///   derive an accurate m/z calibration. Pass "NO_SDK" (or an empty
    ///   string) to skip and use the 2-point boundary m/z model.
    /// * `tof_max_index` - Maximum TOF index (from GlobalMetaData)
    /// * `mz_lower` - Minimum m/z value (from GlobalMetaData)
    /// * `mz_upper` - Maximum m/z value (from GlobalMetaData)
    /// * `im_lookup` - Pre-computed scan→1/K0 lookup table
    ///
    /// # Returns
    /// A new TimsDataLoader with LookupIndexConverter
    pub fn new_lazy_with_calibration(
        data_path: &str,
        bruker_lib_path: &str,
        tof_max_index: u32,
        mz_lower: f64,
        mz_upper: f64,
        im_lookup: Vec<f64>,
    ) -> Self {
        let raw_data_layout = TimsRawDataLayout::new(data_path);

        let index_converter = TimsIndexConverter::Lookup(build_lookup_converter(
            bruker_lib_path,
            data_path,
            tof_max_index,
            mz_lower,
            mz_upper,
            im_lookup,
        ));

        TimsDataLoader::Lazy(TimsLazyLoder {
            raw_data_layout,
            index_converter,
        })
    }

    /// Create an in-memory loader with pre-computed ion mobility calibration lookup table.
    ///
    /// This method enables accurate ion mobility calibration with fast parallel extraction.
    /// The im_lookup table should be pre-computed using the Bruker SDK.
    ///
    /// # Arguments
    /// * `data_path` - Path to the .d folder
    /// * `bruker_lib_path` - Path to the Bruker SDK shared library; used to
    ///   derive an accurate m/z calibration. Pass "NO_SDK" (or an empty
    ///   string) to skip and use the 2-point boundary m/z model.
    /// * `tof_max_index` - Maximum TOF index (from GlobalMetaData)
    /// * `mz_lower` - Minimum m/z value (from GlobalMetaData)
    /// * `mz_upper` - Maximum m/z value (from GlobalMetaData)
    /// * `im_lookup` - Pre-computed scan→1/K0 lookup table
    ///
    /// # Returns
    /// A new TimsDataLoader with LookupIndexConverter
    pub fn new_in_memory_with_calibration(
        data_path: &str,
        bruker_lib_path: &str,
        tof_max_index: u32,
        mz_lower: f64,
        mz_upper: f64,
        im_lookup: Vec<f64>,
    ) -> Self {
        let raw_data_layout = TimsRawDataLayout::new(data_path);

        let index_converter = TimsIndexConverter::Lookup(build_lookup_converter(
            bruker_lib_path,
            data_path,
            tof_max_index,
            mz_lower,
            mz_upper,
            im_lookup,
        ));

        let mut file_path = PathBuf::from(data_path);
        file_path.push("analysis.tdf_bin");
        let mut infile = File::open(file_path).unwrap();
        let mut data = Vec::new();
        infile.read_to_end(&mut data).unwrap();

        TimsDataLoader::InMemory(TimsInMemoryLoader {
            raw_data_layout,
            index_converter,
            compressed_data: data,
        })
    }

    /// Create a lazy loader using the exact SDK-free Bruker calibration formulas.
    ///
    /// Builds a [`BrukerFormulaConverter`] from the `MzCalibration` /
    /// `TimsCalibration` tables (frame `calibration_frame_id`, default 1 — the
    /// coefficients are near-constant per run). Needs no Bruker SDK at build or
    /// runtime; 1/K0 is machine-exact and m/z is bit-exact for MzCalibration
    /// ModelType 1 (~1 ppm for ModelType 2).
    pub fn new_lazy_with_bruker_formula(data_path: &str, calibration_frame_id: u32) -> Self {
        let raw_data_layout = TimsRawDataLayout::new(data_path);
        let index_converter = TimsIndexConverter::BrukerFormula(
            BrukerFormulaConverter::from_d_folder(data_path, calibration_frame_id).unwrap(),
        );
        TimsDataLoader::Lazy(TimsLazyLoder {
            raw_data_layout,
            index_converter,
        })
    }

    /// In-memory counterpart of [`Self::new_lazy_with_bruker_formula`].
    pub fn new_in_memory_with_bruker_formula(data_path: &str, calibration_frame_id: u32) -> Self {
        let raw_data_layout = TimsRawDataLayout::new(data_path);
        let index_converter = TimsIndexConverter::BrukerFormula(
            BrukerFormulaConverter::from_d_folder(data_path, calibration_frame_id).unwrap(),
        );

        let mut file_path = PathBuf::from(data_path);
        file_path.push("analysis.tdf_bin");
        let mut infile = File::open(file_path).unwrap();
        let mut data = Vec::new();
        infile.read_to_end(&mut data).unwrap();

        TimsDataLoader::InMemory(TimsInMemoryLoader {
            raw_data_layout,
            index_converter,
            compressed_data: data,
        })
    }

    /// Create a lazy loader with full calibration (both m/z and IM).
    ///
    /// This method uses regression-derived m/z calibration coefficients instead of
    /// the simple boundary model, providing more accurate m/z conversion.
    ///
    /// # Arguments
    /// * `data_path` - Path to the .d folder
    /// * `tof_intercept` - Intercept for sqrt(mz) = intercept + slope * tof
    /// * `tof_slope` - Slope for sqrt(mz) = intercept + slope * tof
    /// * `im_min` - Minimum 1/K0 value
    /// * `im_max` - Maximum 1/K0 value
    /// * `scan_max_index` - Maximum scan index
    pub fn new_lazy_with_mz_calibration(
        data_path: &str,
        tof_intercept: f64,
        tof_slope: f64,
        im_min: f64,
        im_max: f64,
        scan_max_index: u32,
    ) -> Self {
        let raw_data_layout = TimsRawDataLayout::new(data_path);

        let index_converter = TimsIndexConverter::Calibrated(CalibratedIndexConverter::new(
            tof_intercept,
            tof_slope,
            im_min,
            im_max,
            scan_max_index,
        ));

        TimsDataLoader::Lazy(TimsLazyLoder {
            raw_data_layout,
            index_converter,
        })
    }

    /// Create an in-memory loader with full calibration (both m/z and IM).
    ///
    /// This method uses regression-derived m/z calibration coefficients instead of
    /// the simple boundary model, providing more accurate m/z conversion.
    pub fn new_in_memory_with_mz_calibration(
        data_path: &str,
        tof_intercept: f64,
        tof_slope: f64,
        im_min: f64,
        im_max: f64,
        scan_max_index: u32,
    ) -> Self {
        let raw_data_layout = TimsRawDataLayout::new(data_path);

        let index_converter = TimsIndexConverter::Calibrated(CalibratedIndexConverter::new(
            tof_intercept,
            tof_slope,
            im_min,
            im_max,
            scan_max_index,
        ));

        let mut file_path = PathBuf::from(data_path);
        file_path.push("analysis.tdf_bin");
        let mut infile = File::open(file_path).unwrap();
        let mut data = Vec::new();
        infile.read_to_end(&mut data).unwrap();

        TimsDataLoader::InMemory(TimsInMemoryLoader {
            raw_data_layout,
            index_converter,
            compressed_data: data,
        })
    }

    pub fn get_index_converter(&self) -> &dyn IndexConverter {
        match self {
            TimsDataLoader::InMemory(loader) => &loader.index_converter,
            TimsDataLoader::Lazy(loader) => &loader.index_converter,
        }
    }

    /// Check if the Bruker SDK is being used for index conversion.
    /// The Bruker SDK is NOT thread-safe, so parallel operations that call
    /// the index converter must be disabled when using the SDK.
    pub fn uses_bruker_sdk(&self) -> bool {
        match self {
            TimsDataLoader::InMemory(loader) => matches!(&loader.index_converter, TimsIndexConverter::BrukerLib(_)),
            TimsDataLoader::Lazy(loader) => matches!(&loader.index_converter, TimsIndexConverter::BrukerLib(_)),
        }
    }
}

impl TimsData for TimsDataLoader {
    fn get_frame(&self, frame_id: u32) -> TimsFrame {
        match self {
            TimsDataLoader::InMemory(loader) => loader.get_frame(frame_id),
            TimsDataLoader::Lazy(loader) => loader.get_frame(frame_id),
        }
    }
    fn get_raw_frame(&self, frame_id: u32) -> RawTimsFrame {
        match self {
            TimsDataLoader::InMemory(loader) => loader.get_raw_frame(frame_id),
            TimsDataLoader::Lazy(loader) => loader.get_raw_frame(frame_id),
        }
    }

    fn get_slice(&self, frame_ids: Vec<u32>, num_threads: usize) -> TimsSlice {
        match self {
            TimsDataLoader::InMemory(loader) => loader.get_slice(frame_ids, num_threads),
            TimsDataLoader::Lazy(loader) => loader.get_slice(frame_ids, num_threads),
        }
    }

    fn get_acquisition_mode(&self) -> AcquisitionMode {
        match self {
            TimsDataLoader::InMemory(loader) => loader.get_acquisition_mode(),
            TimsDataLoader::Lazy(loader) => loader.get_acquisition_mode(),
        }
    }

    fn get_frame_count(&self) -> i32 {
        match self {
            TimsDataLoader::InMemory(loader) => loader.get_frame_count(),
            TimsDataLoader::Lazy(loader) => loader.get_frame_count(),
        }
    }

    fn get_data_path(&self) -> &str {
        match self {
            TimsDataLoader::InMemory(loader) => loader.get_data_path(),
            TimsDataLoader::Lazy(loader) => loader.get_data_path(),
        }
    }
}

pub struct SimpleIndexConverter {
    pub tof_intercept: f64,
    pub tof_slope: f64,
    pub scan_intercept: f64,
    pub scan_slope: f64,
}

impl SimpleIndexConverter {
    pub fn from_boundaries(
        mz_min: f64,
        mz_max: f64,
        tof_max_index: u32,
        im_min: f64,
        im_max: f64,
        scan_max_index: u32,
    ) -> Self {
        let tof_intercept: f64 = mz_min.sqrt();
        let tof_slope: f64 = (mz_max.sqrt() - tof_intercept) / tof_max_index as f64;

        let scan_intercept: f64 = im_max;
        let scan_slope: f64 = (im_min - scan_intercept) / scan_max_index as f64;
        Self {
            tof_intercept,
            tof_slope,
            scan_intercept,
            scan_slope,
        }
    }
}

impl IndexConverter for SimpleIndexConverter {
    fn tof_to_mz(&self, _frame_id: u32, _tof_values: &Vec<u32>) -> Vec<f64> {
        let mut mz_values: Vec<f64> = Vec::new();
        mz_values.resize(_tof_values.len(), 0.0);

        for (i, &val) in _tof_values.iter().enumerate() {
            mz_values[i] = (self.tof_intercept + self.tof_slope * val as f64).powi(2);
        }

        mz_values
    }

    fn mz_to_tof(&self, _frame_id: u32, _mz_values: &Vec<f64>) -> Vec<u32> {
        let mut tof_values: Vec<u32> = Vec::new();
        tof_values.resize(_mz_values.len(), 0);

        for (i, &val) in _mz_values.iter().enumerate() {
            tof_values[i] = ((val.sqrt() - self.tof_intercept) / self.tof_slope) as u32;
        }

        tof_values
    }

    fn scan_to_inverse_mobility(&self, _frame_id: u32, _scan_values: &Vec<u32>) -> Vec<f64> {
        let mut inv_mobility_values: Vec<f64> = Vec::new();
        inv_mobility_values.resize(_scan_values.len(), 0.0);

        for (i, &val) in _scan_values.iter().enumerate() {
            inv_mobility_values[i] = self.scan_intercept + self.scan_slope * val as f64;
        }

        inv_mobility_values
    }

    fn inverse_mobility_to_scan(
        &self,
        _frame_id: u32,
        _inverse_mobility_values: &Vec<f64>,
    ) -> Vec<u32> {
        let mut scan_values: Vec<u32> = Vec::new();
        scan_values.resize(_inverse_mobility_values.len(), 0);

        for (i, &val) in _inverse_mobility_values.iter().enumerate() {
            scan_values[i] = ((val - self.scan_intercept) / self.scan_slope) as u32;
        }

        scan_values
    }
}

/// M/z calibrated index converter using regression-derived coefficients.
///
/// This provides accurate TOF to m/z conversion without requiring the Bruker SDK
/// by using linear regression coefficients derived from known precursor m/z values.
///
/// The calibration formula is:
///   sqrt(mz) = tof_intercept + tof_slope * tof_index
///
/// This is similar to SimpleIndexConverter but uses externally-provided coefficients
/// (e.g., from regression on precursor data) rather than boundary-derived values.
pub struct CalibratedIndexConverter {
    pub tof_intercept: f64,
    pub tof_slope: f64,
    pub scan_intercept: f64,
    pub scan_slope: f64,
}

impl CalibratedIndexConverter {
    /// Create a new calibrated converter with regression-derived coefficients.
    ///
    /// # Arguments
    /// * `tof_intercept` - Intercept for sqrt(mz) = intercept + slope * tof
    /// * `tof_slope` - Slope for sqrt(mz) = intercept + slope * tof
    /// * `im_min` - Minimum 1/K0 value
    /// * `im_max` - Maximum 1/K0 value
    /// * `scan_max_index` - Maximum scan index
    pub fn new(
        tof_intercept: f64,
        tof_slope: f64,
        im_min: f64,
        im_max: f64,
        scan_max_index: u32,
    ) -> Self {
        let scan_intercept = im_max;
        let scan_slope = (im_min - scan_intercept) / scan_max_index as f64;
        Self {
            tof_intercept,
            tof_slope,
            scan_intercept,
            scan_slope,
        }
    }
}

impl IndexConverter for CalibratedIndexConverter {
    fn tof_to_mz(&self, _frame_id: u32, tof_values: &Vec<u32>) -> Vec<f64> {
        let mut mz_values: Vec<f64> = Vec::with_capacity(tof_values.len());

        for &tof_index in tof_values.iter() {
            // sqrt(mz) = tof_intercept + tof_slope * tof_index
            let sqrt_mz = self.tof_intercept + self.tof_slope * tof_index as f64;
            mz_values.push(sqrt_mz * sqrt_mz);
        }

        mz_values
    }

    fn mz_to_tof(&self, _frame_id: u32, mz_values: &Vec<f64>) -> Vec<u32> {
        let mut tof_values: Vec<u32> = Vec::with_capacity(mz_values.len());

        for &mz in mz_values.iter() {
            let sqrt_mz = mz.sqrt();
            // tof_index = (sqrt(mz) - tof_intercept) / tof_slope
            tof_values.push(((sqrt_mz - self.tof_intercept) / self.tof_slope) as u32);
        }

        tof_values
    }

    fn scan_to_inverse_mobility(&self, _frame_id: u32, scan_values: &Vec<u32>) -> Vec<f64> {
        let mut inv_mobility_values: Vec<f64> = Vec::with_capacity(scan_values.len());

        for &val in scan_values.iter() {
            inv_mobility_values.push(self.scan_intercept + self.scan_slope * val as f64);
        }

        inv_mobility_values
    }

    fn inverse_mobility_to_scan(
        &self,
        _frame_id: u32,
        inverse_mobility_values: &Vec<f64>,
    ) -> Vec<u32> {
        let mut scan_values: Vec<u32> = Vec::with_capacity(inverse_mobility_values.len());

        for &val in inverse_mobility_values.iter() {
            scan_values.push(((val - self.scan_intercept) / self.scan_slope) as u32);
        }

        scan_values
    }
}

/// Ion mobility index converter using pre-computed lookup table.
///
/// This converter uses a pre-computed scan→1/K0 lookup table extracted from the Bruker SDK.
/// It enables accurate ion mobility calibration with fast parallel extraction.
///
/// Background:
/// - The Bruker calibration formula is patented and proprietary
/// - Using the Bruker SDK gives accurate values but is slow (not thread-safe)
/// - Linear interpolation is fast but inaccurate
/// - This converter uses SDK-probed lookup for accuracy with O(1) thread-safe lookups
///
/// The lookup table is typically small (~8KB for 1000 scans) and constant across all frames.
pub struct LookupIndexConverter {
    // m/z conversion uses simple linear model (accurate enough for most purposes)
    pub tof_intercept: f64,
    pub tof_slope: f64,

    // Ion mobility: pre-computed lookup table from Bruker SDK
    // scan_index → 1/K0 value
    pub im_lookup: Vec<f64>,

    // Fallback for inverse conversion (1/K0 → scan)
    // We store the min/max for binary search bounds
    pub im_min: f64,
    pub im_max: f64,
}

impl LookupIndexConverter {
    /// Create a new LookupIndexConverter with pre-computed ion mobility lookup.
    ///
    /// # Arguments
    /// * `mz_min` - Minimum m/z value for TOF conversion
    /// * `mz_max` - Maximum m/z value for TOF conversion
    /// * `tof_max_index` - Maximum TOF index
    /// * `im_lookup` - Pre-computed scan→1/K0 lookup table from Bruker SDK
    ///
    /// # Returns
    /// A new LookupIndexConverter instance
    pub fn new(
        mz_min: f64,
        mz_max: f64,
        tof_max_index: u32,
        im_lookup: Vec<f64>,
    ) -> Self {
        let tof_intercept: f64 = mz_min.sqrt();
        let tof_slope: f64 = (mz_max.sqrt() - tof_intercept) / tof_max_index as f64;

        // Get IM bounds for inverse conversion
        let im_min = im_lookup.iter().cloned().fold(f64::INFINITY, f64::min);
        let im_max = im_lookup.iter().cloned().fold(f64::NEG_INFINITY, f64::max);

        Self {
            tof_intercept,
            tof_slope,
            im_lookup,
            im_min,
            im_max,
        }
    }

    /// Create a `LookupIndexConverter` with regression-derived m/z coefficients.
    ///
    /// Uses `sqrt(mz) = tof_intercept + tof_slope * tof` with coefficients fitted
    /// from the Bruker SDK (see `derive_mz_calibration`), instead of the 2-point
    /// boundary model in `new`. Preferred whenever the SDK is available.
    ///
    /// # Arguments
    /// * `tof_intercept` - Intercept of the sqrt(mz)-vs-tof regression
    /// * `tof_slope` - Slope of the sqrt(mz)-vs-tof regression
    /// * `im_lookup` - Pre-computed scan→1/K0 lookup table from Bruker SDK
    pub fn with_mz_fit(tof_intercept: f64, tof_slope: f64, im_lookup: Vec<f64>) -> Self {
        // Get IM bounds for inverse conversion
        let im_min = im_lookup.iter().cloned().fold(f64::INFINITY, f64::min);
        let im_max = im_lookup.iter().cloned().fold(f64::NEG_INFINITY, f64::max);

        Self {
            tof_intercept,
            tof_slope,
            im_lookup,
            im_min,
            im_max,
        }
    }
}

impl IndexConverter for LookupIndexConverter {
    fn tof_to_mz(&self, _frame_id: u32, tof_values: &Vec<u32>) -> Vec<f64> {
        let mut mz_values: Vec<f64> = Vec::new();
        mz_values.resize(tof_values.len(), 0.0);

        for (i, &val) in tof_values.iter().enumerate() {
            mz_values[i] = (self.tof_intercept + self.tof_slope * val as f64).powi(2);
        }

        mz_values
    }

    fn mz_to_tof(&self, _frame_id: u32, mz_values: &Vec<f64>) -> Vec<u32> {
        let mut tof_values: Vec<u32> = Vec::new();
        tof_values.resize(mz_values.len(), 0);

        for (i, &val) in mz_values.iter().enumerate() {
            tof_values[i] = ((val.sqrt() - self.tof_intercept) / self.tof_slope) as u32;
        }

        tof_values
    }

    fn scan_to_inverse_mobility(&self, _frame_id: u32, scan_values: &Vec<u32>) -> Vec<f64> {
        // Use the pre-computed lookup table for O(1) conversion
        scan_values
            .iter()
            .map(|&s| {
                self.im_lookup
                    .get(s as usize)
                    .copied()
                    .unwrap_or(f64::NAN)
            })
            .collect()
    }

    fn inverse_mobility_to_scan(
        &self,
        _frame_id: u32,
        inverse_mobility_values: &Vec<f64>,
    ) -> Vec<u32> {
        // Use binary search to find the closest scan index for each 1/K0 value
        // The lookup table is monotonically decreasing (higher scan = lower 1/K0)
        inverse_mobility_values
            .iter()
            .map(|&im| {
                if im.is_nan() || self.im_lookup.is_empty() {
                    return 0;
                }

                // Binary search for the closest value
                // Note: im_lookup is typically monotonically decreasing
                let mut best_scan = 0usize;
                let mut best_diff = f64::INFINITY;

                for (scan, &lookup_im) in self.im_lookup.iter().enumerate() {
                    let diff = (lookup_im - im).abs();
                    if diff < best_diff {
                        best_diff = diff;
                        best_scan = scan;
                    }
                }

                best_scan as u32
            })
            .collect()
    }
}