Source code for tooluniverse.epigenomics_tool

# epigenomics_tool.py
"""
Epigenomics and methylation analysis tools for ToolUniverse.

Integrates data from:
- ENCODE Project (histone ChIP-seq, WGBS methylation, ATAC-seq, DNase-seq, annotations)
- UCSC Genome Browser (CpG islands, ENCODE4 cCREs, TF binding clusters)
- NCBI GEO (methylation array datasets, ChIP-seq datasets)
- Ensembl Regulatory Build (regulatory features, enhancers, promoters)

Optional auth: set ``NCBI_API_KEY`` to lift the NCBI E-utilities rate limit
from 3 req/sec to 10 req/sec (free, register at
https://www.ncbi.nlm.nih.gov/account/settings/). The tool works without it
but the GEO_* endpoints will burst-throttle (HTTP 429) under load.
"""

import json
import os
import random
import time
import requests
from typing import Dict, Any, Optional
from .base_tool import BaseTool
from .ncbi_eutils_tool import esearch_query_disclosure
from .tool_registry import register_tool

ENCODE_BASE_URL = "https://www.encodeproject.org"
UCSC_API_URL = "https://api.genome.ucsc.edu"
NCBI_EUTILS_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
ENSEMBL_REST_URL = "https://rest.ensembl.org"

# Retry on transient upstream issues. NCBI E-utils 429s on burst; ENCODE/UCSC
# occasionally 503. Exponential backoff with jitter.
_RETRY_STATUS_CODES = {408, 429, 500, 502, 503, 504}
_MAX_RETRIES = 3


def _encode_empty_result_note(biosample: Any = None) -> str:
    """Explain a genuinely empty ENCODE result, hinting at ontology spelling."""
    note = (
        "No matching ENCODE experiments. The filters are applied as given -- "
        "this is a real zero-result answer, not a dropped filter."
    )
    if biosample:
        note += (
            f" ENCODE requires exact ontology names, so check biosample='{biosample}' "
            "(e.g. 'K562', 'HepG2', 'liver', 'brain', 'breast epithelium', 'MCF-7'; "
            "'breast' must be spelled 'breast epithelium' or 'mammary gland'). "
            "Use GEO_search_chipseq_datasets for disease-based searches."
        )
    return note


def _inject_ncbi_api_key(params: Dict[str, Any]) -> Dict[str, Any]:
    """If NCBI_API_KEY is set, append it to the params dict. 3→10 req/sec.

    Returns the same params dict (mutated) for chaining.
    """
    key = os.environ.get("NCBI_API_KEY")
    if key:
        params["api_key"] = key
    return params


def _request_with_backoff(url: str, *, timeout: int, **kwargs) -> requests.Response:
    """GET ``url`` with exponential-backoff retry on 429/5xx.

    Re-raises the final response's HTTPError after exhausting retries.
    Sleeps respect a ``Retry-After`` header when the server provides one.
    """
    last_resp = None
    for attempt in range(_MAX_RETRIES + 1):
        last_resp = requests.get(url, timeout=timeout, **kwargs)
        if last_resp.status_code not in _RETRY_STATUS_CODES:
            return last_resp
        if attempt == _MAX_RETRIES:
            break
        # Honour Retry-After when present, otherwise exp-backoff + jitter
        retry_after = last_resp.headers.get("Retry-After")
        if retry_after and retry_after.isdigit():
            delay = min(float(retry_after), 30.0)
        else:
            delay = min(0.5 * (2**attempt) + random.uniform(0, 0.25), 8.0)
        time.sleep(delay)
    last_resp.raise_for_status()
    return last_resp


def _geo_search_metadata(query, term, organism, esearch):
    """Response metadata for a GEO dataset search, including GEO's own report
    that it did not run the query as asked.

    Fix-54A-1: ``search_term`` reported the term as SUBMITTED, but GEO's
    esearch drops phrases it cannot match and answers the remainder, so the
    field naming the executed query was the one field stating it wrongly.
    Measured: query "breast cancer zzzqqqxyz nonexistentterm12345" returned
    total 205972 -- identical to "breast cancer" alone -- with both nonsense
    phrases in errorlist.phrasesnotfound and nothing said to the caller.

    Shared by the four ``_geo_*_search`` methods, which differ only in the
    ``term`` they build; stated once so the disclosure cannot be added to three
    of them and forgotten in the fourth.
    """
    metadata = {
        "source": "NCBI GEO (ncbi.nlm.nih.gov/geo)",
        "query": query,
        "search_term": term,
        "organism": organism,
    }
    disclosure = esearch_query_disclosure(esearch, source="GEO")
    if disclosure:
        metadata["query_disclosure"] = disclosure
    return metadata




[docs]
@register_tool("EpigenomicsTool")
class EpigenomicsTool(BaseTool):
    """
    Tool for epigenomics and methylation analysis across multiple databases.

    Supports:
    - ENCODE histone ChIP-seq, methylation (WGBS/RRBS), chromatin accessibility
    - ENCODE annotations (cCREs, chromatin states)
    - GEO methylation and ChIP-seq dataset search
    - Ensembl regulatory features

    No authentication required.
    """



[docs]
    def __init__(self, tool_config: Dict[str, Any]):
        super().__init__(tool_config)
        self.timeout = tool_config.get("timeout", 30)
        fields = tool_config.get("fields", {})
        self.endpoint = fields.get("endpoint", "histone_chipseq")



    _ORGANISM_ALIASES = {
        "human": "Homo sapiens",
        "homo sapiens": "Homo sapiens",
        "mouse": "Mus musculus",
        "mus musculus": "Mus musculus",
        "rat": "Rattus norvegicus",
        "zebrafish": "Danio rerio",
        "fly": "Drosophila melanogaster",
        "worm": "Caenorhabditis elegans",
    }



[docs]
    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Execute the epigenomics API call."""
        # Normalize organism aliases to scientific names required by ENCODE API
        if "organism" in arguments:
            org = arguments["organism"].lower()
            if org in self._ORGANISM_ALIASES:
                arguments = dict(arguments, organism=self._ORGANISM_ALIASES[org])
        try:
            result = self._dispatch(arguments)
            if isinstance(result, dict) and "status" not in result:
                if "error" in result:
                    return {"status": "error", **result}
                return {"status": "success", **result}
            return result
        except requests.exceptions.Timeout:
            return {
                "status": "error",
                "error": f"API request timed out after {self.timeout}s",
            }
        except requests.exceptions.ConnectionError:
            return {
                "status": "error",
                "error": "Failed to connect to API. Check network connectivity.",
            }
        except requests.exceptions.HTTPError as e:
            status = e.response.status_code if e.response is not None else "unknown"
            return {"status": "error", "error": f"API HTTP error: {status}"}
        except Exception as e:
            return {"status": "error", "error": f"Unexpected error: {str(e)}"}





[docs]
    def _dispatch(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Route to appropriate endpoint based on config."""
        if self.endpoint == "histone_chipseq":
            return self._encode_histone_search(arguments)
        elif self.endpoint == "methylation":
            return self._encode_methylation_search(arguments)
        elif self.endpoint == "chromatin_accessibility":
            return self._encode_chromatin_accessibility_search(arguments)
        elif self.endpoint == "annotations":
            return self._encode_annotations_search(arguments)
        elif self.endpoint == "chromatin_state":
            return self._encode_chromatin_state_search(arguments)
        elif self.endpoint == "geo_methylation_search":
            return self._geo_methylation_search(arguments)
        elif self.endpoint == "geo_chipseq_search":
            return self._geo_chipseq_search(arguments)
        elif self.endpoint == "geo_dataset_details":
            return self._geo_dataset_details(arguments)
        elif self.endpoint == "ensembl_regulatory":
            return self._ensembl_regulatory_features(arguments)
        elif self.endpoint == "geo_rnaseq_search":
            return self._geo_rnaseq_search(arguments)
        elif self.endpoint == "geo_atacseq_search":
            return self._geo_atacseq_search(arguments)
        elif self.endpoint == "encode_rnaseq":
            return self._encode_rnaseq_search(arguments)
        elif self.endpoint == "encode_hic":
            return self._encode_hic_search(arguments)
        elif self.endpoint == "encode_microrna":
            return self._encode_microrna_search(arguments)
        else:
            return {"status": "error", "error": f"Unknown endpoint: {self.endpoint}"}



    # =========================================================================
    # ENCODE Search Tools
    # =========================================================================






[docs]
    @staticmethod
    def _is_histone_mark(target: str) -> bool:
        """Return True if target looks like a histone modification (e.g. H3K27ac, H3K4me3)."""
        import re

        return bool(re.match(r"^H[1-4][A-Za-z0-9]", target))


















    # =========================================================================
    # GEO Search Tools
    # =========================================================================



[docs]
    def _geo_esearch(self, term: str, limit: int = 20) -> Dict[str, Any]:
        """Search GEO datasets via NCBI E-utilities.

        Uses NCBI_API_KEY from env when set (lifts the rate cap from 3 to
        10 req/sec) and retries on 429/5xx with exponential backoff.
        """
        url = f"{NCBI_EUTILS_URL}/esearch.fcgi"
        params = _inject_ncbi_api_key(
            {
                "db": "gds",
                "term": term,
                "retmax": min(int(limit), 100),
                "retmode": "json",
            }
        )
        response = _request_with_backoff(url, params=params, timeout=self.timeout)
        response.raise_for_status()
        return response.json()





[docs]
    def _geo_esummary(self, ids: list) -> Dict[str, Any]:
        """Get summary for GEO dataset IDs via NCBI E-utilities.

        Uses NCBI_API_KEY when set + retries on 429/5xx.
        """
        if not ids:
            return {"result": {}}
        url = f"{NCBI_EUTILS_URL}/esummary.fcgi"
        params = _inject_ncbi_api_key(
            {
                "db": "gds",
                "id": ",".join(str(i) for i in ids),
                "retmode": "json",
            }
        )
        response = _request_with_backoff(url, params=params, timeout=self.timeout)
        response.raise_for_status()
        return response.json()











[docs]
    def _geo_dataset_details(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Get detailed metadata for a GEO dataset."""
        geo_id = arguments.get("geo_id", "")
        if not geo_id:
            return {
                "status": "error",
                "error": "geo_id parameter is required (e.g., '200291249')",
            }

        summary_result = self._geo_esummary([geo_id])
        result = summary_result.get("result", {})
        uid_data = result.get(str(geo_id), {})

        if not isinstance(uid_data, dict) or "accession" not in uid_data:
            return {
                "status": "error",
                "error": f"Dataset with ID '{geo_id}' not found in GEO",
            }

        ftplink = uid_data.get("ftplink", "")
        suppfile = uid_data.get("suppfile", "")
        supp_data = []
        if ftplink:
            supp_data.append(ftplink)
        if suppfile:
            supp_data.append(suppfile)

        return {
            "status": "success",
            "data": {
                "accession": uid_data.get("accession", ""),
                "title": uid_data.get("title", ""),
                "summary": uid_data.get("summary", ""),
                "experiment_type": uid_data.get("gdstype"),
                "platform": uid_data.get("gpl"),
                "organism": uid_data.get("taxon", ""),
                "n_samples": uid_data.get("n_samples", 0),
                "date_published": uid_data.get("pdat"),
                "supplementary_data": supp_data if supp_data else None,
            },
            "metadata": {
                "source": "NCBI GEO (ncbi.nlm.nih.gov/geo)",
                "geo_id": geo_id,
            },
        }












    # =========================================================================
    # Ensembl Regulatory Features
    # =========================================================================



[docs]
    def _ensembl_regulatory_features(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Get Ensembl regulatory features for a genomic region."""
        species = arguments.get("species", "homo_sapiens")
        chrom = arguments.get("chrom", "")
        start = arguments.get("start")
        end = arguments.get("end")

        if not chrom or start is None or end is None:
            return {
                "status": "error",
                "error": "chrom, start, and end parameters are required",
            }

        # Ensure region is not too large (max 5Mb)
        if end - start > 5000000:
            return {
                "status": "error",
                "error": "Region too large. Maximum region size is 5 Mb.",
            }

        url = (
            f"{ENSEMBL_REST_URL}/overlap/region/{species}/{chrom}:{start}-{end}"
            f"?feature=regulatory;content-type=application/json"
        )
        # Ensembl REST API can be slow - use 90s timeout
        response = requests.get(url, timeout=max(self.timeout, 90))
        response.raise_for_status()
        raw = response.json()

        features = []
        for feat in raw:
            features.append(
                {
                    "id": feat.get("id", ""),
                    "description": feat.get("description", ""),
                    "feature_type": feat.get("feature_type", ""),
                    "start": feat.get("start"),
                    "end": feat.get("end"),
                    "strand": feat.get("strand", 0),
                    "seq_region_name": feat.get("seq_region_name", ""),
                }
            )

        return {
            "status": "success",
            "data": {
                "species": species,
                "region": f"{chrom}:{start}-{end}",
                "feature_count": len(features),
                "regulatory_features": features,
            },
            "metadata": {
                "source": "Ensembl Regulatory Build (rest.ensembl.org)",
                "species": species,
                "region": f"{chrom}:{start}-{end}",
            },
        }



    # =========================================================================
    # ENCODE Hi-C / 3D Genome Tools
    # =========================================================================




    # =========================================================================
    # ENCODE microRNA-seq Tools
    # =========================================================================









[docs]
@register_tool("UCSCEpigenomicsTool")
class UCSCEpigenomicsTool(BaseTool):
    """
    UCSC Genome Browser epigenomics-specific tools.

    Provides access to:
    - CpG island annotations
    - ENCODE4 candidate cis-Regulatory Elements (cCREs)
    - Transcription Factor binding site clusters

    No authentication required.
    """



[docs]
    def __init__(self, tool_config: Dict[str, Any]):
        super().__init__(tool_config)
        self.timeout = tool_config.get("timeout", 30)
        fields = tool_config.get("fields", {})
        self.endpoint = fields.get("endpoint", "cpg_islands")





[docs]
    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Execute the UCSC epigenomics API call."""
        try:
            return self._dispatch(arguments)
        except requests.exceptions.Timeout:
            return {
                "status": "error",
                "error": f"UCSC API request timed out after {self.timeout}s",
            }
        except requests.exceptions.ConnectionError:
            return {"status": "error", "error": "Failed to connect to UCSC API."}
        except requests.exceptions.HTTPError as e:
            status = e.response.status_code if e.response is not None else "unknown"
            return {"status": "error", "error": f"UCSC API HTTP error: {status}"}
        except Exception as e:
            return {"status": "error", "error": f"Unexpected error: {str(e)}"}





[docs]
    def _dispatch(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Route to appropriate endpoint."""
        if self.endpoint == "cpg_islands":
            return self._get_cpg_islands(arguments)
        elif self.endpoint == "encode_ccres":
            return self._get_encode_ccres(arguments)
        elif self.endpoint == "tf_binding":
            return self._get_tf_binding_clusters(arguments)
        else:
            return {"status": "error", "error": f"Unknown endpoint: {self.endpoint}"}





[docs]
    def _ucsc_get_track(
        self, genome: str, track: str, chrom: str, start: int, end: int
    ) -> Dict[str, Any]:
        """Helper to fetch UCSC track data."""
        url = (
            f"{UCSC_API_URL}/getData/track"
            f"?genome={genome}&track={track}&chrom={chrom}&start={start}&end={end}"
        )
        response = requests.get(url, timeout=self.timeout)
        response.raise_for_status()
        return response.json()





[docs]
    def _get_cpg_islands(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Get CpG island annotations for a genomic region."""
        genome = arguments.get("genome", "hg38")
        chrom = arguments.get("chrom", "")
        start = arguments.get("start")
        end = arguments.get("end")

        if not chrom or start is None or end is None:
            return {
                "status": "error",
                "error": "chrom, start, and end parameters are required",
            }

        raw = self._ucsc_get_track(genome, "cpgIslandExt", chrom, start, end)
        items = raw.get("cpgIslandExt", [])
        if not isinstance(items, list):
            items = []

        cpg_islands = []
        for item in items:
            cpg_islands.append(
                {
                    "chrom": item.get("chrom", ""),
                    "chromStart": item.get("chromStart"),
                    "chromEnd": item.get("chromEnd"),
                    "name": item.get("name", ""),
                    "length": item.get("length", 0),
                    "cpgNum": item.get("cpgNum", 0),
                    "gcNum": item.get("gcNum", 0),
                    "perCpg": item.get("perCpg", 0),
                    "perGc": item.get("perGc", 0),
                    "obsExp": item.get("obsExp", 0),
                }
            )

        return {
            "status": "success",
            "data": {
                "genome": genome,
                "region": f"{chrom}:{start}-{end}",
                "cpg_island_count": len(cpg_islands),
                "cpg_islands": cpg_islands,
            },
            "metadata": {
                "source": "UCSC Genome Browser (api.genome.ucsc.edu)",
                "track": "cpgIslandExt",
                "genome": genome,
            },
        }





[docs]
    def _get_encode_ccres(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Get ENCODE4 candidate cis-Regulatory Elements for a genomic region."""
        genome = arguments.get("genome", "hg38")
        chrom = arguments.get("chrom", "")
        start = arguments.get("start")
        end = arguments.get("end")

        if not chrom or start is None or end is None:
            return {
                "status": "error",
                "error": "chrom, start, and end parameters are required",
            }

        raw = self._ucsc_get_track(genome, "cCREregistry", chrom, start, end)
        items = raw.get("cCREregistry", [])
        if not isinstance(items, list):
            items = []

        ccres = []
        for item in items:
            ccres.append(
                {
                    "name": item.get("name", ""),
                    "chrom": item.get("chrom", ""),
                    "chromStart": item.get("chromStart"),
                    "chromEnd": item.get("chromEnd"),
                    "cCRE_class": item.get("cCRE_class", ""),
                    "DNase_maxZ": item.get("DNase_maxZ", 0),
                    "H3K4me3_maxZ": item.get("H3K4me3_maxZ", 0),
                    "H3K27ac_maxZ": item.get("H3K27ac_maxZ", 0),
                    "CTCF_maxZ": item.get("CTCF_maxZ", 0),
                }
            )

        return {
            "status": "success",
            "data": {
                "genome": genome,
                "region": f"{chrom}:{start}-{end}",
                "ccre_count": len(ccres),
                "ccres": ccres,
            },
            "metadata": {
                "source": "UCSC Genome Browser / ENCODE4 (api.genome.ucsc.edu)",
                "track": "cCREregistry",
                "genome": genome,
            },
        }





[docs]
    def _get_tf_binding_clusters(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Get TF binding site clusters from ENCODE3."""
        genome = arguments.get("genome", "hg38")
        chrom = arguments.get("chrom", "")
        start = arguments.get("start")
        end = arguments.get("end")

        if not chrom or start is None or end is None:
            return {
                "status": "error",
                "error": "chrom, start, and end parameters are required",
            }

        raw = self._ucsc_get_track(genome, "encRegTfbsClustered", chrom, start, end)
        items = raw.get("encRegTfbsClustered", [])
        if not isinstance(items, list):
            items = []

        tf_clusters = []
        for item in items:
            tf_clusters.append(
                {
                    "name": item.get("name", ""),
                    "chrom": item.get("chrom", ""),
                    "chromStart": item.get("chromStart"),
                    "chromEnd": item.get("chromEnd"),
                    "score": item.get("score", 0),
                    "sourceCount": item.get("sourceCount", 0),
                }
            )

        return {
            "status": "success",
            "data": {
                "genome": genome,
                "region": f"{chrom}:{start}-{end}",
                "tf_cluster_count": len(tf_clusters),
                "tf_clusters": tf_clusters,
            },
            "metadata": {
                "source": "UCSC Genome Browser / ENCODE3 (api.genome.ucsc.edu)",
                "track": "encRegTfbsClustered",
                "genome": genome,
                "description": "340 TFs across 129 cell types",
            },
        }