load_data.py

"""
Part 1: Data Management

Designing the relational database schema with SQLite with data from cell-count.csv. 
Runs with: python load_data.py
 
Schema overview:
  projects: one row per project
  subjects: one row per patient (with clinical metadata)
  samples: one row per timepoint per patient
  cell_counts: long-format cell type measurements per sample
"""


import sqlite3
import csv
import os

DB_PATH = "cell_counts.db"
CSV_PATH = "cell-count.csv"

# Cell type columns present in the CSV
CELL_TYPES = ["b_cell", "cd8_t_cell", "cd4_t_cell", "nk_cell", "monocyte"]


def init_db(conn: sqlite3.Connection) -> None:
    """Creates all the tables."""
    cursor = conn.cursor()

    cursor.executescript("""
        PRAGMA foreign_keys = ON;

        CREATE TABLE IF NOT EXISTS projects (
            project_id  TEXT PRIMARY KEY
        );

        CREATE TABLE IF NOT EXISTS subjects (
            subject_id  TEXT PRIMARY KEY,
            project_id  TEXT    NOT NULL REFERENCES projects(project_id),
            condition   TEXT    NOT NULL,
            age         INTEGER NOT NULL,
            sex         TEXT    NOT NULL,
            treatment   TEXT    NOT NULL,
            response    TEXT                -- NULL for healthy/untreated subjects
        );

        CREATE TABLE IF NOT EXISTS samples (
            sample_id                   TEXT    PRIMARY KEY,
            subject_id                  TEXT    NOT NULL REFERENCES subjects(subject_id),
            sample_type                 TEXT    NOT NULL,
            time_from_treatment_start   INTEGER NOT NULL
        );

        CREATE TABLE IF NOT EXISTS cell_counts (
            count_id    INTEGER PRIMARY KEY AUTOINCREMENT,
            sample_id   TEXT    NOT NULL REFERENCES samples(sample_id),
            cell_type   TEXT    NOT NULL,
            count       INTEGER NOT NULL
        );

        -- Indexes to speed up common analytical queries
        CREATE INDEX IF NOT EXISTS idx_subjects_project
            ON subjects(project_id);
        CREATE INDEX IF NOT EXISTS idx_samples_subject
            ON samples(subject_id);
        CREATE INDEX IF NOT EXISTS idx_cell_counts_sample
            ON cell_counts(sample_id);
        CREATE INDEX IF NOT EXISTS idx_cell_counts_type
            ON cell_counts(cell_type);
    """)

    conn.commit()


def load_csv(conn: sqlite3.Connection, csv_path: str) -> None:
    """
    Reads cell-count.csv and inserts rows into the tables for the schema.
    Uses sets to track already-inserted projects and subjects so we
    only insert each unique entity once.
    """
    cursor = conn.cursor()

    seen_projects: set[str] = set()
    seen_subjects: set[str] = set()

    with open(csv_path, newline="") as f:
        reader = csv.DictReader(f)

        # Batch inserts for performance
        sample_rows = []
        cell_count_rows = []

        for row in reader:
            project_id = row["project"]
            subject_id = row["subject"]
            sample_id  = row["sample"]

            # projects
            if project_id not in seen_projects:
                cursor.execute(
                    "INSERT OR IGNORE INTO projects (project_id) VALUES (?)",
                    (project_id,)
                )
                seen_projects.add(project_id)

            # subjects 
            if subject_id not in seen_subjects:
                # response is empty string for healthy subjects, so we'll store as NULL
                response = row["response"] if row["response"] != "" else None
                cursor.execute(
                    """
                    INSERT OR IGNORE INTO subjects
                        (subject_id, project_id, condition, age, sex, treatment, response)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        subject_id,
                        project_id,
                        row["condition"],
                        int(row["age"]),
                        row["sex"],
                        row["treatment"],
                        response,
                    )
                )
                seen_subjects.add(subject_id)

            # samples 
            sample_rows.append((
                sample_id,
                subject_id,
                row["sample_type"],
                int(row["time_from_treatment_start"]),
            ))

            # cell_counts become long-format with one row per cell type per sample for easy appending and querying
            for cell_type in CELL_TYPES:
                cell_count_rows.append((
                    sample_id,
                    cell_type,
                    int(row[cell_type]),
                ))

    cursor.executemany(
        """
        INSERT OR IGNORE INTO samples
            (sample_id, subject_id, sample_type, time_from_treatment_start)
        VALUES (?, ?, ?, ?)
        """,
        sample_rows,
    )

    cursor.executemany(
        "INSERT INTO cell_counts (sample_id, cell_type, count) VALUES (?, ?, ?)",
        cell_count_rows,
    )

    conn.commit()


def print_summary(conn: sqlite3.Connection) -> None:
    """Print row counts for each table."""
    cursor = conn.cursor()
    tables = ["projects", "subjects", "samples", "cell_counts"]
    print("\nDatabase summary:")
    for table in tables:
        cursor.execute(f"SELECT COUNT(*) FROM {table}")
        count = cursor.fetchone()[0]
        print(f"  {table:<15} {count:>7,} rows")
    print()


def main() -> None:
    if not os.path.exists(CSV_PATH):
        raise FileNotFoundError(
            f"Could not find '{CSV_PATH}'. "
            "Verify that cell-count.csv is in the same directory as load_data.py."
        )

    # Remove populated database if present so we can start fresh each time and not risk duplicate entries
    if os.path.exists(DB_PATH):
        os.remove(DB_PATH)
        print(f"Removed existing {DB_PATH}")

    print(f"Initializing database: {DB_PATH}")
    conn = sqlite3.connect(DB_PATH)

    try:
        init_db(conn)
        print(f"Loading data from: {CSV_PATH}")
        load_csv(conn, CSV_PATH)
        print_summary(conn)
    finally:
        conn.close()


if __name__ == "__main__":
    main()