rustytorch/crates/specialized/rtx-neuro-db/src/schema.rs

//! Database schema definition and migration

use rusqlite::Connection;

use crate::error::Result;

/// Current schema version
pub const SCHEMA_VERSION: i32 = 1;

/// Create the database schema
pub fn create_schema(conn: &Connection) -> Result<()> {
    conn.execute_batch(SCHEMA_SQL)?;

    // Set schema version
    conn.execute(
        "INSERT OR REPLACE INTO schema_version (id, version) VALUES (1, ?1)",
        [SCHEMA_VERSION],
    )?;

    Ok(())
}

/// Check if schema needs migration
pub fn check_schema_version(conn: &Connection) -> Result<Option<i32>> {
    // Check if schema_version table exists
    let exists: bool = conn.query_row(
        "SELECT COUNT(*) > 0 FROM sqlite_master WHERE type='table' AND name='schema_version'",
        [],
        |row| row.get(0),
    )?;

    if !exists {
        return Ok(None);
    }

    let version: i32 = conn.query_row(
        "SELECT version FROM schema_version WHERE id = 1",
        [],
        |row| row.get(0),
    )?;

    Ok(Some(version))
}

/// Schema SQL for all tables
const SCHEMA_SQL: &str = r"
-- Schema version tracking
CREATE TABLE IF NOT EXISTS schema_version (
    id INTEGER PRIMARY KEY,
    version INTEGER NOT NULL
);

-- Protocols (top-level container, like a Brainstorm protocol)
CREATE TABLE IF NOT EXISTS protocols (
    id TEXT PRIMARY KEY,
    name TEXT NOT NULL,
    path TEXT NOT NULL,
    description TEXT DEFAULT '',
    created_at TEXT NOT NULL,
    modified_at TEXT NOT NULL,
    settings_json TEXT NOT NULL DEFAULT '{}'
);

-- Subjects within a protocol
CREATE TABLE IF NOT EXISTS subjects (
    id TEXT PRIMARY KEY,
    protocol_id TEXT NOT NULL REFERENCES protocols(id) ON DELETE CASCADE,
    label TEXT NOT NULL,
    anatomy_path TEXT,
    demographics_json TEXT DEFAULT '{}',
    created_at TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_subjects_protocol ON subjects(protocol_id);

-- Conditions (experimental conditions for grouping)
CREATE TABLE IF NOT EXISTS conditions (
    id TEXT PRIMARY KEY,
    protocol_id TEXT NOT NULL REFERENCES protocols(id) ON DELETE CASCADE,
    name TEXT NOT NULL,
    color TEXT DEFAULT '#4A90D9',
    description TEXT DEFAULT ''
);

CREATE INDEX IF NOT EXISTS idx_conditions_protocol ON conditions(protocol_id);

-- Recordings (loaded data files)
CREATE TABLE IF NOT EXISTS recordings (
    id TEXT PRIMARY KEY,
    subject_id TEXT REFERENCES subjects(id) ON DELETE SET NULL,
    condition_id TEXT REFERENCES conditions(id) ON DELETE SET NULL,
    handle_id TEXT NOT NULL,
    file_path TEXT NOT NULL,
    format TEXT NOT NULL,
    n_channels INTEGER NOT NULL,
    n_samples INTEGER NOT NULL,
    sfreq REAL NOT NULL,
    duration REAL NOT NULL,
    bad_channels_json TEXT DEFAULT '[]',
    annotations_json TEXT DEFAULT '[]',
    created_at TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_recordings_subject ON recordings(subject_id);
CREATE INDEX IF NOT EXISTS idx_recordings_condition ON recordings(condition_id);

-- Epochs sets
CREATE TABLE IF NOT EXISTS epochs (
    id TEXT PRIMARY KEY,
    recording_id TEXT NOT NULL REFERENCES recordings(id) ON DELETE CASCADE,
    handle_id TEXT NOT NULL,
    event_ids_json TEXT NOT NULL,
    tmin REAL NOT NULL,
    tmax REAL NOT NULL,
    baseline_json TEXT,
    n_epochs INTEGER NOT NULL,
    n_channels INTEGER NOT NULL,
    n_times INTEGER NOT NULL,
    reject_criteria_json TEXT DEFAULT '{}',
    created_at TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_epochs_recording ON epochs(recording_id);

-- Forward models
CREATE TABLE IF NOT EXISTS forward_models (
    id TEXT PRIMARY KEY,
    recording_id TEXT NOT NULL REFERENCES recordings(id) ON DELETE CASCADE,
    handle_id TEXT NOT NULL,
    model_type TEXT NOT NULL,
    n_sensors INTEGER NOT NULL,
    n_sources INTEGER NOT NULL,
    config_json TEXT NOT NULL,
    gain_matrix_path TEXT,
    created_at TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_forward_recording ON forward_models(recording_id);

-- Inverse operators
CREATE TABLE IF NOT EXISTS inverse_operators (
    id TEXT PRIMARY KEY,
    forward_id TEXT NOT NULL REFERENCES forward_models(id) ON DELETE CASCADE,
    handle_id TEXT NOT NULL,
    method TEXT NOT NULL,
    lambda2 REAL NOT NULL,
    n_sources INTEGER NOT NULL,
    n_channels INTEGER NOT NULL,
    kernel_path TEXT,
    created_at TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_inverse_forward ON inverse_operators(forward_id);

-- SSP projectors
CREATE TABLE IF NOT EXISTS ssp_projectors (
    id TEXT PRIMARY KEY,
    recording_id TEXT NOT NULL REFERENCES recordings(id) ON DELETE CASCADE,
    handle_id TEXT NOT NULL,
    n_components INTEGER NOT NULL,
    artifact_type TEXT,
    explained_var_json TEXT DEFAULT '[]',
    projector_path TEXT,
    created_at TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_ssp_recording ON ssp_projectors(recording_id);

-- ICA models
CREATE TABLE IF NOT EXISTS ica_models (
    id TEXT PRIMARY KEY,
    recording_id TEXT NOT NULL REFERENCES recordings(id) ON DELETE CASCADE,
    handle_id TEXT NOT NULL,
    n_components INTEGER NOT NULL,
    method TEXT NOT NULL,
    excluded_json TEXT DEFAULT '[]',
    mixing_matrix_path TEXT,
    unmixing_matrix_path TEXT,
    created_at TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_ica_recording ON ica_models(recording_id);
";

#[cfg(test)]
mod tests {
    use super::*;
    use rusqlite::Connection;

    #[test]
    fn test_create_schema() {
        let conn = Connection::open_in_memory().unwrap();
        create_schema(&conn).unwrap();

        // Verify tables exist
        let count: i32 = conn
            .query_row(
                "SELECT COUNT(*) FROM sqlite_master WHERE type='table'",
                [],
                |row| row.get(0),
            )
            .unwrap();

        assert!(count >= 10); // At least 10 tables
    }

    #[test]
    fn test_schema_version() {
        let conn = Connection::open_in_memory().unwrap();

        // Before schema creation
        assert_eq!(check_schema_version(&conn).unwrap(), None);

        // After schema creation
        create_schema(&conn).unwrap();
        assert_eq!(check_schema_version(&conn).unwrap(), Some(SCHEMA_VERSION));
    }
}