src/cache.py

"""
Disk-based dataset cache for DecoupleRpy.

WHY DISK INSTEAD OF MEMORY
---------------------------
The MCP server runs as a short-lived subprocess: every tool call spawns a new
`python server.py` process and exits after returning the result.  A module-level
memory dict is destroyed at the end of each call.

/tmp on HuggingFace Spaces persists for the lifetime of the container (hours to
days), so caching parsed h5ad files there eliminates the ~50MB GEO download on
all calls after the first.

Cache layout
------------
  /tmp/decoupleRpy/cache/
      GSE71729.h5ad          ← parsed AnnData (fast to read_h5ad, no download)
      GSE71729.gpl_cache/
          GPL14951.pkl       ← probe→gene mapping dict
"""

from __future__ import annotations

import json
import re
import threading
from pathlib import Path
from typing import Any

CACHE_DIR = Path("/tmp/decoupleRpy/cache")


# ── In-memory AnnData cache (process-lifetime) ───────────────────────────────
# When server.py runs as a persistent HTTP process (`--transport http`), the
# same h5ad path is read on nearly every tool call of a multi-step run. Parsing
# a 50MB h5ad from disk each time is pure waste once the process is long-lived.
# This module-level cache holds the parsed AnnData for the process lifetime,
# keyed by (abspath, mtime, size) so a rewritten file is transparently
# re-read. Each read returns a `.copy()`, so callers may freely mutate (and
# write back) without corrupting the shared, read-mostly cache — this is what
# makes it safe to share one resident server across concurrent Gradio sessions.
# Under stdio transport (one process per call) it simply never gets a second
# hit, so it's a no-op there rather than a correctness risk.
_MEM_LOCK = threading.Lock()
_MEM_CACHE: dict[str, Any] = {}
_MEM_MAX_ENTRIES = 8  # bound memory; bulk h5ads are large


def _mem_key(path: str) -> str | None:
    try:
        st = Path(path).stat()
    except OSError:
        return None
    return f"{Path(path).resolve()}::{int(st.st_mtime)}::{st.st_size}"


def read_h5ad_cached(path: str) -> Any:
    """Read an h5ad, reusing a process-lifetime in-memory parse when possible.

    Returns a fresh `.copy()` every call, so the caller owns the object and may
    mutate or write it back without affecting the cache. Falls back to a plain
    `sc.read_h5ad` if the path can't be stat'd (e.g. a URL or missing file) so
    the caller sees the normal error.
    """
    import scanpy as sc

    key = _mem_key(path)
    if key is None:
        return sc.read_h5ad(path)

    with _MEM_LOCK:
        cached = _MEM_CACHE.get(key)
    if cached is not None:
        return cached.copy()

    adata = sc.read_h5ad(path)
    with _MEM_LOCK:
        # Drop stale entries for the same resolved path (older mtime/size).
        resolved = str(Path(path).resolve())
        for k in [k for k in _MEM_CACHE if k.startswith(resolved + "::") and k != key]:
            _MEM_CACHE.pop(k, None)
        _MEM_CACHE[key] = adata
        # Evict oldest if over the bound (dict preserves insertion order).
        while len(_MEM_CACHE) > _MEM_MAX_ENTRIES:
            _MEM_CACHE.pop(next(iter(_MEM_CACHE)))
    return adata.copy()


# ── Cache key helpers ────────────────────────────────────────────────────────


def url_to_cache_key(url_or_path: str) -> str:
    """Return a stable GSE accession key from any URL or path string.
    Picks the match with the most digits to avoid partial matches like
    GSE71 from GSE71nnn/GSE71729."""
    matches = re.findall(r"GSE(\d+)", url_or_path, re.IGNORECASE)
    if not matches:
        return url_or_path
    return "GSE" + max(matches, key=len)


# ── AnnData (dataset) cache ──────────────────────────────────────────────────


def _dataset_cache_path(cache_key: str) -> Path:
    return CACHE_DIR / f"{cache_key}.h5ad"


def get_dataset(cache_key: str) -> Any | None:
    """Return a cached AnnData, or None if not cached."""
    path = _dataset_cache_path(cache_key)
    if path.exists() and path.stat().st_size > 0:
        try:
            import scanpy as sc

            print(f"[cache] Disk hit for {cache_key} — loading from {path}")
            return sc.read_h5ad(path)
        except Exception as exc:
            print(f"[cache] Disk read failed for {cache_key}: {exc} — will re-download")
            path.unlink(missing_ok=True)
    return None


def set_dataset(cache_key: str, adata: Any) -> None:
    """Write an AnnData to the disk cache."""
    try:
        CACHE_DIR.mkdir(parents=True, exist_ok=True)
        path = _dataset_cache_path(cache_key)
        adata.write_h5ad(path)
        print(f"[cache] Wrote {cache_key} to disk cache ({path.stat().st_size // 1024} KB)")
    except Exception as exc:
        print(f"[cache] Failed to write {cache_key} to disk cache: {exc}")


def is_loaded(cache_key: str) -> bool:
    return _dataset_cache_path(cache_key).exists()


# ── GPL probe mapping cache ──────────────────────────────────────────────────


def _gpl_cache_path(accession: str) -> Path:
    # JSON (not pickle): the mapping is a plain dict[str, str], so JSON is safe
    # to deserialize and human-inspectable. Old .pkl files are simply ignored
    # (cache miss → regenerated).
    return CACHE_DIR / "gpl" / f"{accession}.json"


def get_gpl(accession: str) -> dict[str, str] | None:
    """Return a cached probe→gene mapping, or None if not cached."""
    path = _gpl_cache_path(accession)
    if path.exists() and path.stat().st_size > 0:
        try:
            print(f"[cache] Disk hit for GPL {accession}")
            with open(path, encoding="utf-8") as f:
                return json.load(f)
        except Exception as exc:
            print(f"[cache] Disk read failed for GPL {accession}: {exc}")
            path.unlink(missing_ok=True)
    return None


def set_gpl(accession: str, mapping: dict[str, str]) -> None:
    """Write a probe→gene mapping to the disk cache."""
    try:
        _gpl_cache_path(accession).parent.mkdir(parents=True, exist_ok=True)
        with open(_gpl_cache_path(accession), "w", encoding="utf-8") as f:
            json.dump(mapping, f)
        print(f"[cache] Wrote GPL {accession} to disk cache")
    except Exception as exc:
        print(f"[cache] Failed to write GPL {accession} to disk cache: {exc}")


# ── Preloader (optional warm-up, still useful for first-query speed) ─────────


def preload_datasets() -> None:
    """
    Pre-populate the disk cache at startup by downloading all registered
    geo_series_matrix datasets.  Runs in a background thread so the server
    stays responsive.  No-ops if the cache file already exists.
    """
    import pandas as pd
    import scanpy as sc

    from src.datasets.registry import get_registry
    from src.tools.rna import decode_geo_numeric_codes
    from src.workflows.geo import load_geo_series_matrix_lines, parse_geo_series_matrix_lines

    registry = get_registry()

    for dataset_id in registry.list():
        raw = registry.get(dataset_id)
        if raw is None:
            continue
        expr_src = raw.get("expression_source", {})
        if expr_src.get("type") != "geo_series_matrix":
            continue
        url = expr_src.get("url")
        if not url:
            continue

        cache_key = url_to_cache_key(url)
        if is_loaded(cache_key):
            print(f"[cache] {dataset_id} already on disk — skipping preload")
            continue

        print(f"[cache] Preloading {dataset_id} from {url} …")
        try:
            lines = load_geo_series_matrix_lines(url)
            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 = sc.AnnData(X=X, obs=obs_df)
            adata.var.index = pd.Index(probe_ids, name="probe_id")

            set_dataset(cache_key, adata)
            print(f"[cache] {dataset_id} preloaded: {adata.n_obs} × {adata.n_vars}")
        except Exception as exc:
            print(f"[cache] Failed to preload {dataset_id}: {exc}")