src/workflows/geo.py

"""
GEO (Gene Expression Omnibus) data loading and GPL platform annotation helpers.

Public API
----------
GEO series matrix
~~~~~~~~~~~~~~~~~
load_geo_series_matrix_lines   Download or read a GEO series matrix file; return raw lines.
parse_geo_series_matrix_lines  Parse raw lines; return sample IDs, probe IDs, expression
                               matrix, and sample characteristics.
decode_geo_numeric_codes       Decode GEO characteristic columns whose name embeds a
                               numeric-to-label mapping (e.g. tumor_subtype_0na_1classical_2basal).

GPL platform annotation
~~~~~~~~~~~~~~~~~~~~~~~
GPL_GENE_SYMBOL_CANDIDATES     Ordered list of candidate gene-symbol column names in GPL
                               annotation/SOFT tables.
gpl_accession_to_url           Convert a GPL accession to its NCBI .annot.gz annotation URL.
load_gpl_soft_lines            Load GPL annotation/SOFT file lines from an accession, URL, or
                               local path.
parse_gpl_soft                 Parse GPL annotation/SOFT lines (.annot.gz or _family.soft.gz
                               table format); return probe-to-gene mapping dict.
load_gpl_mapping_via_gse       Fallback for platforms with no .annot.gz (e.g. GPL13667):
                               look up the platform table via GEOparse + a GSE-level family
                               SOFT file; return a probe-to-gene mapping dict (same shape
                               as parse_gpl_soft's).
"""

from __future__ import annotations

import re
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd

# Fixed, non-exhaustive list of known gene-symbol column names. Platforms
# using other naming conventions need gene_symbol_column passed explicitly
# (see gpl_columns_available in load_gpl_mapping_via_gse / the
# decoupler_annotate_probes_with_gpl tool's return dict).
GPL_GENE_SYMBOL_CANDIDATES: list[str] = [
    "Gene symbol",
    "GENE_SYMBOL",
    "Gene Symbol",
    "SYMBOL",
    "symbol",
    "gene_symbol",
    "GENE_NAME",
    "Gene_Name",
    "gene_name",
]

# Tier 2 fallback: when a bare GPL accession's .annot.gz is unavailable and
# no gse_accession was passed, look up a registered GSE known to use that
# platform here instead of raising. Seeded from biodata-registry manifests'
# accession/platform fields (one representative GSE per GPL is enough — the
# family SOFT file embeds the full platform table regardless of which GSE
# it came from). An explicit gse_accession argument always takes priority
# over this table.
GPL_TO_KNOWN_GSE: dict[str, str] = {
    "GPL6244": "GSE28735",  # gse28735_pdac (Affymetrix Human Gene 1.0 ST)
    "GPL13667": "GSE62165",  # gse62165_jiang (Affymetrix Human Genome U219)
    "GPL15048": "GSE57495",  # gse57495 (Affymetrix Human Exon 1.0 ST, reannotated)
    "GPL10558": "GSE50827",  # gse50827_nones (Illumina HumanHT-12 V4.0)
    "GPL4133": "GSE21501",  # gse21501_stratford (Agilent-014850)
    "GPL570": "GSE15471",  # gse15471_badea (Affymetrix HG-U133 Plus 2.0)
}


def decode_geo_numeric_codes(col_name: str, series: pd.Series) -> tuple[str, pd.Series] | None:
    """
    Detect GEO characteristic columns whose name embeds a numeric-to-label
    mapping, e.g. 'tumor_subtype_0na_1classical_2basal', and return
    (base_name, decoded_series).  Returns None if the pattern is not found.

    Convention: the suffix after the base name consists of one or more segments
    of the form <digits><word>, e.g. 0na, 1classical, 2basal.  "na" decodes to
    "NA"; everything else is Title-cased.  Values in the series are coerced to
    int-strings before lookup so "0", "0.0", and " 0 " all map correctly.
    """
    match = re.search(r"^(.+?)((?:_\d+[a-z]+)+)$", col_name)
    if not match:
        return None

    base_name = match.group(1)
    code_map = {
        num: ("NA" if label == "na" else label.title())
        for num, label in re.findall(r"(\d+)([a-z]+)", match.group(2))
    }
    if not code_map:
        return None

    def _to_key(v: object) -> str:
        try:
            return str(int(float(str(v).strip())))
        except (ValueError, TypeError):
            return str(v).strip()

    decoded = series.apply(lambda v: code_map.get(_to_key(v), _to_key(v)))
    return base_name, decoded


def load_geo_series_matrix_lines(url_or_path: str) -> list[str]:
    """
    Download or read a GEO series matrix file and return its lines.

    Accepts http://, https://, and ftp:// URLs, or a local file path.
    Handles .gz compressed files automatically.

    Parameters
    ----------
    url_or_path:
        URL or local path to the series matrix file.

    Returns
    -------
    List of decoded text lines (no trailing newlines).

    Raises
    ------
    FileNotFoundError  if a local path does not exist.
    """
    import gzip
    import urllib.request

    # ADR-0010: this loader bypasses src/core/data_io.resolve_to_local_path (it
    # streams the series matrix straight into memory), so run the same on-load
    # integrity check here. verify_* is a no-op unless a manifest baselines this
    # exact URL/path, so untracked GEO fetches are unaffected; a baselined source
    # whose bytes have been altered is refused before parsing (IntegrityError).
    from src.core.integrity import verify_bytes, verify_file

    if url_or_path.startswith(("http://", "https://", "ftp://")):
        with urllib.request.urlopen(url_or_path) as resp:
            raw = resp.read()
        # Verify the raw response as served (still compressed for .gz), matching
        # how the baseline was recorded from this URL — before decompression.
        verify_bytes(raw, url_or_path)
        if url_or_path.rstrip("?").endswith(".gz"):
            raw = gzip.decompress(raw)
        return raw.decode("utf-8", errors="replace").splitlines()

    src = Path(url_or_path)
    if not src.exists():
        raise FileNotFoundError(f"File not found: {url_or_path}")
    verify_file(src, url_or_path)
    if src.suffix == ".gz":
        with gzip.open(src, "rt", encoding="utf-8", errors="replace") as fh:
            return fh.read().splitlines()
    with open(src, encoding="utf-8", errors="replace") as fh:
        return fh.read().splitlines()


def parse_geo_series_matrix_lines(lines: list[str]) -> dict[str, Any]:
    """
    Parse lines from a GEO series matrix file.

    Extracts sample characteristics from !Sample_characteristics_ch* lines,
    GEO accession IDs from !Sample_geo_accession, and the expression matrix
    between !series_matrix_table_begin / !series_matrix_table_end markers.

    The returned matrix X is transposed from GEO convention (probes × samples)
    to AnnData convention (samples × probes).

    Parameters
    ----------
    lines:
        Lines from a GEO series matrix file (from load_geo_series_matrix_lines).

    Returns
    -------
    dict with keys:
        sample_ids             (list[str])             — GEO sample IDs from the matrix header.
        probe_ids              (list[str])             — probe / feature IDs.
        X                      (np.ndarray, float32)   — (n_samples, n_probes) expression matrix.
        sample_characteristics (dict[str, list[str]])  — characteristic_key → per-sample values.
        geo_accessions         (list[str])             — GSM accessions from !Sample_geo_accession.

    Raises
    ------
    ValueError  if no !series_matrix_table_begin marker is found.
    """
    sample_characteristics: dict[str, list[str]] = {}
    geo_accessions: list[str] = []
    in_matrix = False
    matrix_lines: list[str] = []

    for line in lines:
        if line.startswith("!series_matrix_table_begin"):
            in_matrix = True
            continue
        if line.startswith("!series_matrix_table_end"):
            in_matrix = False
            continue
        if in_matrix:
            matrix_lines.append(line)
            continue

        if line.startswith("!Sample_geo_accession"):
            parts = line.split("\t")
            geo_accessions = [p.strip().strip('"') for p in parts[1:] if p.strip()]

        elif line.startswith("!Sample_characteristics_ch"):
            parts = line.split("\t")
            raw_values = [p.strip().strip('"') for p in parts[1:]]
            key = next((v.split(":", 1)[0].strip() for v in raw_values if ":" in v), None)
            if key:
                col_values = [v.split(":", 1)[1].strip() if ":" in v else v for v in raw_values]
                sample_characteristics[key] = col_values

    if not matrix_lines:
        raise ValueError(
            "No expression matrix found — missing !series_matrix_table_begin marker. "
            "Verify the file is a valid GEO series matrix."
        )

    header_parts = matrix_lines[0].split("\t")
    sample_ids = [h.strip().strip('"') for h in header_parts[1:] if h.strip()]

    probe_ids: list[str] = []
    expr_rows: list[list[float]] = []
    for line in matrix_lines[1:]:
        if not line.strip():
            continue
        parts = line.split("\t")
        probe_ids.append(parts[0].strip().strip('"'))
        vals: list[float] = []
        for v in parts[1 : len(sample_ids) + 1]:
            v = v.strip().strip('"')
            try:
                vals.append(float(v))
            except ValueError:
                vals.append(np.nan)
        while len(vals) < len(sample_ids):
            vals.append(np.nan)
        expr_rows.append(vals)

    # GEO matrix is probes × samples — transpose to samples × probes for AnnData
    X = np.array(expr_rows, dtype=np.float32).T  # (n_samples, n_probes)

    return {
        "sample_ids": sample_ids,
        "probe_ids": probe_ids,
        "X": X,
        "sample_characteristics": sample_characteristics,
        "geo_accessions": geo_accessions,
    }


def gpl_accession_to_url(accession: str) -> str:
    """
    Convert a GPL accession string to its NCBI .annot.gz platform annotation URL.

    NCBI directory structure uses the accession number with the last 3 digits
    replaced by 'nnn' as the parent directory name.

    .annot.gz is NCBI's "GEO2R"-generated platform annotation file — typically
    1-10 MB, with a 'Gene symbol' column (multi-gene probes '///'-joined with
    no surrounding spaces, e.g. 'MIR4640///DDR1'). This is much smaller than
    the legacy '_family.soft.gz' (can be GBs for widely-used platforms such as
    GPL570/GPL6244). Some platforms do not have a .annot.gz file (404, e.g.
    GPL13667) — for those, pass a direct URL or local path to a SOFT file to
    load_gpl_soft_lines instead of a bare GPL accession.

    Examples
    --------
    "GPL14951" → https://ftp.ncbi.nlm.nih.gov/geo/platforms/GPL14nnn/GPL14951/annot/GPL14951.annot.gz
    "GPL96"    → https://ftp.ncbi.nlm.nih.gov/geo/platforms/GPLnnn/GPL96/annot/GPL96.annot.gz
    """
    stripped = accession.strip().upper()
    num_str = stripped[3:] if stripped.startswith("GPL") else stripped
    if not num_str.isdigit():
        raise ValueError(
            f"Cannot parse GPL accession '{accession}'. Expected 'GPL14951' or '14951'."
        )
    dir_prefix = (num_str[:-3] + "nnn") if len(num_str) > 3 else "nnn"
    return (
        f"https://ftp.ncbi.nlm.nih.gov/geo/platforms/"
        f"GPL{dir_prefix}/GPL{num_str}/annot/GPL{num_str}.annot.gz"
    )


def load_gpl_soft_lines(source: str) -> tuple[str, list[str]]:
    """
    Load GPL SOFT file lines from a GPL accession, URL, or local path.

    Returns (source_used, lines) where source_used is the resolved URL or path.
    """
    import gzip as _gz
    import urllib.request as _ur

    s = source.strip()
    upper = s.upper()

    if upper.startswith("GPL") and s[3:].isdigit():
        resolved = gpl_accession_to_url(s)
    elif "://" in s:
        resolved = s
    else:
        local = Path(s)
        if not local.exists():
            raise FileNotFoundError(f"GPL file not found: {s}")
        if local.suffix == ".gz":
            with _gz.open(local, "rt", encoding="utf-8", errors="replace") as fh:
                return s, fh.read().splitlines()
        else:
            with open(local, encoding="utf-8", errors="replace") as fh:
                return s, fh.read().splitlines()

    with _ur.urlopen(resolved) as resp:
        raw = resp.read()
    content = (
        _gz.decompress(raw).decode("utf-8", errors="replace")
        if resolved.rstrip("?").endswith(".gz")
        else raw.decode("utf-8", errors="replace")
    )
    return resolved, content.splitlines()


def parse_gpl_soft(
    lines: list[str],
    gene_symbol_col: str | None,
) -> tuple[dict[str, str], str]:
    """
    Parse GPL annotation/SOFT lines and return ({probe_id: gene_symbol}, column_used).

    Recognizes the data table in either of NCBI's two table formats:
      - .annot.gz (GEO2R platform annotation): !platform_table_begin /
        !platform_table_end, with a 'Gene symbol' column.
      - _family.soft.gz (legacy platform SOFT): !Platform_data_table_begin /
        !Platform_data_table_end, with a 'GENE_SYMBOL'/'Gene Symbol'-style column.

    Probe IDs come from 'ID' (or 'ID_REF'/'Probe_Id'/'Probe ID', else the first
    column) and gene symbols from gene_symbol_col or auto-detection via
    GPL_GENE_SYMBOL_CANDIDATES. Multi-gene values (e.g. Affymetrix
    'GENE1///GENE2' or 'GENE1 /// GENE2') are passed through unchanged —
    decoupler_collapse_probes_to_genes splits these per its multi_gene_policy.
    """
    import io as _io

    table_lines: list[str] = []
    in_table = False
    for line in lines:
        if line.startswith("!Platform_data_table_begin") or line.startswith(
            "!platform_table_begin"
        ):
            in_table = True
            continue
        if line.startswith("!Platform_data_table_end") or line.startswith("!platform_table_end"):
            break
        if in_table:
            table_lines.append(line)

    if not table_lines:
        raise ValueError(
            "No data table found in GPL file — expected a "
            "!Platform_data_table_begin (_family.soft) or !platform_table_begin "
            "(.annot) block. Verify the file is a valid GPL platform record."
        )

    df = pd.read_csv(
        _io.StringIO("\n".join(table_lines)),
        sep="\t",
        low_memory=False,
        dtype=str,
    )

    probe_col = df.columns[0]
    for candidate in ("ID", "ID_REF", "Probe_Id", "Probe ID"):
        if candidate in df.columns:
            probe_col = candidate
            break

    if gene_symbol_col:
        if gene_symbol_col not in df.columns:
            raise ValueError(
                f"Specified gene_symbol_column '{gene_symbol_col}' not found. "
                f"Available columns: {list(df.columns)}"
            )
        sym_col = gene_symbol_col
    else:
        sym_col = None
        for candidate in GPL_GENE_SYMBOL_CANDIDATES:
            if candidate in df.columns:
                sym_col = candidate
                break
        if sym_col is None:
            raise ValueError(
                f"Could not auto-detect a gene symbol column. "
                f"Available columns: {list(df.columns)}. "
                "Pass gene_symbol_column explicitly."
            )

    mapping: dict[str, str] = {}
    for _, row in df.iterrows():
        probe = str(row[probe_col]).strip()
        sym = str(row[sym_col]).strip()
        if probe and sym and sym not in ("", "nan", "NaN", "None", "---"):
            mapping[probe] = sym

    return mapping, sym_col


# Hard limit on a single family SOFT download. Observed family SOFT files
# for the platforms in GPL_TO_KNOWN_GSE are well under 100MB (e.g. ~96MB for
# GSE62165/GPL13667); widely-used platforms like GPL570 can have family SOFT
# files in the multi-GB range. Reject anything over this limit up front
# rather than failing unpredictably mid-download or filling the cache disk.
FAMILY_SOFT_MAX_DOWNLOAD_BYTES = 150 * 1024 * 1024  # 150 MB

# Total size cap for cached *_family.soft.gz files. When a new download
# would push the cache over this limit, the oldest (by mtime) files are
# evicted first.
FAMILY_SOFT_CACHE_MAX_BYTES = 2 * 1024 * 1024 * 1024  # 2 GB


def _gse_range_subdir(gse_accession: str) -> str:
    """e.g. 'GSE62165' -> 'GSE62nnn', 'GSE71989' -> 'GSE71nnn'."""
    digits = gse_accession[len("GSE") :]
    return f"GSE{digits[:-3]}nnn"


def _evict_family_soft_cache(cache_dir: Path, max_bytes: int = FAMILY_SOFT_CACHE_MAX_BYTES) -> None:
    """
    Delete cached *_family.soft.gz files, oldest (by mtime) first, until the
    total size of remaining files is <= max_bytes.

    Called after a new family SOFT download. A no-op if the cache is
    already within budget.
    """
    files = sorted(cache_dir.glob("*_family.soft.gz"), key=lambda p: p.stat().st_mtime)
    total = sum(f.stat().st_size for f in files)
    for f in files:
        if total <= max_bytes:
            break
        total -= f.stat().st_size
        f.unlink()


def _ensure_gse_family_soft(gse_accession: str, cache_dir: Path) -> Path:
    """
    Pre-place <cache_dir>/<gse_accession>_family.soft.gz via HTTPS so
    GEOparse.get_GEO's FTP download is skipped entirely.

    GEOparse.get_GEO_file's skip-logic is: "if not path.isfile(filepath):
    download via FTP; else: use local version". The FTP download
    (ftp://ftp.ncbi.nlm.nih.gov/geo/series/.../soft/{gse}_family.soft.gz) is
    unreliable for files beyond a few MB (intermittent "Downloaded size do
    not match the expected size"). The HTTPS mirror of the same path is
    reliable, so fetch it that way first.

    Before downloading, checks the response's Content-Length header against
    FAMILY_SOFT_MAX_DOWNLOAD_BYTES and raises OSError if it's too large,
    rather than risking a slow/partial download. After a successful
    download, evicts old cache entries via _evict_family_soft_cache so the
    cache doesn't grow unbounded.

    Returns the local path (existing or newly downloaded). Raises
    requests.HTTPError (an OSError subclass) if the GSE has no family SOFT
    file at the expected URL, or OSError if it exceeds the size limit.
    """
    import requests

    cache_dir.mkdir(parents=True, exist_ok=True)
    filename = f"{gse_accession}_family.soft.gz"
    filepath = cache_dir / filename
    if filepath.is_file():
        return filepath

    range_subdir = _gse_range_subdir(gse_accession)
    url = f"https://ftp.ncbi.nlm.nih.gov/geo/series/{range_subdir}/{gse_accession}/soft/{filename}"
    resp = requests.get(url, stream=True, timeout=180)
    resp.raise_for_status()

    content_length = resp.headers.get("Content-Length")
    if content_length is not None and int(content_length) > FAMILY_SOFT_MAX_DOWNLOAD_BYTES:
        resp.close()
        raise OSError(
            f"{gse_accession} family SOFT file is "
            f"{int(content_length) / (1024 * 1024):.1f}MB, exceeding the "
            f"{FAMILY_SOFT_MAX_DOWNLOAD_BYTES / (1024 * 1024):.0f}MB limit "
            "(FAMILY_SOFT_MAX_DOWNLOAD_BYTES). Pass a direct URL or local "
            "path to a smaller annotation file instead of relying on the "
            "GSE family SOFT fallback for this platform."
        )

    tmp_path = filepath.with_suffix(filepath.suffix + ".tmp")
    with open(tmp_path, "wb") as f:
        for chunk in resp.iter_content(chunk_size=1024 * 1024):
            f.write(chunk)
    tmp_path.rename(filepath)
    _evict_family_soft_cache(cache_dir)
    return filepath


def _first_gene_symbol_from_assignment(value: object) -> str:
    """
    Parse the first gene symbol from a GPL 'gene_assignment'-style column
    (Affymetrix Gene/Exon ST arrays, e.g. GPL6244).

    Format: a '///'-delimited list of groups, each
    '<transcript_id> // <gene_symbol> // <description> // <chr_location> //
    <entrez_id>' (or '---' if the probe has no annotation). Returns the gene
    symbol (2nd field) from the first group that has a usable assignment
    ('---' or malformed groups are skipped in favor of a later group), or ""
    if no group is usable.

    Improves on a previous version that only looked at the first group: a
    probe whose first group is '---' but a later group has a real assignment
    was previously treated as unannotated, capping coverage at ~87.6% for
    GPL6244 on gse28735_pdac. Skipping unusable leading groups raises that
    coverage.
    """
    if not isinstance(value, str):
        return ""
    for group in value.split(" /// "):
        group = group.strip()
        if not group or group == "---":
            continue
        fields = group.split(" // ")
        if len(fields) < 2:
            continue
        symbol = fields[1].strip()
        if symbol:
            return symbol
    return ""


def load_gpl_mapping_via_gse(
    gse_accession: str,
    gpl_accession: str,
    gene_symbol_col: str | None,
    cache_dir: Path,
) -> tuple[dict[str, str], str, str, list[str]]:
    """
    Fallback probe-to-gene lookup for platforms with no .annot.gz (e.g.
    GPL13667): download the GSE-level family SOFT file for gse_accession via
    GEOparse and read the embedded gpl_accession platform table.

    Column selection mirrors parse_gpl_soft: gene_symbol_col if given (must
    exist), else the first GPL_GENE_SYMBOL_CANDIDATES match, else (for
    Affymetrix Gene/Exon ST platforms such as GPL6244) a 'gene_assignment'
    column parsed via _first_gene_symbol_from_assignment, returned as
    sym_col="gene_symbol". Multi-gene values (e.g. 'GENE1 /// GENE2') are
    passed through unchanged, same as parse_gpl_soft.

    Returns (mapping, sym_col_used, source_used, gpl_columns_available) —
    same mapping/sym_col shape as parse_gpl_soft, plus diagnostics.

    Raises ValueError if gpl_accession's table is missing/empty in the GSE
    family file, gene_symbol_col was given but not found, or no usable
    column could be auto-detected. Raises OSError (e.g.
    requests.HTTPError) if the family SOFT file cannot be downloaded.
    """
    import GEOparse

    _ensure_gse_family_soft(gse_accession, cache_dir)
    gse = GEOparse.get_GEO(geo=gse_accession, destdir=str(cache_dir), silent=True)

    source_used = f"GEOparse:{gse_accession}:{gpl_accession}"
    gpl = gse.gpls.get(gpl_accession)
    if gpl is None or gpl.table is None or gpl.table.empty:
        raise ValueError(
            f"{gpl_accession} table not found/empty in {gse_accession} family "
            f"SOFT file (platforms available: {list(gse.gpls.keys())})"
        )

    table = gpl.table
    gpl_columns_available = list(table.columns)

    probe_col = table.columns[0]
    for candidate in ("ID", "ID_REF", "Probe_Id", "Probe ID"):
        if candidate in table.columns:
            probe_col = candidate
            break

    if gene_symbol_col:
        if gene_symbol_col not in table.columns:
            raise ValueError(
                f"Specified gene_symbol_column '{gene_symbol_col}' not found in "
                f"{gpl_accession} table (from {gse_accession}). "
                f"Available columns: {gpl_columns_available}"
            )
        sym_col = gene_symbol_col
        symbols = table[sym_col]
    else:
        sym_col = next((c for c in GPL_GENE_SYMBOL_CANDIDATES if c in table.columns), None)
        if sym_col is not None:
            symbols = table[sym_col]
        elif "gene_assignment" in table.columns:
            sym_col = "gene_symbol"
            symbols = table["gene_assignment"].map(_first_gene_symbol_from_assignment)
        else:
            raise ValueError(
                f"Could not auto-detect a gene symbol column in {gpl_accession} "
                f"table (from {gse_accession}). Available columns: "
                f"{gpl_columns_available}. Pass gene_symbol_column explicitly."
            )

    probe_ids = table[probe_col].astype(str).str.strip()
    sym_values = symbols.astype(str).str.strip()

    mapping: dict[str, str] = {}
    for probe, sym in zip(probe_ids, sym_values, strict=True):
        if probe and sym and sym not in ("", "nan", "NaN", "None", "---"):
            mapping[probe] = sym

    return mapping, sym_col, source_used, gpl_columns_available