scripts/load_and_validate_geo.py

"""
One-off script to load all 4 GEO datasets, save as h5ad, and validate
against their manifests. Run from the project root with:

    .venv/bin/python scripts/load_and_validate_geo.py
"""

import json
import os
import sys
from pathlib import Path

# Run from project root
ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(ROOT))
os.environ.setdefault("RNA_OUTPUT_DIR", str(ROOT / "tmp" / "datasets"))

from src.tools.dataset_tools import dataset_validate_manifest_against_data
from src.tools.rna import decoupler_load_geo_series_matrix

DATASETS = [
    {
        "dataset_id": "gse71989_chen",
        "url": "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE71nnn/GSE71989/matrix/GSE71989_series_matrix.txt.gz",
        "condition_column": "tissue_subtype",
        "out_prefix": "gse71989_chen",
    },
    {
        "dataset_id": "gse62165_jiang",
        "url": "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE62nnn/GSE62165/matrix/GSE62165_series_matrix.txt.gz",
        "condition_column": "tissue",
        "out_prefix": "gse62165_jiang",
    },
    {
        "dataset_id": "gse16515_mayo",
        "url": "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE16nnn/GSE16515/matrix/GSE16515_series_matrix.txt.gz",
        "condition_column": "tissue",
        "out_prefix": "gse16515_mayo",
    },
    {
        "dataset_id": "gse28735_pdac",
        "url": "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE28nnn/GSE28735/matrix/GSE28735_series_matrix.txt.gz",
        "condition_column": "tissue",
        "out_prefix": "gse28735_pdac",
    },
]


def run():
    results = {}
    for ds in DATASETS:
        did = ds["dataset_id"]
        print(f"\n{'=' * 60}")
        print(f"Loading {did} ...")
        print(f"{'=' * 60}")

        # Step 1: Load from GEO
        try:
            load_result = decoupler_load_geo_series_matrix(
                url_or_path=ds["url"],
                condition_column=ds["condition_column"],
                out_prefix=ds["out_prefix"],
            )
            print(f"  Load: {load_result.get('message', load_result)}")
            adata_path = load_result.get("output_path")
            print(f"  Saved to: {adata_path}")
            print(f"  OBS columns: {load_result.get('obs_columns', [])}")
            print("  Sample characteristics:")
            for k, v in (load_result.get("sample_characteristics") or {}).items():
                vals = v if isinstance(v, list) else list(v)[:5]
                print(f"    {k}: {vals}")
        except Exception as e:
            print(f"  ERROR loading: {e}")
            results[did] = {"load_error": str(e)}
            continue

        if not adata_path or not Path(adata_path).exists():
            print(f"  ERROR: h5ad not found at {adata_path}")
            results[did] = {"load_error": "h5ad not written"}
            continue

        # Step 2: Validate against manifest
        print("\n  Validating against manifest ...")
        try:
            val_result = dataset_validate_manifest_against_data(
                dataset_id=did,
                adata_path=adata_path,
            )
            print(f"  overall_valid: {val_result.get('overall_valid')}")
            print(f"  n_errors: {val_result.get('n_errors')}")
            print(f"  n_warnings: {val_result.get('n_warnings')}")
            for err in val_result.get("errors", []):
                print(f"  ERROR: {err}")
            for warn in val_result.get("warnings", []):
                print(f"  WARNING: {warn}")
            checks = val_result.get("checks", {})
            for check_name, check in checks.items():
                status = check.get("status", "?")
                msg = check.get("message", "")
                print(f"  [{status}] {check_name}: {msg[:120]}")
        except Exception as e:
            print(f"  ERROR validating: {e}")
            val_result = {"validation_error": str(e)}

        results[did] = {
            "load": load_result,
            "validation": val_result,
        }

    # Summary
    print(f"\n{'=' * 60}")
    print("SUMMARY")
    print(f"{'=' * 60}")
    for did, r in results.items():
        if "load_error" in r:
            print(f"  {did}: LOAD FAILED — {r['load_error']}")
        elif "validation_error" in r.get("validation", {}):
            print(f"  {did}: VALIDATION FAILED — {r['validation']['validation_error']}")
        else:
            v = r.get("validation", {})
            print(
                f"  {did}: valid={v.get('overall_valid')} errors={v.get('n_errors')} warnings={v.get('n_warnings')}"
            )

    # Save full results as JSON for inspection
    out = ROOT / "tmp" / "datasets" / "validation_results.json"
    with open(out, "w") as f:
        json.dump(results, f, indent=2, default=str)
    print(f"\nFull results written to: {out}")


if __name__ == "__main__":
    run()