rustytorch/demos/ui/src-tauri/src/neuro_commands/service.rs

//! Neuro service for state management.

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

use super::types::*;

/// Loaded recording with metadata
pub struct LoadedRecording {
    pub handle: RecordingHandle,
    pub data: Vec<Vec<f64>>,
    pub sfreq: f64,
    pub channel_names: Vec<String>,
    pub channel_types: Vec<String>,
    pub events: Vec<EventDto>,
}

/// Loaded BIDS dataset storage
pub struct LoadedBidsDataset {
    pub handle: BidsDatasetHandle,
    pub dataset: rtx_neuro::io::bids::BidsDataset,
}

/// Epochs data storage
pub struct EpochsData {
    pub handle: EpochsHandle,
    /// [n_epochs x n_channels x n_times]
    pub data: Vec<Vec<Vec<f64>>>,
    pub times: Vec<f64>,
    pub channel_names: Vec<String>,
}

/// Forward model storage
pub struct ForwardModel {
    pub handle: ForwardHandle,
    /// Gain matrix [n_sensors x n_sources*3]
    pub gain: Vec<Vec<f64>>,
    /// Source positions [n_sources x 3]
    pub source_positions: Vec<[f64; 3]>,
    /// Sensor positions [n_sensors x 3]
    pub sensor_positions: Vec<[f64; 3]>,
}

/// Inverse operator storage
pub struct InverseOperator {
    pub handle: InverseHandle,
    /// Inverse kernel [n_sources x n_channels]
    pub kernel: Vec<Vec<f64>>,
    /// Source positions
    pub source_positions: Vec<[f64; 3]>,
}

/// SSP projector storage
pub struct SspProjectorData {
    pub handle: SspHandle,
    /// Projector vectors [n_projectors x n_channels]
    pub vectors: Vec<Vec<f64>>,
    /// Active flags
    pub active: Vec<bool>,
}

/// ICA model storage
#[derive(Clone)]
pub struct IcaModelData {
    pub handle: IcaHandle,
    /// Unmixing matrix [n_components x n_channels]
    pub unmixing: Vec<Vec<f64>>,
    /// Mixing matrix [n_channels x n_components]
    pub mixing: Vec<Vec<f64>>,
}

/// Loaded FreeSurfer subject storage
pub struct LoadedFreeSurferSubject {
    pub handle: FreeSurferHandleDto,
    pub subject: rtx_neuro_anatomy::FreeSurferSubject,
}

/// Handle to artifact detector
pub struct ArtifactDetectorHandle {
    pub config: ArtifactDetectorConfigDto,
}

/// Holds GNN model and metadata
pub struct GnnModelHolder {
    pub config: GnnModelConfigDto,
    pub n_parameters: usize,
}

/// Neuro service managing recordings and processing
pub struct NeuroService {
    pub(crate) recordings: HashMap<String, LoadedRecording>,
    pub(crate) epochs: HashMap<String, EpochsData>,
    pub(crate) forward_models: HashMap<String, ForwardModel>,
    pub(crate) inverse_operators: HashMap<String, InverseOperator>,
    pub(crate) ssp_projectors: HashMap<String, SspProjectorData>,
    pub(crate) ica_models: HashMap<String, IcaModelData>,
    pub(crate) bids_datasets: HashMap<String, LoadedBidsDataset>,
    pub(crate) freesurfer_subjects: HashMap<String, LoadedFreeSurferSubject>,
    pub(crate) active_recording: Option<String>,
    pub(crate) counter: usize,
}

impl NeuroService {
    /// Create a new neuro service
    pub fn new() -> Self {
        Self {
            recordings: HashMap::new(),
            epochs: HashMap::new(),
            forward_models: HashMap::new(),
            inverse_operators: HashMap::new(),
            ssp_projectors: HashMap::new(),
            ica_models: HashMap::new(),
            bids_datasets: HashMap::new(),
            freesurfer_subjects: HashMap::new(),
            active_recording: None,
            counter: 0,
        }
    }

    /// Generate a unique ID
    pub fn next_id(&mut self, prefix: &str) -> String {
        self.counter += 1;
        format!("{}_{}", prefix, self.counter)
    }

    /// Get service status
    pub fn status(&self) -> NeuroStatus {
        let memory_usage = self.recordings.values().map(|r| {
            r.data.iter().map(|ch| ch.len() * std::mem::size_of::<f64>()).sum::<usize>()
        }).sum::<usize>() + self.epochs.values().map(|e| {
            e.data.iter().map(|ep| {
                ep.iter().map(|ch| ch.len() * std::mem::size_of::<f64>()).sum::<usize>()
            }).sum::<usize>()
        }).sum::<usize>();

        NeuroStatus {
            n_recordings: self.recordings.len(),
            active_recording: self.active_recording.clone(),
            n_epochs_sets: self.epochs.len(),
            memory_usage,
        }
    }

    /// Reset the service, clearing all data
    pub fn reset(&mut self) {
        self.recordings.clear();
        self.epochs.clear();
        self.forward_models.clear();
        self.inverse_operators.clear();
        self.ssp_projectors.clear();
        self.ica_models.clear();
        self.bids_datasets.clear();
        self.freesurfer_subjects.clear();
        self.active_recording = None;
    }
}

impl Default for NeuroService {
    fn default() -> Self {
        Self::new()
    }
}

/// Type alias for the neuro state used by Tauri commands
pub type NeuroState = Arc<RwLock<NeuroService>>;

/// State for LSL service
pub type LslState = Arc<RwLock<rtx_neuro_lsl::LslService>>;

/// BCI pipeline state
pub type BciState = Arc<RwLock<Option<rtx_neuro_realtime::BciPipeline>>>;

/// State for artifact detector
pub type ArtifactDetectorState = Arc<RwLock<Option<ArtifactDetectorHandle>>>;

/// GNN state for managing models
pub type GnnModelState = Arc<RwLock<Option<GnnModelHolder>>>;

/// Separate state for the database to avoid async/sync conflicts
pub type DatabaseState = Arc<RwLock<Option<rtx_neuro_db::NeuroDatabase>>>;