scripts/_precompute_common.py

"""
Shared helpers for precompute-cache assemble_*.py scripts.

These mirror the live GEO-loading code paths in src/workflows/geo.py and
src/tools/rna.py (decoupler_load_geo_series_matrix,
decoupler_annotate_probes_with_gpl) so precomputed h5ads have the same
shape/obs/var as a live load -- just without the per-query reasoning steps.
"""

import sys
from datetime import datetime, timezone
from pathlib import Path

import anndata as ad
import pandas as pd
from huggingface_hub import HfApi

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from src.workflows.geo import (
    decode_geo_numeric_codes,
    load_geo_series_matrix_lines,
    parse_geo_series_matrix_lines,
)

HF_REPO = "anne-voigt/pdac-research-data"


def load_series_matrix_to_anndata(url_or_path: str) -> ad.AnnData:
    """
    Load a GEO series matrix into an AnnData with the same obs columns
    (including decoded numeric-code companions) that
    decoupler_load_geo_series_matrix produces for the same URL.
    """
    lines = load_geo_series_matrix_lines(url_or_path)
    parsed = parse_geo_series_matrix_lines(lines)
    sample_ids = parsed["sample_ids"]
    probe_ids = parsed["probe_ids"]
    X = parsed["X"]
    sample_characteristics = parsed["sample_characteristics"]

    obs_df = pd.DataFrame(index=sample_ids)
    for key, values in sample_characteristics.items():
        col_name = key.lower().replace(" ", "_").replace("-", "_")
        if len(values) == len(sample_ids):
            obs_df[col_name] = values

    for col in list(obs_df.columns):
        result = decode_geo_numeric_codes(col, obs_df[col])
        if result is not None:
            base_name, decoded = result
            if base_name not in obs_df.columns:
                obs_df[base_name] = decoded

    adata = ad.AnnData(X=X, obs=obs_df)
    adata.var.index = pd.Index(probe_ids, name="probe_id")
    return adata


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 _ensure_gse_family_soft(gse_accession: str, out_dir: Path) -> Path:
    """
    Pre-place <out_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 in this environment for files beyond a few MB (intermittent
    "Downloaded size do not match the expected size" -- confirmed for both a
    19MB and a 96MB file). The HTTPS mirror of the same path is reliable
    (confirmed via curl -sIL: 200 + correct Content-Length for all of
    GSE71989/GSE62165/GSE16515/GSE28735), so fetch it that way first.
    """
    import requests

    filename = f"{gse_accession}_family.soft.gz"
    filepath = out_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}"
    print(f"  Pre-fetching {url} ...")
    resp = requests.get(url, stream=True, timeout=180)
    resp.raise_for_status()
    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)
    print(f"  Saved {filepath} ({filepath.stat().st_size / 1e6:.1f} MB)")
    return filepath


def annotate_with_gpl(adata, gse_accession, gpl_accession, gene_symbol_column="Gene Symbol"):
    """
    Annotate adata.var with a gene-symbol column via GEOparse, using the
    platform table embedded in the GSE-level family SOFT file
    (gse.gpls[gpl_accession].table) -- the same source as
    scripts/assemble_gse15471.py / assemble_gse17891.py.

    NOTE: this deliberately does NOT use src.workflows.geo's
    gpl_accession_to_url/load_gpl_soft_lines. As of 2026-06,
    gpl_accession_to_url targets the platform-level ".annot.gz" file
    (small, 1-10MB) rather than "_family.soft.gz" (multi-gigabyte for
    heavily-reused platforms -- confirmed: GPL570_family.soft.gz=71.9GB,
    GPL6244_family.soft.gz=8.7GB; GPL13667_family.soft.gz=6GB).
    decoupler_annotate_probes_with_gpl now works directly for GPL570/GPL6244
    via .annot.gz (confirmed: 45118/54675 and 22195/28869 probes annotated
    respectively), but GPL13667 has no .annot.gz (404), so it still needs a
    fallback. The GSE-level family file (tens of MB) embeds a scoped copy of
    the platform table and remains the proven-working source for this script
    (and for GPL13667-style platforms generally).

    Unlike the live tool, this does NOT raise if gene_symbol_column is
    missing -- it returns sym_col_used=None plus gpl_columns_available so
    the caller can inspect available columns and fall back.

    Returns (adata, sym_col_used_or_None, stats_dict).
    """
    import GEOparse

    out_dir = ROOT / "tmp" / "datasets"
    out_dir.mkdir(parents=True, exist_ok=True)
    _ensure_gse_family_soft(gse_accession, out_dir)
    gse = GEOparse.get_GEO(geo=gse_accession, destdir=str(out_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:
        return adata, None, {
            "source_used": source_used,
            "gpl_columns_available": [],
            "requested_column_found": False,
            "error": f"{gpl_accession} table not found/empty in {gse_accession} "
                     f"(available GPLs: {list(gse.gpls.keys())})",
        }

    gpl_columns_available = list(gpl.table.columns)
    if gene_symbol_column not in gpl_columns_available:
        return adata, None, {
            "source_used": source_used,
            "gpl_columns_available": gpl_columns_available,
            "requested_column_found": False,
        }

    probe_col = "ID" if "ID" in gpl.table.columns else gpl.table.columns[0]
    gpl_map = gpl.table.set_index(probe_col)[gene_symbol_column]
    gpl_map = gpl_map[~gpl_map.index.duplicated()]

    probe_ids = adata.var.index.astype(str)
    mapped = gpl_map.reindex(probe_ids).fillna("")
    adata.var[gene_symbol_column] = mapped.values

    n_total = len(probe_ids)
    n_annotated = int((mapped != "").sum())
    return adata, gene_symbol_column, {
        "source_used": source_used,
        "gpl_columns_available": gpl_columns_available,
        "n_total": n_total,
        "n_annotated": n_annotated,
        "n_unannotated": n_total - n_annotated,
        "coverage_pct": round(100 * n_annotated / n_total, 1) if n_total else 0.0,
        "requested_column_found": True,
    }


def annotate_with_gpl_gene_assignment(adata, gse_accession, gpl_accession, assignment_column="gene_assignment"):
    """
    Annotate adata.var with a "gene_symbol" column parsed from a GPL
    'gene_assignment'-style column (Affymetrix Gene/Exon ST arrays, e.g.
    GPL6244), via the same GSE-level family SOFT file as annotate_with_gpl.

    GPL6244's platform table has no "Gene Symbol" column. Instead its
    "gene_assignment" column is a '///'-delimited list of groups, each
    '<transcript_id> // <gene_symbol> // <description> // <chr_location> //
    <entrez_id>' (or '---' if the probe has no annotation). This takes the
    gene symbol (2nd field) from the FIRST group only.

    Returns (adata, "gene_symbol", stats_dict) -- same shape as
    annotate_with_gpl's return (sym_col_used is always "gene_symbol" on
    success, or None if assignment_column is missing).
    """
    import GEOparse

    out_dir = ROOT / "tmp" / "datasets"
    out_dir.mkdir(parents=True, exist_ok=True)
    _ensure_gse_family_soft(gse_accession, out_dir)
    gse = GEOparse.get_GEO(geo=gse_accession, destdir=str(out_dir), silent=True)

    source_used = f"GEOparse:{gse_accession}:{gpl_accession}:{assignment_column}"
    gpl = gse.gpls.get(gpl_accession)
    if gpl is None or gpl.table is None or gpl.table.empty:
        return adata, None, {
            "source_used": source_used,
            "gpl_columns_available": [],
            "requested_column_found": False,
            "error": f"{gpl_accession} table not found/empty in {gse_accession} "
                     f"(available GPLs: {list(gse.gpls.keys())})",
        }

    gpl_columns_available = list(gpl.table.columns)
    if assignment_column not in gpl_columns_available:
        return adata, None, {
            "source_used": source_used,
            "gpl_columns_available": gpl_columns_available,
            "requested_column_found": False,
        }

    probe_col = "ID" if "ID" in gpl.table.columns else gpl.table.columns[0]

    def _first_gene_symbol(value):
        if not isinstance(value, str):
            return ""
        first_group = value.split(" /// ")[0].strip()
        if first_group == "---" or not first_group:
            return ""
        fields = first_group.split(" // ")
        if len(fields) < 2:
            return ""
        return fields[1].strip()

    symbols = gpl.table[assignment_column].map(_first_gene_symbol)
    gpl_map = pd.Series(symbols.values, index=gpl.table[probe_col].astype(str))
    gpl_map = gpl_map[~gpl_map.index.duplicated()]

    probe_ids = adata.var.index.astype(str)
    mapped = gpl_map.reindex(probe_ids).fillna("")
    adata.var["gene_symbol"] = mapped.values

    n_total = len(probe_ids)
    n_annotated = int((mapped != "").sum())
    return adata, "gene_symbol", {
        "source_used": source_used,
        "gpl_columns_available": gpl_columns_available,
        "n_total": n_total,
        "n_annotated": n_annotated,
        "n_unannotated": n_total - n_annotated,
        "coverage_pct": round(100 * n_annotated / n_total, 1) if n_total else 0.0,
        "requested_column_found": True,
    }


def stamp_provenance(adata, *, source_url, dataset_id, biodata_registry_commit, script_name, extra=None):
    """Write plain-string .uns provenance keys (h5ad-safe)."""
    adata.uns["precompute_source_url"] = source_url
    adata.uns["precompute_built_at"] = datetime.now(timezone.utc).isoformat()
    adata.uns["precompute_biodata_registry_commit"] = biodata_registry_commit
    adata.uns["precompute_dataset_id"] = dataset_id
    adata.uns["precompute_script"] = script_name
    if extra:
        for k, v in extra.items():
            adata.uns[k] = str(v)
    return adata


def write_and_upload(adata, dataset_id, dry_run=False):
    """Write to tmp/datasets/{dataset_id}.h5ad and upload to HF_REPO."""
    out_dir = ROOT / "tmp" / "datasets"
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / f"{dataset_id}.h5ad"
    adata.write_h5ad(out_path)
    print(f"Wrote {out_path} ({out_path.stat().st_size / 1e6:.1f} MB)")

    if dry_run:
        print(f"[dry-run] Skipping upload to {HF_REPO}/{dataset_id}.h5ad")
        return out_path

    api = HfApi()
    api.upload_file(
        path_or_fileobj=str(out_path),
        path_in_repo=f"{dataset_id}.h5ad",
        repo_id=HF_REPO,
        repo_type="dataset",
        commit_message=f"Add precomputed {dataset_id} h5ad (precompute cache phase 1)",
    )
    print(f"Uploaded to {HF_REPO}/{dataset_id}.h5ad")
    return out_path