data_pipeline/build.py

"""Build processed dataset from data/raw/.

Run as:
    python -m data_pipeline.build

Reads the cBioPortal clinical and expression files, joins sample <-> patient,
derives an MSI-H / MSS label, normalises confounders, writes:

    data/processed/clinical.parquet     # one row per sample
    data/processed/expression.parquet   # genes x samples

Prints a provenance report naming every file + column it used, the MSI threshold,
and cohort sizes (total / COAD / READ / usable / READ-only-usable).
"""

from __future__ import annotations

import re
import sys
from pathlib import Path

import pandas as pd

from data_pipeline import schema


# --- column-level helpers ---------------------------------------------------

def label_msi(score: float) -> str:
    """Map an MSIsensor score to MSI-H / MSS / MSI-Indeterminate / NA.

    Thresholds come from the cBioPortal clinical sample file header itself.
    """
    if pd.isna(score):
        return "NA"
    if score >= schema.MSI_SENSOR_HIGH:
        return "MSI-H"
    if score < schema.MSI_SENSOR_LOW:
        return "MSS"
    return "MSI-Indeterminate"


_STAGE_RE = re.compile(r"STAGE\s+(IV|III|II|I)", re.IGNORECASE)


def normalise_stage(raw: object) -> str:
    """Bucket AJCC strings like 'STAGE IIIB' -> 'III'. NA for anything unparseable."""
    if not isinstance(raw, str):
        return "NA"
    m = _STAGE_RE.match(raw.strip())
    return m.group(1).upper() if m else "NA"


def parse_os_status(raw: object) -> float:
    """'0:LIVING' -> 0, '1:DECEASED' -> 1, otherwise NaN."""
    if not isinstance(raw, str):
        return float("nan")
    head = raw.split(":", 1)[0].strip()
    if head == "0":
        return 0.0
    if head == "1":
        return 1.0
    return float("nan")


# --- IO ---------------------------------------------------------------------

def read_clinical(path: Path, required: list[str]) -> pd.DataFrame:
    """Read a cBioPortal clinical TSV (which prefixes 4 metadata lines with '#').

    Raises KeyError with a helpful message if any required column is missing.
    """
    df = pd.read_csv(path, sep="\t", comment="#", dtype=str, low_memory=False)
    missing = [c for c in required if c not in df.columns]
    if missing:
        raise KeyError(
            f"{path.name}: missing required columns {missing}. "
            f"Columns present: {list(df.columns)}"
        )
    return df


def build_clinical() -> tuple[pd.DataFrame, dict]:
    sample_path = schema.RAW_DIR / schema.FILES["sample"][0]
    patient_path = schema.RAW_DIR / schema.FILES["patient"][0]

    sample = read_clinical(sample_path, schema.REQUIRED_SAMPLE_COLS)
    patient = read_clinical(patient_path, schema.REQUIRED_PATIENT_COLS)

    # Numeric coercion on fields we need to compare or filter on.
    sample["MSI_SENSOR_SCORE"] = pd.to_numeric(
        sample["MSI_SENSOR_SCORE"], errors="coerce"
    )
    sample["TMB_NONSYNONYMOUS"] = pd.to_numeric(
        sample["TMB_NONSYNONYMOUS"], errors="coerce"
    )
    patient["AGE"] = pd.to_numeric(patient["AGE"], errors="coerce")
    patient["OS_MONTHS"] = pd.to_numeric(patient["OS_MONTHS"], errors="coerce")

    merged = sample.merge(
        patient, on="PATIENT_ID", how="left", suffixes=("", "_pat")
    )

    out = pd.DataFrame({
        "sample_id":   merged["SAMPLE_ID"],
        "patient_id":  merged["PATIENT_ID"],
        # COAD (colon adeno), READ (rectal adeno), MACR (mucinous adeno of colon
        # and rectum — also CRC, kept as its own bucket). Anything else (none in
        # the current study but defensive) -> OTHER.
        "site":        merged["ONCOTREE_CODE"].map(
            lambda x: x if x in {"COAD", "READ", "MACR"} else "OTHER"
        ),
        "msi_sensor_score": merged["MSI_SENSOR_SCORE"],
        "msi_status":  merged["MSI_SENSOR_SCORE"].map(label_msi),
        "tmb":         merged["TMB_NONSYNONYMOUS"],
        "age":         merged["AGE"],
        "sex":         merged["SEX"],
        "stage_raw":   merged["AJCC_PATHOLOGIC_TUMOR_STAGE"],
        "stage":       merged["AJCC_PATHOLOGIC_TUMOR_STAGE"].map(normalise_stage),
        "os_event":    merged["OS_STATUS"].map(parse_os_status),
        "os_months":   merged["OS_MONTHS"],
        "subtype":     merged.get("SUBTYPE"),
    })

    provenance = {
        "sample_file":       sample_path.name,
        "patient_file":      patient_path.name,
        "msi_label_column":  "MSI_SENSOR_SCORE (derived; threshold >= "
                             f"{schema.MSI_SENSOR_HIGH} -> MSI-H, < "
                             f"{schema.MSI_SENSOR_LOW} -> MSS, between -> Indeterminate)",
        "tmb_column":        "TMB_NONSYNONYMOUS (data_clinical_sample.txt)",
        "stage_column":      "AJCC_PATHOLOGIC_TUMOR_STAGE (data_clinical_patient.txt)",
        "age_column":        "AGE (data_clinical_patient.txt)",
        "sex_column":        "SEX (data_clinical_patient.txt)",
        "survival_columns":  "OS_STATUS, OS_MONTHS (data_clinical_patient.txt)",
        "site_column":       "ONCOTREE_CODE (data_clinical_sample.txt; COAD vs READ)",
    }
    return out, provenance


def build_expression(clinical: pd.DataFrame) -> tuple[pd.DataFrame, int, int]:
    """Read expression TSV, keep `Hugo_Symbol`-indexed gene x sample matrix.

    Returns (matrix, n_dropped_blank_symbol, n_dropped_duplicate_symbol).
    """
    expr_path = schema.RAW_DIR / schema.FILES["expression"][0]
    df = pd.read_csv(expr_path, sep="\t", low_memory=False)

    if "Hugo_Symbol" not in df.columns:
        raise KeyError(
            f"{expr_path.name}: missing 'Hugo_Symbol' column. "
            f"Columns present: {list(df.columns)[:8]}..."
        )

    # Drop the Entrez ID column; we key on gene symbol.
    if "Entrez_Gene_Id" in df.columns:
        df = df.drop(columns=["Entrez_Gene_Id"])

    blank_mask = df["Hugo_Symbol"].isna() | (df["Hugo_Symbol"].astype(str).str.strip() == "")
    n_blank = int(blank_mask.sum())
    df = df.loc[~blank_mask]

    dup_mask = df["Hugo_Symbol"].duplicated(keep=False)
    n_dup = int(dup_mask.sum())
    df = df.loc[~dup_mask]

    df = df.set_index("Hugo_Symbol")

    # Restrict sample columns to those we have clinical data for.
    keep_samples = [c for c in df.columns if c in set(clinical["sample_id"])]
    df = df[keep_samples]

    # Coerce to float (RSEM values are numeric).
    df = df.apply(pd.to_numeric, errors="coerce")
    return df, n_blank, n_dup


# --- main -------------------------------------------------------------------

def main() -> int:
    schema.PROCESSED_DIR.mkdir(parents=True, exist_ok=True)

    print("=" * 75)
    print(f"OncoDSL build  —  study: {schema.STUDY}")
    print(f"Source URL    : {schema.BASE_URL}")
    print(f"Raw dir       : {schema.RAW_DIR}")
    print(f"Processed dir : {schema.PROCESSED_DIR}")
    print("=" * 75)

    # Files present check (required only).
    missing_required = [
        fname for fname, req in schema.FILES.values()
        if req and not (schema.RAW_DIR / fname).exists()
    ]
    if missing_required:
        print(f"\nABORT: missing required files in {schema.RAW_DIR}: {missing_required}")
        print("Run: python -m data_pipeline.download")
        return 1

    print("\n[1/3] Reading clinical data ...")
    clinical, provenance = build_clinical()
    print(f"      sample rows: {len(clinical)}, "
          f"unique patients: {clinical['patient_id'].nunique()}")

    print("\n[2/3] Reading expression matrix ...")
    expression, n_blank, n_dup = build_expression(clinical)
    print(f"      gene x sample matrix: {expression.shape[0]} genes x "
          f"{expression.shape[1]} samples")
    if n_blank or n_dup:
        print(f"      dropped {n_blank} blank-symbol rows, "
              f"{n_dup} duplicate-symbol rows")

    # Mark which samples have expression — used downstream.
    expr_samples = set(expression.columns)
    clinical["has_expression"] = clinical["sample_id"].isin(expr_samples)

    print("\n[3/3] Writing parquet outputs ...")
    clin_out = schema.PROCESSED_DIR / "clinical.parquet"
    expr_out = schema.PROCESSED_DIR / "expression.parquet"
    clinical.to_parquet(clin_out, index=False)
    expression.to_parquet(expr_out)
    print(f"      wrote {clin_out}")
    print(f"      wrote {expr_out}")

    # ----- Provenance report -----
    print("\n" + "-" * 75)
    print("PROVENANCE (which file/column was used for each field)")
    print("-" * 75)
    for k, v in provenance.items():
        print(f"  {k:18s} : {v}")
    print("  NOTE             : MSI label is DERIVED from the continuous "
          "MSIsensor score — cBioPortal does not ship a clean MSI-H/MSS column.")

    # ----- Cohort report -----
    n_total = len(clinical)
    n_coad = int((clinical["site"] == "COAD").sum())
    n_read = int((clinical["site"] == "READ").sum())
    n_macr = int((clinical["site"] == "MACR").sum())
    n_other = int((clinical["site"] == "OTHER").sum())

    missing_counts = {
        "msi_status (=NA)":           int((clinical["msi_status"] == "NA").sum()),
        "msi_status (=Indeterm.)":    int((clinical["msi_status"] == "MSI-Indeterminate").sum()),
        "tmb":                        int(clinical["tmb"].isna().sum()),
        "age":                        int(clinical["age"].isna().sum()),
        "sex (blank)":                int(clinical["sex"].isna().sum() + (clinical["sex"] == "").sum()),
        "stage (=NA)":                int((clinical["stage"] == "NA").sum()),
        "os_event":                   int(clinical["os_event"].isna().sum()),
        "os_months":                  int(clinical["os_months"].isna().sum()),
        "no expression":              int((~clinical["has_expression"]).sum()),
    }

    usable_mask = (
        clinical["msi_status"].isin(["MSI-H", "MSS"])
        & clinical["has_expression"]
        & (clinical["stage"] != "NA")
        & clinical["age"].notna()
    )
    n_usable = int(usable_mask.sum())
    n_usable_read = int((usable_mask & (clinical["site"] == "READ")).sum())
    n_usable_coad = int((usable_mask & (clinical["site"] == "COAD")).sum())

    # MSI label agreement with patient SUBTYPE *_MSI as a sanity check.
    subtype_agreement_lines = []
    if "subtype" in clinical.columns and clinical["subtype"].notna().any():
        sub = clinical[clinical["subtype"].notna()].copy()
        sub["subtype_msi"] = sub["subtype"].str.endswith("_MSI")
        sub["derived_msi"] = sub["msi_status"] == "MSI-H"
        ct = pd.crosstab(
            sub["derived_msi"].map({True: "derived=MSI-H", False: "derived=MSS/Other"}),
            sub["subtype_msi"].map({True: "SUBTYPE=*_MSI", False: "SUBTYPE!=*_MSI"}),
        )
        subtype_agreement_lines = ct.to_string().splitlines()

    print("\n" + "-" * 75)
    print("COHORT")
    print("-" * 75)
    print(f"  Total samples     : {n_total}")
    print(f"    COAD            : {n_coad}   (colon adenocarcinoma)")
    print(f"    READ            : {n_read}   (rectal adenocarcinoma)")
    print(f"    MACR            : {n_macr}   (mucinous adenocarcinoma of colon/rectum)")
    if n_other:
        print(f"    OTHER           : {n_other}")
    print(f"  Genes (expression): {expression.shape[0]}")
    print()
    print("  Missing / dropped per field:")
    for k, v in missing_counts.items():
        print(f"    {k:30s} {v}")
    print()
    print(f"  USABLE cohort N     : {n_usable}   "
          "(MSI label in {MSI-H, MSS} AND expression AND stage AND age)")
    print(f"    of which COAD     : {n_usable_coad}")
    print(f"    of which READ     : {n_usable_read}")

    if subtype_agreement_lines:
        print("\n  Agreement: derived MSI label vs patient SUBTYPE ending in '_MSI':")
        for line in subtype_agreement_lines:
            print(f"    {line}")
    print("-" * 75)
    print("Done.")
    return 0


if __name__ == "__main__":
    sys.exit(main())