rustytorch/crates/specialized/rtx-neuro-python/src/io.rs

//! I/O module - File format readers

use pyo3::prelude::*;
use numpy::PyArray2;
use rtx_neuro_io::{
    EdfReader, FifReader, CtfReader, BtiReader, KitReader, EgiReader,
    NeuroReader, IoError,
};
use std::path::PathBuf;

/// Create the io submodule
pub fn create_module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
    let m = PyModule::new(py, "io")?;

    m.add_class::<PyRawData>()?;
    m.add_function(wrap_pyfunction!(read_raw_edf, &m)?)?;
    m.add_function(wrap_pyfunction!(read_raw_fif, &m)?)?;
    m.add_function(wrap_pyfunction!(read_raw_ctf, &m)?)?;
    m.add_function(wrap_pyfunction!(read_raw_bti, &m)?)?;
    m.add_function(wrap_pyfunction!(read_raw_kit, &m)?)?;
    m.add_function(wrap_pyfunction!(read_raw_egi, &m)?)?;

    Ok(m)
}

/// Convert IoError to Python exception
fn io_err_to_py(e: IoError) -> PyErr {
    PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string())
}

/// Raw neuroimaging data container
#[pyclass]
pub struct PyRawData {
    /// Channel names
    channel_names: Vec<String>,
    /// Sampling frequency
    sfreq: f64,
    /// Number of channels
    n_channels: usize,
    /// Number of samples
    n_samples: usize,
    /// File path
    path: PathBuf,
    /// Reader type for lazy loading
    reader_type: ReaderType,
}

#[derive(Clone)]
enum ReaderType {
    Edf,
    Fif,
    Ctf,
    Bti,
    Kit,
    Egi,
}

#[pymethods]
impl PyRawData {
    /// Get sampling frequency in Hz
    #[getter]
    fn sfreq(&self) -> f64 {
        self.sfreq
    }

    /// Get number of channels
    #[getter]
    fn n_channels(&self) -> usize {
        self.n_channels
    }

    /// Get number of samples
    #[getter]
    fn n_samples(&self) -> usize {
        self.n_samples
    }

    /// Get channel names
    #[getter]
    fn ch_names(&self) -> Vec<String> {
        self.channel_names.clone()
    }

    /// Get duration in seconds
    #[getter]
    fn duration(&self) -> f64 {
        self.n_samples as f64 / self.sfreq
    }

    /// Get file path
    #[getter]
    fn path(&self) -> String {
        self.path.to_string_lossy().to_string()
    }

    /// Get data as numpy array
    ///
    /// # Arguments
    /// * `tmin` - Start time in seconds (default: 0.0)
    /// * `tmax` - End time in seconds (default: None = end of recording)
    ///
    /// # Returns
    /// 2D numpy array [n_channels x n_samples]
    fn get_data<'py>(
        &self,
        py: Python<'py>,
        tmin: Option<f64>,
        tmax: Option<f64>,
    ) -> PyResult<Bound<'py, PyArray2<f64>>> {
        let tmin = tmin.unwrap_or(0.0);
        let tmax = tmax.unwrap_or(self.duration());

        let data = match &self.reader_type {
            ReaderType::Edf => {
                let mut reader = EdfReader::open(&self.path).map_err(io_err_to_py)?;
                reader.read_raw_data(tmin, tmax).map_err(io_err_to_py)?
            }
            ReaderType::Fif => {
                let mut reader = FifReader::open(&self.path).map_err(io_err_to_py)?;
                reader.read_raw_data(tmin, tmax).map_err(io_err_to_py)?
            }
            ReaderType::Ctf => {
                let mut reader = CtfReader::open(&self.path).map_err(io_err_to_py)?;
                reader.read_raw_data(tmin, tmax).map_err(io_err_to_py)?
            }
            ReaderType::Bti => {
                let mut reader = BtiReader::open(&self.path).map_err(io_err_to_py)?;
                reader.read_raw_data(tmin, tmax).map_err(io_err_to_py)?
            }
            ReaderType::Kit => {
                let mut reader = KitReader::open(&self.path).map_err(io_err_to_py)?;
                reader.read_raw_data(tmin, tmax).map_err(io_err_to_py)?
            }
            ReaderType::Egi => {
                let mut reader = EgiReader::open(&self.path).map_err(io_err_to_py)?;
                reader.read_raw_data(tmin, tmax).map_err(io_err_to_py)?
            }
        };

        // Calculate actual dimensions
        let n_samples = data.len() / self.n_channels;

        // Reshape from flat to 2D [n_channels x n_samples]
        // Convert channel-major flat array to Vec<Vec<f64>>
        let mut rows = Vec::with_capacity(self.n_channels);
        for ch in 0..self.n_channels {
            let start = ch * n_samples;
            let end = start + n_samples;
            rows.push(data[start..end].to_vec());
        }

        PyArray2::from_vec2(py, &rows)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
    }

    /// Pick specific channels by name
    fn pick_channels(&self, ch_names: Vec<String>) -> PyResult<Vec<usize>> {
        let mut indices = Vec::new();
        for name in &ch_names {
            if let Some(idx) = self.channel_names.iter().position(|n| n == name) {
                indices.push(idx);
            } else {
                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                    format!("Channel '{}' not found", name)
                ));
            }
        }
        Ok(indices)
    }

    /// String representation
    fn __repr__(&self) -> String {
        format!(
            "RawData({} channels, {} samples, {:.1} Hz, {:.1}s)",
            self.n_channels, self.n_samples, self.sfreq, self.duration()
        )
    }
}

/// Read EDF/EDF+ file
///
/// # Arguments
/// * `path` - Path to the EDF file
///
/// # Returns
/// PyRawData object
#[pyfunction]
fn read_raw_edf(path: &str) -> PyResult<PyRawData> {
    let path_buf = PathBuf::from(path);
    let mut reader = EdfReader::open(&path_buf).map_err(io_err_to_py)?;
    reader.read_header().map_err(io_err_to_py)?;

    Ok(PyRawData {
        channel_names: reader.channel_names(),
        sfreq: reader.sfreq(),
        n_channels: reader.n_channels(),
        n_samples: reader.n_samples(),
        path: path_buf,
        reader_type: ReaderType::Edf,
    })
}

/// Read Elekta/Neuromag FIF file
///
/// # Arguments
/// * `path` - Path to the FIF file
///
/// # Returns
/// PyRawData object
#[pyfunction]
fn read_raw_fif(path: &str) -> PyResult<PyRawData> {
    let path_buf = PathBuf::from(path);
    let mut reader = FifReader::open(&path_buf).map_err(io_err_to_py)?;
    reader.read_header().map_err(io_err_to_py)?;

    Ok(PyRawData {
        channel_names: reader.channel_names(),
        sfreq: reader.sfreq(),
        n_channels: reader.n_channels(),
        n_samples: reader.n_samples(),
        path: path_buf,
        reader_type: ReaderType::Fif,
    })
}

/// Read CTF MEG dataset (.ds directory)
///
/// # Arguments
/// * `path` - Path to the .ds directory
///
/// # Returns
/// PyRawData object
#[pyfunction]
fn read_raw_ctf(path: &str) -> PyResult<PyRawData> {
    let path_buf = PathBuf::from(path);
    let mut reader = CtfReader::open(&path_buf).map_err(io_err_to_py)?;
    reader.read_header().map_err(io_err_to_py)?;

    Ok(PyRawData {
        channel_names: reader.channel_names(),
        sfreq: reader.sfreq(),
        n_channels: reader.n_channels(),
        n_samples: reader.n_samples(),
        path: path_buf,
        reader_type: ReaderType::Ctf,
    })
}

/// Read 4D-Neuroimaging/BTi MEG data
///
/// # Arguments
/// * `path` - Path to the BTi data directory or PDF file
///
/// # Returns
/// PyRawData object
#[pyfunction]
fn read_raw_bti(path: &str) -> PyResult<PyRawData> {
    let path_buf = PathBuf::from(path);
    let mut reader = BtiReader::open(&path_buf).map_err(io_err_to_py)?;
    reader.read_header().map_err(io_err_to_py)?;

    Ok(PyRawData {
        channel_names: reader.channel_names(),
        sfreq: reader.sfreq(),
        n_channels: reader.n_channels(),
        n_samples: reader.n_samples(),
        path: path_buf,
        reader_type: ReaderType::Bti,
    })
}

/// Read Yokogawa/KIT MEG data (.con, .sqd)
///
/// # Arguments
/// * `path` - Path to the KIT file
///
/// # Returns
/// PyRawData object
#[pyfunction]
fn read_raw_kit(path: &str) -> PyResult<PyRawData> {
    let path_buf = PathBuf::from(path);
    let mut reader = KitReader::open(&path_buf).map_err(io_err_to_py)?;
    reader.read_header().map_err(io_err_to_py)?;

    Ok(PyRawData {
        channel_names: reader.channel_names(),
        sfreq: reader.sfreq(),
        n_channels: reader.n_channels(),
        n_samples: reader.n_samples(),
        path: path_buf,
        reader_type: ReaderType::Kit,
    })
}

/// Read EGI data (.raw, .mff)
///
/// # Arguments
/// * `path` - Path to the EGI file or .mff directory
///
/// # Returns
/// PyRawData object
#[pyfunction]
fn read_raw_egi(path: &str) -> PyResult<PyRawData> {
    let path_buf = PathBuf::from(path);
    let mut reader = EgiReader::open(&path_buf).map_err(io_err_to_py)?;
    reader.read_header().map_err(io_err_to_py)?;

    Ok(PyRawData {
        channel_names: reader.channel_names(),
        sfreq: reader.sfreq(),
        n_channels: reader.n_channels(),
        n_samples: reader.n_samples(),
        path: path_buf,
        reader_type: ReaderType::Egi,
    })
}