src/etl/extract.py

"""
src/etl/extract.py
────────────────────────────────────────────────────────────────
Data extraction layer — the "E" in ETL.

Responsibilities
────────────────
  - Load MTSamples clinical notes from CSV
  - Load the ICD-10 reference table
  - Detect whether MIMIC-III data is available (optional)
  - Cache raw downloads so re-runs are fast

Design
──────
  Extractors are intentionally dumb: they load data and add a
  ``_source`` provenance column, but they do not clean, filter,
  or reshape anything.  All of that happens in transform.py.

  Every public function returns a DataFrame or raises a clear
  exception — never returns None.  Callers can always rely on
  the return type.

  MTSamples is freely available on Kaggle:
    https://www.kaggle.com/datasets/tboyle10/medicaltranscriptions
  Download ``mtsamples.csv`` and place it in ``data/raw/``.

  The ICD-10 code file (icd10_codes.csv) is included in the
  project under ``data/raw/`` — it is a static reference table
  that rarely changes.
────────────────────────────────────────────────────────────────
"""

from __future__ import annotations

from pathlib import Path

import pandas as pd

from src.utils.config import Paths
from src.utils.logger import get_logger

logger = get_logger(__name__)


# ── Column name aliases ────────────────────────────────────────────
# MTSamples columns vary slightly across Kaggle versions.
# We normalise to these names during extraction.
_MTSAMPLES_COLUMN_MAP: dict[str, str] = {
    "description":        "description",
    "medical_specialty":  "specialty",
    "sample_name":        "sample_name",
    "transcription":      "transcription",
    "keywords":           "keywords",
}

_ICD10_COLUMN_MAP: dict[str, str] = {
    "code":        "icd10_code",
    "description": "description",
    "category":    "category",
}


# ── MTSamples ─────────────────────────────────────────────────────

def load_mtsamples(
    path: Path | None = None,
    force_reload: bool = False,
) -> pd.DataFrame:
    """Load the MTSamples clinical notes dataset from CSV.

    MTSamples contains 4,999 de-identified medical transcriptions
    across 40 clinical specialties.  The raw CSV is expected at
    ``data/raw/mtsamples.csv``.

    Args:
        path: Override the default file location.  Useful in tests.
        force_reload: If True, bypass the cached version and re-read
            from disk.  Defaults to False.

    Returns:
        DataFrame with columns::

            description     : short note title
            specialty       : medical specialty (40 categories)
            sample_name     : note type (e.g. "Discharge Summary")
            transcription   : full clinical note text
            keywords        : comma-separated clinical keywords
            _source         : "mtsamples" (provenance marker)

    Raises:
        FileNotFoundError: If the CSV does not exist at the
            expected path.  The error message includes the
            download URL.

    Example::

        df = load_mtsamples()
        print(df.shape)          # (4999, 6)
        print(df["specialty"].nunique())  # 40
    """
    csv_path = path or Paths.mtsamples_csv

    if not csv_path.exists():
        raise FileNotFoundError(
            f"MTSamples CSV not found at: {csv_path}\n"
            "Download it from Kaggle:\n"
            "  https://www.kaggle.com/datasets/tboyle10/medicaltranscriptions\n"
            f"Then save it to: {Paths.raw}/"
        )

    logger.info("Loading MTSamples from %s", csv_path)
    df = pd.read_csv(csv_path, low_memory=False)

    # Rename columns to our standard names, ignoring any extras
    rename_map = {
        col: _MTSAMPLES_COLUMN_MAP[col]
        for col in df.columns
        if col in _MTSAMPLES_COLUMN_MAP
    }
    df = df.rename(columns=rename_map)

    # Keep only columns we use; add the rest as-is if not in our map
    expected = list(_MTSAMPLES_COLUMN_MAP.values())
    available = [c for c in expected if c in df.columns]
    df = df[available].copy()

    df["_source"] = "mtsamples"

    logger.info(
        "MTSamples loaded: %d notes, %d specialties",
        len(df),
        df["specialty"].nunique() if "specialty" in df.columns else 0,
    )
    return df


# ── ICD-10 reference table ─────────────────────────────────────────

def load_icd10_codes(path: Path | None = None) -> pd.DataFrame:
    """Load the ICD-10 diagnostic code reference table.

    The ICD-10 CSV is included in the repository under
    ``data/raw/icd10_codes.csv``.  It contains all billable
    ICD-10-CM codes with descriptions.

    Args:
        path: Override the default file location.  Useful in tests.

    Returns:
        DataFrame with columns::

            icd10_code  : ICD-10-CM code (e.g. "I10", "J18.9")
            description : Human-readable description
            category    : Broad category (e.g. "Circulatory System")
            _source     : "icd10_cms" (provenance marker)

    Raises:
        FileNotFoundError: If the CSV is missing.

    Example::

        icd = load_icd10_codes()
        icd[icd["icd10_code"] == "I10"]
        # icd10_code | description            | category
        # I10        | Essential hypertension | Circulatory System
    """
    csv_path = path or Paths.icd10_csv

    if not csv_path.exists():
        raise FileNotFoundError(
            f"ICD-10 codes file not found at: {csv_path}\n"
            "This file should be included in the repository.\n"
            "Check that you have the full project download."
        )

    logger.info("Loading ICD-10 reference table from %s", csv_path)
    df = pd.read_csv(csv_path, dtype=str, low_memory=False)

    rename_map = {
        col: _ICD10_COLUMN_MAP[col]
        for col in df.columns
        if col in _ICD10_COLUMN_MAP
    }
    df = df.rename(columns=rename_map)

    # Ensure the code column is clean
    if "icd10_code" in df.columns:
        df["icd10_code"] = df["icd10_code"].str.strip().str.upper()

    df["_source"] = "icd10_cms"

    logger.info("ICD-10 table loaded: %d codes", len(df))
    return df


# ── MIMIC-III (optional) ───────────────────────────────────────────

def load_mimic_notes(
    path: Path | None = None,
) -> pd.DataFrame | None:
    """Load MIMIC-III discharge summaries if available.

    MIMIC-III requires credentialed access via PhysioNet:
      https://physionet.org/content/mimiciii/

    This function gracefully returns None if the data is not
    present rather than raising an error — MIMIC-III is optional.
    The pipeline uses MTSamples when MIMIC-III is absent.

    Args:
        path: Path to ``NOTEEVENTS.csv`` from MIMIC-III.
            Defaults to ``data/raw/mimic_notes.csv``.

    Returns:
        DataFrame with columns ``transcription``, ``specialty``,
        ``_source`` if the file exists; None otherwise.

    Example::

        notes = load_mimic_notes()
        if notes is None:
            print("MIMIC-III not available, using MTSamples")
    """
    default_path = Paths.raw / "mimic_notes.csv"
    csv_path     = path or default_path

    if not csv_path.exists():
        logger.info(
            "MIMIC-III notes not found at %s — using MTSamples only. "
            "See https://physionet.org/content/mimiciii/ for access.",
            csv_path,
        )
        return None

    logger.info("Loading MIMIC-III notes from %s", csv_path)
    df = pd.read_csv(
        csv_path,
        usecols=["TEXT", "CATEGORY"],
        low_memory=False,
    )

    df = df.rename(columns={"TEXT": "transcription", "CATEGORY": "specialty"})
    df["description"] = "MIMIC-III discharge note"
    df["_source"]     = "mimic3"

    # Keep only discharge summaries for comparability with MTSamples
    df = df[df["specialty"].str.lower() == "discharge summary"].copy()

    logger.info("MIMIC-III notes loaded: %d discharge summaries", len(df))
    return df


# ── Auto-detect best available source ─────────────────────────────

def load_clinical_notes() -> pd.DataFrame:
    """Load clinical notes from the best available source.

    Tries sources in order of preference:
      1. MTSamples (primary — free, always available once downloaded)
      2. MIMIC-III (optional — richer, requires credentialed access)

    If both are available, MTSamples is used by default.
    Set ``DATA_SOURCE=mimic`` in your ``.env`` to prefer MIMIC-III.

    Returns:
        DataFrame with at minimum the columns ``transcription``,
        ``specialty``, and ``_source``.

    Raises:
        FileNotFoundError: If no data source is available.
    """
    import os
    preferred = os.getenv("DATA_SOURCE", "mtsamples").lower()

    if preferred == "mimic":
        mimic_df = load_mimic_notes()
        if mimic_df is not None:
            return mimic_df
        logger.warning(
            "DATA_SOURCE=mimic but MIMIC-III not found. "
            "Falling back to MTSamples."
        )

    return load_mtsamples()