# AUTOGENERATED! DO NOT EDIT! File to edit: ../01_core.ipynb.

# %% auto 0
__all__ = [
    "to_ensembl_prefix",
    "get_ensembl_dataset_name",
    "get_cache_dir_path",
    "query_genes",
]

# %% ../01_core.ipynb 3
## helper functions
import logging
import numpy as np
from roux.lib.io import read_dict, to_table, read_table

from pathlib import Path

try:
    from dups.utils import get_cache_dir
except:
    from chrov.utils import get_cache_dir

## ensembl
## prefixes
# obtained from
# https://useast.ensembl.org/info/website/archives/index.html
to_ensembl_prefix = {
    # 113: 'sep2024.archive',
    112: "may2024.archive",
    111: "jan2024.archive",
    110: "jul2023.archive",
    109: "feb2023.archive",
    100: "apr2020.archive",
    101: "aug2020.archive",
    102: "nov2020.archive",
    103: "feb2021.archive",
    104: "may2021.archive",
    98: "sep2019.archive",
    93: "jul2018.archive",
    75: "feb2014.archive",
    "grch37": "grch37",
}

## BioMart


def get_ensembl_dataset_name(
    x: str,
) -> str:
    """Get the name of the Ensembl dataset.

    Args:
        x (str): species name.

    Returns:
        str: output.
    """
    l = x.lower().split(" ")
    assert len(l) == 2, l
    return f"{l[0][0]}{l[1]}_gene_ensembl"


def get_cache_dir_path(
    cache_dir_path=None,  ## force custim dir
    suffix="",
    source=None,  # e.g. biomart
    ensembl_release=None,
    species=None,
):
    ## cache file path
    if cache_dir_path is None:
        cache_dir_path = get_cache_dir()

    if species is not None:
        dataset_name = get_ensembl_dataset_name(species)
    else:
        dataset_name = None

    for s in [source, dataset_name, ensembl_release]:
        if s is not None:
            suffix += f"/{s}/"

    return f"{cache_dir_path}/{suffix}/"


def query_genes(
    species,
    release,
    chromosomes=None,
    cols=[
        "ensembl_gene_id",
        "external_gene_name",
        "chromosome_name",
        "start_position",
        "end_position",
        "strand",
    ],
    cols_contain=None,
    biotypes=["protein_coding"],
    verbose=False,
    cache_dir_path=None,
    force=False,
    filters={},
    kws_get_cache={},
    **kws_query,
):
    dataset_name = get_ensembl_dataset_name(species)
    from pybiomart import Server, Dataset

    host = f"http://{to_ensembl_prefix[release]}.ensembl.org/"
    server = Server(host=host)
    release = server["ENSEMBL_MART_ENSEMBL"].display_name.split(" ")[-1]
    logging.warning(f"{dataset_name} version: {release}")
    dataset = Dataset(name=dataset_name, host=host)
    try:
        df0 = dataset.list_attributes()
    except:
        logging.error(dataset_name)
    attributes = df0.log.query(expr=f"`name` in {cols}")["name"].tolist() + (
        []
        if cols_contain is None
        else df0.log.query(expr=f"`name`.str.contains('{cols_contain}')")[
            "name"
        ].tolist()
    )
    logging.warning(f"attributes: {attributes}")
    if chromosomes is None:
        ## list of chromosomes without scaffolds
        chrom_url = f"https://e{release}.rest.ensembl.org/info/assembly/{species.replace(' ', '_')}?content-type=application/json"
        logging.info(f"getting chroms from {chrom_url}")
        chromosomes = read_dict(chrom_url)["karyotype"]

    ## filters
    filters = {
        **{
            "chromosome_name": chromosomes,
            "biotype": biotypes,
        },
        **filters,
    }

    ## cache file path
    # pkg_name=get_pkg_name()

    if cache_dir_path is None:
        # from .utils import get_cache_dir
        # cache_dir_path = get_cache_dir(cache_dir_path)
        # from roux.lib.str import encode

        # outp = f"{cache_dir_path}/www.ensembl.org/biomart/{dataset_name}/{release}/{encode(dict(attributes=attributes,filters=filters),short=True)}.pqt"

        cache_dir_path = get_cache_dir_path(
            species=species,
            ensembl_release=release,
            source="www.ensembl.org/biomart",
            **kws_get_cache,
        )
        from roux.lib.str import encode

        outp = f"{cache_dir_path}{encode(dict(cols=cols), short=True)}.pqt"

    if not Path(outp).exists() or force:
        df1 = dataset.query(
            attributes=np.unique(attributes),
            filters=filters,
            only_unique=True,
            **kws_query,
        )
        if "Chromosome/scaffold name" in df1:
            df1 = df1.astype({"Chromosome/scaffold name": str})
        to_table(df1, outp)
        logging.info(f"cache path: {outp}")
    else:
        logging.warning(f"read from cache: {outp}")
        df1 = read_table(outp)
    return df1.rename(
        columns={
            "Gene stable ID": "gene id",
            "Chromosome/scaffold name": "chromosome",
            "Gene name": "gene symbol",
            "Gene description": "gene description",
            "Ensembl Family Description": "gene family description",
            # 'Gene Synonym':'gene synonyms',
            "Transcript stable ID": "transcript id",
            "Protein stable ID": "protein id",
            "UniProtKB/Swiss-Prot ID": "protein id (UniProtKB/Swiss-Prot)",
            "Peptide": "protein sequence",
            "Gene start (bp)": "gene start",
            "Gene end (bp)": "gene end",
            "Strand": "gene strand",
            "Gene type": "biotype",
        },
        errors="ignore",
    )
