PDGrapher/scripts/fetch_chembl_smiles.py

"""Fetch drug SMILES from ChEMBL by ChEMBL ID.

Downloads SMILES for SciPlex3 drugs from the ChEMBL database via their
REST API. Results are cached to a CSV file for reuse.

Usage:
    python scripts/fetch_chembl_smiles.py \
        --chembl_ids CHEMBL3544932 CHEMBL1257042 \
        --output data/chembl_smiles.csv

    # Or auto-detect from SciPlex3 h5ad:
    python scripts/fetch_chembl_smiles.py \
        --h5ad data/raw/scPerturb/rna_protein/SrivatsanTrapnell2020_sciplex3.h5ad \
        --output data/chembl_smiles.csv
"""

import argparse
import csv
import logging
import time
from pathlib import Path

import anndata as ad
import numpy as np
import pandas as pd
import requests

logger = logging.getLogger(__name__)

CHUNK_SIZE = 50
RATE_LIMIT_DELAY = 0.5  # seconds between API calls


def fetch_smiles_from_chembl(chembl_id: str) -> str:
    """Fetch canonical SMILES for a ChEMBL compound.

    Parameters
    ----------
    chembl_id : str
        ChEMBL compound ID (e.g. "CHEMBL3544932")

    Returns
    -------
    smiles : str
        Canonical SMILES string, or empty string if not found.
    """
    url = f"https://www.ebi.ac.uk/chembl/api/data/molecule/{chembl_id}.json"
    try:
        resp = requests.get(url, timeout=30)
        resp.raise_for_status()
        data = resp.json()
        # ChEMBL API returns flat structure (not nested under 'molecule')
        structures = data.get("molecule_structures", {})
        smiles = structures.get("canonical_smiles", "")
        return smiles if smiles else ""
    except Exception as e:
        logger.warning("Failed to fetch %s: %s", chembl_id, e)
        return ""


def fetch_smiles_batch(chembl_ids: list[str], cache_path: Path) -> dict[str, str]:
    """Fetch SMILES for a list of ChEMBL IDs, using cache.

    Parameters
    ----------
    chembl_ids : list of str
    cache_path : Path
        CSV file for caching results.

    Returns
    -------
    smiles_dict : dict mapping chembl_id → SMILES
    """
    # Load existing cache
    smiles_dict: dict[str, str] = {}
    if cache_path.exists():
        with open(cache_path) as f:
            reader = csv.DictReader(f)
            for row in reader:
                smiles_dict[row["chembl_id"]] = row["smiles"]
        logger.info("Loaded %d entries from cache %s", len(smiles_dict), cache_path)

    # Find IDs not in cache
    to_fetch = [cid for cid in chembl_ids if cid not in smiles_dict or not smiles_dict[cid]]
    logger.info("Need to fetch %d / %d ChEMBL IDs", len(to_fetch), len(chembl_ids))

    # Fetch in chunks
    for i in range(0, len(to_fetch), CHUNK_SIZE):
        chunk = to_fetch[i:i + CHUNK_SIZE]
        for chembl_id in chunk:
            smiles = fetch_smiles_from_chembl(chembl_id)
            smiles_dict[chembl_id] = smiles
            time.sleep(RATE_LIMIT_DELAY)

        # Save cache after each chunk
        _save_cache(smiles_dict, cache_path)
        logger.info("  Fetched %d/%d", min(i + CHUNK_SIZE, len(to_fetch)), len(to_fetch))

    return smiles_dict


def _save_cache(smiles_dict: dict[str, str], cache_path: Path) -> None:
    """Save SMILES dictionary to CSV."""
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    with open(cache_path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=["chembl_id", "smiles"])
        writer.writeheader()
        for cid, smiles in sorted(smiles_dict.items()):
            writer.writerow({"chembl_id": cid, "smiles": smiles})


def extract_chembl_ids_from_h5ad(h5ad_path: str) -> list[str]:
    """Extract unique ChEMBL IDs from SciPlex3 h5ad file."""
    adata = ad.read_h5ad(h5ad_path)
    chembl_ids = adata.obs["chembl-ID"].dropna().unique().tolist()
    # Clean up (some have multiple IDs separated by ;)
    all_ids = []
    for cid in chembl_ids:
        for part in str(cid).split(";"):
            part = part.strip()
            if part.startswith("CHEMBL"):
                all_ids.append(part)
    return sorted(set(all_ids))


def build_drug_name_to_smiles(
    h5ad_path: str, smiles_csv: Path
) -> dict[str, str]:
    """Build drug_name → SMILES mapping from SciPlex3 h5ad + ChEMBL cache.

    Parameters
    ----------
    h5ad_path : str
    smiles_csv : Path
        CSV with columns chembl_id, smiles

    Returns
    -------
    drug_smiles : dict mapping drug_name → SMILES
    """
    adata = ad.read_h5ad(h5ad_path)

    # Load SMILES cache
    smiles_by_chembl: dict[str, str] = {}
    if smiles_csv.exists():
        with open(smiles_csv) as f:
            reader = csv.DictReader(f)
            for row in reader:
                smiles_by_chembl[row["chembl_id"]] = row["smiles"]

    # Map drug names to SMILES
    drug_smiles: dict[str, str] = {}
    for _, row in adata.obs[["perturbation", "chembl-ID"]].drop_duplicates().iterrows():
        drug_name = row["perturbation"]
        chembl_id = row["chembl-ID"]
        if pd.isna(chembl_id):
            continue
        # Handle multiple IDs (take first)
        chembl_id = str(chembl_id).split(";")[0].strip()
        if chembl_id in smiles_by_chembl and smiles_by_chembl[chembl_id]:
            drug_smiles[drug_name] = smiles_by_chembl[chembl_id]

    return drug_smiles


def main():
    parser = argparse.ArgumentParser(description="Fetch ChEMBL SMILES for drug encoder")
    parser.add_argument("--chembl_ids", nargs="*", help="ChEMBL IDs to fetch")
    parser.add_argument("--h5ad", help="SciPlex3 h5ad path (auto-extract ChEMBL IDs)")
    parser.add_argument("--output", required=True, help="Output CSV path")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO)
    cache_path = Path(args.output)

    if args.h5ad:
        logger.info("Extracting ChEMBL IDs from %s", args.h5ad)
        chembl_ids = extract_chembl_ids_from_h5ad(args.h5ad)
        logger.info("Found %d unique ChEMBL IDs", len(chembl_ids))
    elif args.chembl_ids:
        chembl_ids = args.chembl_ids
    else:
        parser.error("Provide --chembl_ids or --h5ad")

    smiles_dict = fetch_smiles_batch(chembl_ids, cache_path)

    # Stats
    found = sum(1 for v in smiles_dict.values() if v)
    logger.info("Results: %d/%d SMILES found", found, len(smiles_dict))
    logger.info("Saved to %s", cache_path)


if __name__ == "__main__":
    main()