utils.py

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import os
import json
from glob import glob
import polars as pl

"""
Utility functions for interacting with the ENCODE and EpiRR APIs.
Fields fetched adhere to IHEC standards where possible. 
Needed, but missing fields can be left empty and we fill everything we get:
CHIP_ANTIBODY_PROVIDER - The name of the company, laboratory or person that provided the antibody. Leave empty for 'ChIP-Seq Input'.
CHIP_ANTIBODY_CATALOG - The catalog from which the antibody was purchased. Leave empty for 'ChIP-Seq Input'.
CHIP_ANTIBODY_LOT - The lot identifier of the antibody. Leave empty for 'ChIP-Seq Input'.
"""


# Create a requests session that retries on certain HTTP status codes
retry_strategy = Retry(
    total=5,
    backoff_factor=0.5,
    status_forcelist=(500, 502, 503, 504),
    allowed_methods=("GET",),
    raise_on_status=False,
)
adapter = HTTPAdapter(max_retries=retry_strategy)

retry_session = requests.Session()
retry_session.mount("https://", adapter)
retry_session.mount("http://", adapter)

encode_base_url = "https://www.encodeproject.org"


def get_antibody_from_experiment(experiment_id: str) -> str:
    """Get the antibody from the ENCODE experiment ID.
    Args:
        experiment_id (str): The ENCODE experiment ID.
    Returns:
        str: The antibody URL or raise an error if not found.
    """
    url = f"{encode_base_url}/experiments/{experiment_id}/?format=json"
    response = retry_session.get(url, headers={"Accept": "application/json"})
    response.raise_for_status()
    data = response.json()
    # directly get all antibody info from this data
    antibody_list = []
    replicate_data = data["replicates"]
    for rep in replicate_data:
        antibody = rep.get("antibody", {})
        antibody_list.append(
            {
                "primary_id": experiment_id,
                "CHIP_ANTIBODY_PROVIDER": antibody.get("source"),
                "CHIP_ANTIBODY_CATALOG": antibody.get("product_id"),
                "CHIP_ANTIBODY_LOT": antibody.get("lot_id"),
            }
        )
    return antibody_list


epirr_base_url = "https://www.ebi.ac.uk/epirr/api/v1/epigenome?accession="


def get_epirr_experiment_metadata(epirr_id: str) -> list:
    """Get experiment metadata from EpiRR for a given EpiRR ID.
    Args:
        epirr_id (str): The EpiRR ID.
    Returns:
        a list of dicts: The experiment metadata.
    """
    url = f"{epirr_base_url}{epirr_id}"
    response = retry_session.get(url)
    response.raise_for_status()
    data = response.json()
    local_name = data.get("local_name", None)
    experimental_data = data.get("experimental_data", [])
    for exp_data in experimental_data:
        exp_data["local_name"] = local_name
    return experimental_data


def get_deep_metadata(experiment_url: str) -> dict:
    """Fetch antibody information from local files or DEEP Portal given an antibody URL.
    Args:
        experiment_url (str): The antibody URL (or local file name).
    Returns:
        dict: The antibody information as a dictionary.
    """
    # first check if experiment_url is a local file name or an url
    data = {}
    if not experiment_url.startswith("http"):
        # read the whole tsv file as a dictionary with Key\tvalue per line
        with open(experiment_url, "r") as f:
            data = {line.split("\t")[0]: line.split("\t")[1].strip() for line in f}
            # split away '_emd.tsv' from both sides and compare
            exp_id = data.get("EXPERIMENT_ID").replace("_emd.tsv", "")
            inferred_exp_id = os.path.basename(experiment_url).replace("_emd.tsv", "")
            assert exp_id == inferred_exp_id  # sanity check
            data["experiment_url"] = experiment_url
            return data
            return {
                k: data.get(k)
                for k in (
                    "CHIP_ANTIBODY_PROVIDER",
                    "CHIP_ANTIBODY_CATALOG",
                    "CHIP_ANTIBODY_LOT",
                    "experiment_url",
                )
            }
    else:
        return {"experiment_url": experiment_url}
        response = antibody_session.get(
            experiment_url, headers={"Accept": "application/json"}
        )
        response.raise_for_status()
        data = response.json().get("experiments", [])[0]
        assert (
            data["EXPERIMENT_ID"]
            == os.path.splitext(os.path.basename(experiment_url))[0]
        )  # sanity check
        data = data.get("attributes", {})
        data["experiment_url"] = [experiment_url]
        return data
        return {
            k: ";".join(data.get(k, []))
            for k in (
                "CHIP_ANTIBODY_PROVIDER",
                "CHIP_ANTIBODY_CATALOG",
                "CHIP_ANTIBODY_LOT",
                "experiment_url",
            )
        }


def get_deep_file_name(local_name: str, experiment_type: str) -> str:
    """Get the deep file name from EpiRR for a given local_name.
    Args:
        local_name (str): The local_name.
    Returns:
        str: The DEEP file name or None if not found.
    """
    # use case:
    deep_experiment_type = experiment_type
    if experiment_type == "PBAT":
        deep_experiment_type = "WGBS"
    elif experiment_type == "mRNA-Seq":
        deep_experiment_type = "mRNA"
    elif experiment_type == "total-RNA-Seq":
        deep_experiment_type = "tRNA"

    return_list = []
    for file in glob(
        f"deep_experiments/{local_name}/{local_name}_{deep_experiment_type}*.tsv"
    ):
        curr_data_dict = {
            "LOCAL_NAME": local_name,
            "experiment_type": experiment_type,
            "file": file,
        }
        curr_data_dict.update(get_deep_metadata(file))
        return_list.append(curr_data_dict)
    return return_list


import xml.etree.ElementTree as ET


def flatten_tree(xml_file: str) -> dict:
    """Flatten an XML tree into a list of dictionaries.
    Processes all EXPERIMENT_SET elements found anywhere in the XML (not just the first).
    Args:
        xml_file (str): The path to the XML file.
    Returns:
        list: A list of flattened dictionaries, one per experiment.
    """
    tree = ET.parse(xml_file)

    def flatten_element(
        element: ET.Element, parent_key: str = "", sep: str = "_"
    ) -> dict:
        items = {}
        for child in element:
            new_key = f"{parent_key}{sep}{child.tag}" if parent_key else child.tag
            if child.tag == "EXPERIMENT_ATTRIBUTE":
                # guard missing subelements
                tag_el = child.find("TAG")
                val_el = child.find("VALUE")
                attr_tag = tag_el.text if tag_el is not None else None
                attr_value = val_el.text if val_el is not None else None
                if attr_tag:
                    if attr_tag in items:
                        # append to existing value with
                        items[f"{attr_tag}_2"] = attr_value
                        # still print a warning
                        print(
                            f"Warning: Duplicate EXPERIMENT_ATTRIBUTE tag: {attr_tag} in {xml_file}"
                        )
                    else:
                        # assert (
                        #     attr_tag not in items
                        # ), f"Duplicate EXPERIMENT_ATTRIBUTE tag: {attr_tag} in {xml_file}"
                        items[attr_tag] = attr_value
                    # assert (
                    #     attr_tag not in items
                    # ), f"Duplicate EXPERIMENT_ATTRIBUTE tag: {attr_tag}"
                    # items[attr_tag] = attr_value
                continue
            if len(child):
                items.update(flatten_element(child, new_key, sep=sep))
            else:
                assert not new_key in items, f"Duplicate key found: {new_key}"
                items[new_key] = child.text
        return items

    root = tree.getroot()
    flattened_data = []
    # Find all EXPERIMENT_SET elements anywhere in the tree also itself
    experiment_sets = root.findall("EXPERIMENT_SET")
    if root.tag == "EXPERIMENT_SET":
        experiment_sets.append(root)
    for experiment_set in experiment_sets:
        # create one list per experiment_set
        for child in experiment_set:
            record = flatten_element(child)
            flattened_data.append(record)
    return {xml_file: flattened_data}


def load_and_flatten_json_file(json_file: str) -> dict:
    """Load and flatten a JSON file that has a structure similar to ENA/EGA experiment metadata,
    processing all EXPERIMENT_SET occurrences found anywhere in the file.
    Args:
        f (str): The path to the JSON file.
    Returns:
        list[dict]: Flattened list of experiments.
    """

    def flatten_element(d, parent_key="") -> dict:
        """Recursively flatten nested dicts/lists and handle EXPERIMENT_ATTRIBUTE as in flatten_tree."""
        items = {}

        if not isinstance(d, dict):
            assert not parent_key in items, f"Duplicate key found: {parent_key}"
            items[parent_key] = d
            return items

        for k, v in d.items():
            new_key = f"{parent_key}_{k}" if parent_key else k
            if isinstance(v, dict):
                update_dict = flatten_element(v, new_key)
                assert not (
                    update_dict.keys() & items.keys()
                ), f"Duplicate keys found when flattening dict at key: {new_key}"
                items.update(update_dict)
            elif isinstance(v, list):
                if k == "EXPERIMENT_ATTRIBUTE" or k == "EXPERIMENT_ATTRIBUTES":
                    for attr in v:
                        tag = attr.get("TAG")
                        value = attr.get("VALUE")
                        if tag and value:
                            if tag in items:
                                # append to existing value with
                                items[f"{tag}_2"] = value
                                # still print a warning
                                print(
                                    f"Warning: Duplicate EXPERIMENT_ATTRIBUTE tag: {tag} in {json_file}"
                                )
                            else:
                                # assert (
                                #     tag not in items
                                # ), f"Duplicate EXPERIMENT_ATTRIBUTE tag: {tag} in {json_file}"
                                items[tag] = value
                else:
                    # join scalar lists, otherwise flatten items
                    # if all(isinstance(x, (str, int, float)) for x in v):
                    #     items[new_key] = ";".join(map(str, v))
                    # else:
                    for i, item in enumerate(v):
                        items.update(flatten_element(item, f"{new_key}_{i}"))
            else:
                assert not new_key in items, f"Duplicate key found: {new_key}"
                items[new_key] = v
        return items

    def find_experiment_sets(obj):
        """Recursively find and yield EXPERIMENT_SET objects anywhere in obj.
        Yields the raw value of each EXPERIMENT_SET found.
        """

        if isinstance(obj, dict):
            for k, v in obj.items():
                if k == "EXPERIMENT_SET":
                    yield v
                else:
                    yield from find_experiment_sets(v)
        elif isinstance(obj, list):
            for item in obj:
                yield from find_experiment_sets(item)

    with open(json_file, "r") as file:
        data = json.load(file)

        flattened_records = []
        # Walk through top-level keys (e.g., accession keys)
        for top_id, top_entry in data.items():
            if top_id == "EXPERIMENT_SET":
                experiment_sets = [top_entry]
            else:
                experiment_sets = find_experiment_sets(top_entry)
            # find all experiment_set occurrences under this top entry
            for experiment_set in experiment_sets:
                # e.g. { 'EXPERIMENT': {...} } or keyed by accessions
                for experiment in experiment_set.values():
                    flat_experiment = flatten_element(experiment)
                    flattened_records.append(flat_experiment)

        return {json_file: flattened_records}


def merge_columns(df: pl.DataFrame, groups: dict) -> pl.DataFrame:
    """Merge case-insensitive duplicate columns in a Polars DataFrame,
    preserving all unique non-null values as lists.
    Args:
        df (pl.DataFrame): The input Polars DataFrame.
        groups (dict): A dictionary mapping group names to lists of column names to merge.
    Returns:
        pl.DataFrame: The DataFrame with merged columns.
    """

    for group_name, cols in groups.items():
        # check if anything is null in these columns and if so fill with fill_null

        for col in cols:
            if df[col].is_null().any():
                # if type of column is list, fill with empty list, else with empty string
                fill_null = [] if str(df[col].dtype).startswith("List") else ""
                df = df.with_columns(pl.col(col).fill_null(fill_null).alias(col))

        df = df.with_columns(pl.concat_list(pl.col(cols)).alias(group_name))
        df = df.drop([c for c in cols if c != group_name])
    return df