import os

import re

import traceback

from abc import abstractmethod

from collections import defaultdict, OrderedDict

from typing import Union, List, Callable, Dict, Tuple, Optional, Iterable

import numpy as np

import pandas as pd

import tqdm

from Bio import SeqIO

from Bio.SeqFeature import ExactPosition

from dask import dataframe as dd

from logzero import logger

from pyfaidx import Fasta

from six.moves import intern

import openomics

from openomics.io.read_gtf import read_gtf

from .base import Database

from ..io.files import select_files_with_ext

from ..transforms.agg import get_agg_func

from ..transforms.df import drop_duplicate_columns

__all__ = ['GENCODE', 'UniProt', 'MirBase', 'RNAcentral']

SEQUENCE_COL = 'sequence'

class SequenceDatabase(Database):

    """Provides a series of methods to extract sequence data from

    SequenceDataset.

    """

    def __init__(self, **kwargs):

        """

        Args:

            **kwargs:

        """

        super().__init__(**kwargs)

        self.close()

    @abstractmethod

    def load_sequences(self, fasta_file: str, index=None, keys: Union[pd.Index, List[str]] = None, blocksize=None) \

        -> pd.DataFrame:

        """Returns a pandas DataFrame containing the fasta sequence entries.

        With a column named 'sequence'.

        Args:

            index ():

            fasta_file (str): path to the fasta file, usually as

                self.file_resources[<file_name>]

            keys (pd.Index): list of keys to

            blocksize:

        """

        raise NotImplementedError

    @abstractmethod

    def get_sequences(self, index: str, omic: str, agg: str, **kwargs) -> Union[pd.Series, Dict]:

        """Returns a dictionary where keys are 'index' and values are

        sequence(s).

        Args:

            index (str): {"gene_id", "gene_name", "transcript_id",

                "transcript_name"}

            omic (str): {"lncRNA", "microRNA", "messengerRNA"}

            agg (str): {"all", "shortest", "longest"}

            **kwargs: any additional argument to pass to

                SequenceDataset.get_sequences()

        """

        raise NotImplementedError

    @staticmethod

    def aggregator_fn(agg: Union[str, Callable] = None) -> Callable:

        """Returns a function used aggregate a list of sequences from a groupby

        on a given key.

        Args:

            agg: One of ("all", "shortest", "longest", "first"), default "all". If "all",

                then return a list of sequences.

        """

        if agg == "all":

            agg_func = lambda x: list(x) if not isinstance(x, str) else x

        elif agg == "shortest":

            agg_func = lambda x: min(x, key=len) if isinstance(x, list) else x

        elif agg == "longest":

            agg_func = lambda x: max(x, key=len) if isinstance(x, list) else x

        elif agg == 'first':

            agg_func = lambda x: x[0] if isinstance(x, list) else x

        elif callable(agg):

            return agg

        else:

            raise Exception(

                "agg_sequences argument must be one of {'all', 'shortest', 'longest'}"

            )

        return agg_func

class GENCODE(SequenceDatabase):

    """Loads the GENCODE database from https://www.gencodegenes.org/ .

    Default path: ftp://ftp.ebi.ac.uk/pub/databases/gencode/Gencode_human/release_32/ .

    Default file_resources: {

        "basic.annotation.gtf": "gencode.v32.basic.annotation.gtf.gz",

        "long_noncoding_RNAs.gtf": "gencode.v32.long_noncoding_RNAs.gtf.gz",

        "lncRNA_transcripts.fa": "gencode.v32.lncRNA_transcripts.fa.gz",

        "transcripts.fa": "gencode.v32.transcripts.fa.gz",

    }

    """

    def __init__(

        self,

        path="ftp://ftp.ebi.ac.uk/pub/databases/gencode/Gencode_human/release_32/",

        file_resources=None,

        col_rename=None,

        blocksize=0,

        remove_version_num=False,

        **kwargs

    ):

        """

        Args:

            path:

            file_resources:

            col_rename:

            blocksize:

            remove_version_num (bool): Whether to drop the version number on the

                ensembl ID.

        """

        if file_resources is None:

            file_resources = {

                "basic.annotation.gtf.gz": "gencode.v32.basic.annotation.gtf.gz",

                "long_noncoding_RNAs.gtf.gz": "gencode.v32.long_noncoding_RNAs.gtf.gz",

                "lncRNA_transcripts.fa.gz": "gencode.v32.lncRNA_transcripts.fa.gz",

                "transcripts.fa.gz": "gencode.v32.transcripts.fa.gz",

            }

        self.remove_version_num = remove_version_num

        super().__init__(path=path, file_resources=file_resources, col_rename=col_rename, blocksize=blocksize, **kwargs)

    def load_dataframe(self, file_resources, blocksize=None):

        """

        Loads the GENCODE annotation file into a pandas or dask DataFrame.

        Args:

            file_resources:

            blocksize:

        Returns:

        """

        gtf_files = select_files_with_ext(file_resources, ".gtf")

        if not gtf_files:

            gtf_files = select_files_with_ext(file_resources, ".gtf.gz")

        ddfs = []

        for filename, filepath in gtf_files.items():

            if blocksize and not isinstance(filepath, str): continue

            ddf = read_gtf(filepath, blocksize=blocksize,

                           compression="gzip" if filename.endswith(".gz") else None)

            ddfs.append(ddf)

        annotation_df = dd.concat(ddfs, interleave_partitions=True,

                                  ignore_unknown_divisions=True) if blocksize else pd.concat(ddfs)

        if self.remove_version_num:

            annotation_df = annotation_df.assign(

                gene_id=annotation_df["gene_id"].str.replace("[.]\d*", "", regex=True),

                transcript_id=annotation_df["transcript_id"].str.replace("[.]\d*", "", regex=True))

        return annotation_df

    def load_sequences(self, fasta_file: str, index=None, keys: pd.Index = None, blocksize=None):

        """

        Args:

            index ():

            keys ():

            fasta_file:

            blocksize:

        """

        if hasattr(self, '_seq_df_dict') and fasta_file in self._seq_df_dict:

            return self._seq_df_dict[fasta_file]

        def get_transcript_id(x):

            key = x.split('|')[0]  # transcript_id

            if self.remove_version_num:

                return re.sub("[.]\d*", "", key)

            else:

                return key

        fa = Fasta(fasta_file, key_function=get_transcript_id, as_raw=True)

        entries = []

        for key, record in tqdm.tqdm(fa.items(), desc=str(fasta_file)):

            if keys is not None and key not in keys: continue

            attrs = record.long_name.split("|")

            record_dict = {

                "transcript_id": attrs[0],

                "gene_id": attrs[1],

                "gene_name": attrs[5],

                "transcript_name": attrs[4],

                "transcript_length": attrs[6],

                "transcript_biotype": intern(attrs[7]),

                SEQUENCE_COL: str(record),

            }

            entries.append(record_dict)

        seq_df = pd.DataFrame(entries)

        if blocksize:

            seq_df = dd.from_pandas(seq_df, chunksize=blocksize)

        if self.remove_version_num:

            seq_df["gene_id"] = seq_df["gene_id"].str.replace("[.]\d*", "", regex=True)

            seq_df["transcript_id"] = seq_df["transcript_id"].str.replace("[.]\d*", "", regex=True)

        # Cache the seq_df

        if not hasattr(self, '_seq_df_dict'):

            self._seq_df_dict = {}

        if keys is not None:

            self._seq_df_dict[fasta_file] = seq_df

        return seq_df

    def get_sequences(self, index: Union[str, Tuple[str]], omic: str, agg: str = 'all', biotypes: List[str] = None):

        """

        Args:

            index (str):

            omic (str):

            agg (str):

            biotypes (List[str]):

        """

        agg_func = self.aggregator_fn(agg)

        # Parse lncRNA & mRNA fasta

        if omic == openomics.MessengerRNA.name():

            fasta_file = self.file_resources["transcripts.fa"]

        elif omic == openomics.LncRNA.name():

            fasta_file = self.file_resources["lncRNA_transcripts.fa"]

        else:

            raise Exception("omic argument must be one of {'MessengerRNA', 'LncRNA'}")

        assert isinstance(fasta_file, str), \

            f"Fasta file provided in `file_resources` must be an uncompressed .fa file. Given {fasta_file}."

        seq_df = self.load_sequences(fasta_file)

        if "gene" in index:

            if biotypes:

                seq_df = seq_df[seq_df["transcript_biotype"].isin(biotypes)]

            else:

                print("INFO: You can pass in a list of transcript biotypes to filter using the argument 'biotypes'.")

            return seq_df.groupby(index)[SEQUENCE_COL].agg(agg_func)

        else:

            return seq_df.groupby(index)[SEQUENCE_COL].first()

    def get_rename_dict(self, from_index="gene_id", to_index="gene_name"):

        """

        Args:

            from_index:

            to_index:

        """

        ensembl_id_to_gene_name = pd.Series(

            self.data[to_index].values, index=self.data[from_index]).to_dict()

        return ensembl_id_to_gene_name

class UniProt(SequenceDatabase):

    COLUMNS_RENAME_DICT = {

        # idmapping_selected.tab

        "UniProtKB-AC": 'protein_id',

        "UniProtKB-ID": 'protein_name',

        "Ensembl": "gene_id",

        "Ensembl_TRS": "transcript_id",

        "Ensembl_PRO": "protein_embl_id",

        "NCBI-taxon": "species_id",

        "GeneID (EntrezGene)": "entrezgene_id",

        "GO": "go_id",

        # FASTA headers

        "OS": 'species', "OX": 'species_id', 'GN': 'gene_name', 'PE': 'ProteinExistence', 'SV': "version",

        # UniProt XML headers

        "accession": "UniProtKB-AC", "name": "protein_name", "gene": "gene_name", "keyword": "keywords",

        "geneLocation": "subcellular_location",

    }

    SPECIES_ID_NAME = {

        '10090': 'MOUSE', '10116': 'RAT', '226900': 'BACCR', '243273': 'MYCGE', '284812': 'SCHPO', '287': 'PSEAI',

        '3702': 'ARATH', '99287': 'SALTY', '44689': 'DICDI', '4577': 'MAIZE', '559292': 'YEAST', '6239': 'CAEEL',

        '7227': 'DROME', '7955': 'DANRE', '83333': 'ECOLI', '9606': 'HUMAN', '9823': 'PIG', }

    SPECIES_ID_TAXONOMIC = {

        'HUMAN': 'human', 'MOUSE': 'rodents', 'RAT': 'rodents', 'BACCR': 'bacteria', 'MYCGE': 'bacteria',

        'SCHPO': 'fungi', 'PSEAI': 'bacteria', 'ARATH': 'plants', 'SALTY': 'bacteria', 'DICDI': 'bacteria',

        'MAIZE': 'plants', 'YEAST': 'fungi', 'CAEEL': 'vertebrates', 'DROME': 'invertebrates', 'DANRE': 'vertebrates',

        'ECOLI': 'bacteria', 'PIG': 'mammals',

    }

    def __init__(self, path="https://ftp.uniprot.org/pub/databases/uniprot/current_release/",

                 file_resources: Dict[str, str] = None,

                 species_id: str = "9606", remove_version_num=True,

                 index_col='UniProtKB-AC', keys=None,

                 col_rename=COLUMNS_RENAME_DICT,

                 **kwargs):

        """

        Loads the UniProt database from https://uniprot.org/ .

        Default path: 'https://ftp.uniprot.org/pub/databases/uniprot/current_release/'

        Default file_resources: {

            file_resources['uniprot_sprot.xml.gz'] = "knowledgebase/complete/uniprot_sprot.xml.gz

            file_resources['uniprot_trembl.xml.gz'] = "knowledgebase/complete/uniprot_trembl.xml.gz

            file_resources["idmapping_selected.tab.gz"] = "knowledgebase/idmapping/idmapping_selected.tab.gz'

            file_resources["proteomes.tsv"] = "https://rest.uniprot.org/proteomes/stream?compressed=true&

                fields=upid%2Corganism%2Corganism_id&format=tsv&query=%28%2A%29%20AND%20%28proteome_type%3A1%29"

            file_resources['speclist.txt'] = 'https://ftp.uniprot.org/pub/databases/uniprot/current_release/

                knowledgebase/complete/docs/speclist'

        }

        Args:

            path:

            file_resources:

            col_rename:

            verbose:

            blocksize:

        """

        self.species_id = species_id

        self.species = UniProt.SPECIES_ID_NAME[species_id] if species_id in UniProt.SPECIES_ID_NAME else None

        self.taxonomic_id = UniProt.SPECIES_ID_TAXONOMIC[

            self.species] if self.species in UniProt.SPECIES_ID_TAXONOMIC else None

        self.remove_version_num = remove_version_num

        if file_resources is None:

            file_resources = {}

            file_resources['uniprot_sprot.xml.gz'] = os.path.join(path, "knowledgebase/complete/uniprot_sprot.xml.gz")

            file_resources['uniprot_trembl.xml.gz'] = os.path.join(path, "knowledgebase/complete/uniprot_trembl.xml.gz")

            file_resources["idmapping_selected.tab.gz"] = os.path.join(

                path, "knowledgebase/idmapping/idmapping_selected.tab.gz")

            if self.species:

                file_resources['uniprot_sprot.xml.gz'] = os.path.join(

                    path, "knowledgebase/taxonomic_divisions/", f'uniprot_sprot_{self.taxonomic_id}.xml.gz')

                file_resources['uniprot_trembl.xml.gz'] = os.path.join(

                    path, "knowledgebase/taxonomic_divisions/", f'uniprot_trembl_{self.taxonomic_id}.xml.gz')

                file_resources["idmapping_selected.tab.gz"] = os.path.join(

                    path, "knowledgebase/idmapping/by_organism/",

                    f'{self.species}_{self.species_id}_idmapping_selected.tab.gz')

            file_resources["proteomes.tsv"] = \

                "https://rest.uniprot.org/proteomes/stream?compressed=true&fields=upid%2Corganism%2Corganism_id&format=tsv&query=%28%2A%29%20AND%20%28proteome_type%3A1%29"

        super().__init__(path=path, file_resources=file_resources, index_col=index_col, keys=keys,

                         col_rename=col_rename,

                         **kwargs)

    def load_dataframe(self, file_resources, blocksize=None):

        """

        Args:

            file_resources:

            blocksize:

        """

        # Load idmapping_selected.tab

        args = dict(

            names=['UniProtKB-AC', 'UniProtKB-ID', 'GeneID (EntrezGene)', 'RefSeq', 'GI', 'PDB', 'GO', 'UniRef100',

                   'UniRef90', 'UniRef50', 'UniParc', 'PIR', 'NCBI-taxon', 'MIM', 'UniGene', 'PubMed', 'EMBL',

                   'EMBL-CDS', 'Ensembl', 'Ensembl_TRS', 'Ensembl_PRO', 'Additional PubMed'],

            usecols=['UniProtKB-AC', 'UniProtKB-ID', 'GeneID (EntrezGene)', 'RefSeq', 'GI', 'PDB', 'GO',

                     'NCBI-taxon', 'Ensembl', 'Ensembl_TRS', 'Ensembl_PRO'],

            dtype='str')

        if blocksize:

            if "idmapping_selected.parquet" in file_resources and \

                isinstance(file_resources["idmapping_selected.parquet"], str):

                idmapping = dd.read_parquet(file_resources["idmapping_selected.parquet"])

            elif "idmapping_selected.tab" in file_resources and \

                isinstance(file_resources["idmapping_selected.tab"], str):

                idmapping = dd.read_table(file_resources["idmapping_selected.tab"], blocksize=blocksize, **args)

            else:

                idmapping = dd.read_table(file_resources["idmapping_selected.tab.gz"], compression="gzip", **args, )

            idmapping: dd.DataFrame

        else:

            if "idmapping_selected.parquet" in file_resources and \

                isinstance(file_resources["idmapping_selected.parquet"], str):

                idmapping = pd.read_parquet(file_resources["idmapping_selected.parquet"])

            else:

                idmapping = pd.read_table(file_resources["idmapping_selected.tab"], index_col=self.index_col, **args)

        # Filter UniProt accession keys

        if self.keys is not None and idmapping.index.name == self.index_col:

            idmapping = idmapping.loc[idmapping.index.isin(self.keys)]

        elif self.keys is not None and idmapping.index.name != self.index_col:

            idmapping = idmapping.loc[idmapping[self.index_col].isin(self.keys)]

        if idmapping.index.name != self.index_col:

            idmapping = idmapping.set_index(self.index_col, sorted=False)

        if not idmapping.known_divisions:

            idmapping.divisions = idmapping.compute_current_divisions()

        # Transform list columns

        if isinstance(idmapping, dd.DataFrame):

            idmapping = idmapping.assign(**self.assign_transforms(idmapping))

        else:

            idmapping = idmapping.assign(**self.assign_transforms(idmapping))

        # Join metadata from uniprot_sprot.parquet

        if any(fn.startswith('uniprot') and fn.endswith('.parquet') for fn in file_resources):

            uniprot_anns = self.load_uniprot_parquet(file_resources, blocksize=blocksize)

            uniprot_anns = uniprot_anns[uniprot_anns.columns.difference(idmapping.columns)]

            uniprot_anns = drop_duplicate_columns(uniprot_anns)

            assert idmapping.index.name == uniprot_anns.index.name, f"{idmapping.index.name} != {uniprot_anns.index.name}"

            idmapping = idmapping.join(uniprot_anns, on=idmapping.index.name, how='left')

        # Load proteome.tsv

        if "proteomes.tsv" in file_resources:

            proteomes = pd.read_table(file_resources["proteomes.tsv"],

                                      usecols=['Organism Id', 'Proteome Id'],

                                      dtype={'Organism Id': 'str', 'Proteome Id': 'str'}) \

                .rename(columns={'Organism Id': 'NCBI-taxon', 'Proteome Id': 'proteome_id'}) \

                .dropna().set_index('NCBI-taxon')

            idmapping = idmapping.join(proteomes, on='NCBI-taxon')

        # Load species info from speclist.txt

        if 'speclist.txt' in file_resources:

            speclist = self.get_species_list(file_resources['speclist.txt'])

            idmapping = idmapping.join(speclist, on='NCBI-taxon')

        return idmapping

    def assign_transforms(self, idmapping: pd.DataFrame) -> Dict[str, Union[dd.Series, pd.Series]]:

        # Convert string of list elements to a np.array

        list2array = lambda x: np.array(x) if isinstance(x, Iterable) else x

        assign_fn = {}

        for col in {'PDB', 'GI', 'GO', 'RefSeq'}.intersection(idmapping.columns):

            try:

                # Split string to list

                assign_fn[col] = idmapping[col].str.split("; ").map(list2array)

            except:

                continue

        for col in {'Ensembl', 'Ensembl_TRS', 'Ensembl_PRO'}.intersection(idmapping.columns):

            # Removing .# ENGS gene version number at the end

            try:

                if self.remove_version_num:

                    series = idmapping[col].str.replace("[.]\d*", "", regex=True)

                else:

                    series = idmapping[col]

                assign_fn[col] = series.str.split("; ").map(list2array)

                if col == 'Ensembl_PRO':

                    # Prepend species_id to ensembl protein ids to match with STRING PPI

                    concat = dd.concat([idmapping["NCBI-taxon"], assign_fn[col]]) \

                        if isinstance(idmapping, dd.DataFrame) else \

                        pd.concat([idmapping["NCBI-taxon"], assign_fn[col]])

                    assign_fn['protein_external_id'] = concat.apply(

                        lambda row: np.char.add(row['NCBI-taxon'] + ".", row['Ensembl_PRO']) \

                            if isinstance(row['Ensembl_PRO'], Iterable) else None,

                        axis=1)

            except:

                continue

        return assign_fn

    def load_uniprot_parquet(self, file_resources: Dict[str, str], blocksize=None) -> Union[dd.DataFrame, pd.DataFrame]:

        dfs = []

        for filename, file_path in file_resources.items():

            if not ('uniprot' in filename and filename.endswith('.parquet')): continue

            if blocksize:

                df: dd.DataFrame = dd.read_parquet(file_path) \

                    .rename(columns=UniProt.COLUMNS_RENAME_DICT)

                if df.index.name in UniProt.COLUMNS_RENAME_DICT:

                    df.index = df.index.rename(UniProt.COLUMNS_RENAME_DICT[df.index.name])

                if self.keys is not None:

                    if self.index_col in df.columns:

                        df = df.loc[df[self.index_col].isin(self.keys)]

                    elif df.index.name == self.index_col:

                        df = df.loc[df.index.isin(self.keys)]

                if df.index.size.compute() == 0: continue

                if df.index.name != self.index_col:

                    try:

                        df = df.set_index(self.index_col, sorted=True)

                    except Exception as e:

                        print(file_path, e)

                        df = df.set_index(self.index_col, sorted=False)

                if not df.known_divisions:

                    df.divisions = df.compute_current_divisions()

            else:

                df = pd.read_parquet(file_path).rename(columns=UniProt.COLUMNS_RENAME_DICT).set_index(self.index_col)

                if self.keys is not None:

                    df_keys = df.index if df.index.name == self.index_col else df[self.index_col]

                    df = df.loc[df_keys.isin(self.keys)]

                if df.index.size == 0: continue

            dfs.append(df)

        if dfs:

            dfs = dd.concat(dfs, interleave_partitions=True) if blocksize else pd.concat(dfs)

            return dfs

        else:

            return None

    def load_uniprot_xml(self, file_path: str, keys=None, blocksize=None) -> pd.DataFrame:

        records = []

        seqfeats = []

        if isinstance(keys, str):

            index = keys

            keys_set = self.data.index if keys == self.data.index.name else self.data[keys]

        elif isinstance(keys, (dd.Index, dd.Series)):

            index = keys.name

            keys_set = keys.compute()

        else:

            index = keys_set = None

        for record in tqdm.tqdm(SeqIO.parse(file_path, "uniprot-xml"), desc=str(file_path)):

            # Sequence features

            annotations = defaultdict(None, record.annotations)

            record_dict = {

                'protein_id': record.id,

                "protein_name": record.name,

                'gene_name': annotations['gene_name_primary'],

                'description': record.description,

                'molecule_type': annotations['molecule_type'],

                'created': annotations['created'],

                'ec_id': annotations['type'],

                'subcellular_location': annotations['comment_subcellularlocation_location'],

                'taxonomy': annotations['taxonomy'],

                'keywords': annotations['keywords'],

                'sequence_mass': annotations['sequence_mass'],

                SEQUENCE_COL: str(record.seq),

            }

            if index is not None:

                if record_dict[keys] not in keys_set: continue

            records.append(record_dict)

            # Sequence interval features

            _parse_interval = lambda sf: pd.Interval(left=sf.location.start, right=sf.location.end, )

            feature_type_intervals = defaultdict(lambda: [])

            for sf in record.features:

                if isinstance(sf.location.start, ExactPosition) and isinstance(sf.location.end, ExactPosition):

                    feature_type_intervals[sf.type].append(_parse_interval(sf))

            features_dict = {type: pd.IntervalIndex(intervals, name=type) \

                             for type, intervals in feature_type_intervals.items()}

            seqfeats.append({"protein_id": record.id, **features_dict})

        records_df = pd.DataFrame(records) if not blocksize else dd.from_pandas(records, chunksize=blocksize)

        records_df = records_df.set_index(['protein_id'])

        seqfeats_df = pd.DataFrame(seqfeats) if not blocksize else dd.from_pandas(seqfeats, chunksize=blocksize)

        seqfeats_df = seqfeats_df.set_index(['protein_id'])

        seqfeats_df.columns = [f"seq/{col}" for col in seqfeats_df.columns]

        # Join new metadata to self.data

        if SEQUENCE_COL not in self.data.columns:

            exclude_cols = records_df.columns.intersection(self.data.columns)

            self.data = self.data.join(records_df.drop(columns=exclude_cols, errors="ignore"),

                                       on='protein_id', how="left")

        else:

            self.data.update(records_df, overwrite=False)

        # Add new seq features

        if len(seqfeats_df.columns.difference(self.data.columns)):

            self.data = self.data.join(seqfeats_df.drop(columns=seqfeats_df.columns.intersection(self.data.columns)),

                                       on='protein_id', how="left")

        # Fillna seq features

        if len(seqfeats_df.columns.intersection(self.data.columns)):

            self.data.update(seqfeats_df.filter(seqfeats_df.columns.intersection(self.data.columns), axis='columns'),

                             overwrite=False)

        return records_df

    @classmethod

    def get_species_list(cls, file_path):

        speclist = pd.read_fwf(file_path,

                               names=['species_code', 'Taxon', 'species_id', 'attr'],

                               comment="==", skipinitialspace=True, skiprows=59, skipfooter=4)

        speclist = speclist.drop(index=speclist.index[~speclist['attr'].str.contains("=")])

        speclist['species_id'] = speclist['species_id'].str.rstrip(":")

        speclist = speclist.fillna(method='ffill')

        speclist = speclist.groupby(speclist.columns[:3].tolist())['attr'] \

            .apply('|'.join) \

            .apply(lambda s: dict(map(str.strip, sub.split('=', 1)) for sub in s.split("|") if '=' in sub)) \

            .apply(pd.Series)

        speclist = speclist.rename(columns={'N': 'Official (scientific) name', 'C': 'Common name', 'S': 'Synonym'}) \

            .reset_index() \

            .set_index('species_id')

        speclist['Taxon'] = speclist['Taxon'].replace(

            {'A': 'archaea', 'B': 'bacteria', 'E': 'eukaryota', 'V': 'viruses', 'O': 'others'})

        speclist.index.name = 'NCBI-taxon'

        return speclist

    def load_sequences(self, fasta_file: str, index=None, keys: Union[pd.Index, List[str]] = None, blocksize=None) \

        -> OrderedDict:

        def get_id(s: str):

            if index == 'protein_id':

                return s.split('|')[1]

            elif index == 'protein_name':

                return s.split('|')[2]

            else:

                return s.split('|')[1]

        fa = Fasta(fasta_file, key_function=get_id, as_raw=True, )

        return fa.records

    def get_sequences(self, index: str, omic: str = None, agg: str = "first", **kwargs):

        assert index, '`index` must be either "protein_id" or "protein_name"'

        # Parse lncRNA & mRNA fasta

        seq_df = self.load_sequences(self.file_resources["uniprot_sprot.fasta"], index=index, blocksize=self.blocksize)

        if "uniprot_trembl.fasta" in self.file_resources:

            trembl_seq_df = self.load_sequences(self.file_resources["uniprot_trembl.fasta"], index=index,

                                                blocksize=self.blocksize)

            seq_df.update(trembl_seq_df)

        return seq_df

class MirBase(SequenceDatabase):

    """Loads the MirBase database from https://mirbase.org .

    Default path: "ftp://mirbase.org/pub/mirbase/CURRENT/" .

    Default file_resources: {

        "aliases.txt": "aliases.txt.gz",

        "mature.fa": "mature.fa.gz",

        "hairpin.fa": "hairpin.fa.gz",

        "rnacentral.mirbase.tsv": "ftp://ftp.ebi.ac.uk/pub/databases/RNAcentral/current_release/id_mapping/database_mappings/mirbase.tsv",

    }

    """

    def __init__(

        self,

        path="http://mirbase.org/ftp/CURRENT/",

        file_resources=None,

        species_id: Optional[str] = '9606',

        index_col: str = "mirbase_id",

        col_rename=None,

        **kwargs,

    ):

        """

        Args:

            path:

            file_resources:

            sequence (str):

            species_id (str): Species code, e.g., 9606 for human

            col_rename:

            blocksize:

        """

        if file_resources is None:

            file_resources = {}

            file_resources["aliases.txt.gz"] = "aliases.txt.gz"

            file_resources["mature.fa.gz"] = "mature.fa.gz"

            file_resources["hairpin.fa.gz"] = "hairpin.fa.gz"

        if 'rnacentral.mirbase.tsv' not in file_resources:

            file_resources["rnacentral.mirbase.tsv"] = "ftp://ftp.ebi.ac.uk/pub/databases/RNAcentral/current_release/" \

                                                       "id_mapping/database_mappings/mirbase.tsv"

        self.species_id = species_id

        super().__init__(path=path, file_resources=file_resources, index_col=index_col, col_rename=col_rename, **kwargs)

    def load_dataframe(self, file_resources, blocksize=None):

        """

        Args:

            file_resources: dict of file name and path

            blocksize:

        """

        rnacentral_mirbase = pd.read_table(

            file_resources["rnacentral.mirbase.tsv"], low_memory=True, header=None,

            names=["RNAcentral id", "database", "mirbase_id", "species_id", "RNA type", "NA"],

            usecols=["RNAcentral id", "database", "mirbase_id", "species_id", "RNA type"],

            index_col="RNAcentral id",

            dtype={'mirbase_id': 'str', "species_id": "category", 'database': 'category', 'RNA type': 'category'})

        if isinstance(self.species_id, str):

            rnacentral_mirbase = rnacentral_mirbase[rnacentral_mirbase["species_id"] == self.species_id]

        elif isinstance(self.species_id, Iterable):

            rnacentral_mirbase = rnacentral_mirbase[rnacentral_mirbase["species_id"].isin(set(self.species_id))]

        mirbase_df = pd.read_table(file_resources["aliases.txt"], low_memory=True, header=None,

                                   names=["mirbase_id", "mirbase_name"], index_col=self.index_col,

                                   dtype='str', )

        if mirbase_df.index.name == 'mirbase id':

            mirbase_df = mirbase_df.join(rnacentral_mirbase, on=self.index_col, how="left", rsuffix='_rnacentral')

        else:

            mirbase_df = mirbase_df.merge(rnacentral_mirbase, on=self.index_col, how="left")

        # Expanding miRNA names in each MirBase Ascension ID

        mirbase_df['mirbase_name'] = mirbase_df['mirbase_name'].str.rstrip(";").str.split(";")

        seq_dfs = []

        for filename in file_resources:

            if filename.endswith('.fa') or filename.endswith('.fasta'):

                assert isinstance(file_resources[filename], str), f"Must provide a path to an uncompressed .fa file. " \

                                                                  f"Given {file_resources[filename]}"

                df = self.load_sequences(file_resources[filename], index=self.index_col, keys=self.keys)

                seq_dfs.append(df)

        if len(seq_dfs):

            seq_dfs = pd.concat(seq_dfs, axis=0)

            mirbase_df = mirbase_df.join(seq_dfs, how='left', on=self.index_col)

        else:

            logger.info('Missing sequence data because no "hairpin.fa" or "mature.fa" file were given.')

        # mirbase_df = mirbase_df.explode(column='gene_name')

        # mirbase_name["miRNA name"] = mirbase_name["miRNA name"].str.lower()

        # mirbase_name["miRNA name"] = mirbase_name["miRNA name"].str.replace("-3p.*|-5p.*", "")

        return mirbase_df

    def load_sequences(self, fasta_file, index=None, keys=None, blocksize=None):

        """

        Args:

            fasta_file:

            index ():

            keys ():

            blocksize:

        """

        if hasattr(self, '_seq_df_dict') and fasta_file in self._seq_df_dict:

            return self._seq_df_dict[fasta_file]

        fa = Fasta(fasta_file, read_long_names=True, as_raw=True)

        mirna_types = {'stem-loop', 'stem', 'type', 'loop'}

        entries = []

        for key, record in tqdm.tqdm(fa.items(), desc=str(fasta_file)):

            attrs: List[str] = record.long_name.split(" ")

            if attrs[-1] in mirna_types:

                if attrs[-2] in mirna_types:

                    mirna_type = intern(' '.join(attrs[-2:]))

                    gene_name_idx = -3

                else:

                    mirna_type = intern(attrs[-1])

                    gene_name_idx = -2

            else:

                mirna_type = None

                gene_name_idx = -1

            record_dict = {

                "gene_id": attrs[0],

                "mirbase_id": attrs[1],

                "species": intern(" ".join(attrs[2:gene_name_idx])),

                "gene_name": attrs[gene_name_idx],

                "mirna_type": mirna_type,

                SEQUENCE_COL: str(record),

            }

            if keys is not None and index:

                if record_dict[index] not in keys:

                    del record_dict

                    continue

            entries.append(record_dict)

        df = pd.DataFrame(entries)

        if index:

            df = df.set_index(index)

        # if blocksize:

        #     df = dd.from_pandas(df, chunksize=blocksize)

        if not hasattr(self, '_seq_df_dict'):

            self._seq_df_dict = {}

        self._seq_df_dict[fasta_file] = df

        return df

    def get_sequences(self,

                      index="gene_name",

                      omic=None,

                      agg="all",

                      **kwargs):

        """

        Args:

            index:

            omic:

            agg:

            **kwargs:

        """

        dfs = []

        for filename in self.file_resources:

            if filename.endswith('.fa'):

                seq_df = self.load_sequences(self.file_resources[filename])

                dfs.append(seq_df)

        seq_df = pd.concat(dfs, axis=0)

        seq_df = seq_df.groupby(index)[SEQUENCE_COL].agg(self.aggregator_fn(agg))

        return seq_df

class RNAcentral(SequenceDatabase):

    """

    Loads the RNAcentral database from https://rnacentral.org/ and provides a series of methods to extract sequence data from it.

    Default path: https://ftp.ebi.ac.uk/pub/databases/RNAcentral/current_release/ .

    Default file_resources: {

        "rnacentral_rfam_annotations.tsv": "go_annotations/rnacentral_rfam_annotations.tsv.gz",

        "database_mappings/gencode.tsv": "id_mapping/database_mappings/gencode.tsv",

        "gencode.fasta": "sequences/by-database/gencode.fasta",

        ...

    }

    """

    COLUMNS_RENAME_DICT = {

        'ensembl_gene_id': 'gene_id',

        'external id': 'transcript_id',

        'GO terms': 'go_id',

        'gene symbol': 'gene_id',

    }

    def __init__(self, path="https://ftp.ebi.ac.uk/pub/databases/RNAcentral/current_release/", file_resources=None,

                 col_rename=COLUMNS_RENAME_DICT, species_id: Union[List[str], str, None] = None,

                 index_col="RNAcentral id", keys=None,

                 remove_version_num=True, remove_species_suffix=True, **kwargs):

        """

        Provide

        Args:

            path ():

            file_resources ():

            col_rename ():

            species_id ():

            index_col ():

            keys ():

            remove_version_num ():

            remove_species_suffix ():

            **kwargs ():

        """

        self.species_id = species_id

        self.remove_version_num = remove_version_num

        self.remove_species_suffix = remove_species_suffix

        if file_resources is None:

            file_resources = {}

            file_resources["rnacentral_rfam_annotations.tsv.gz"] = "go_annotations/rnacentral_rfam_annotations.tsv.gz"

            file_resources["database_mappings/ensembl_gencode.tsv"] = "id_mapping/database_mappings/ensembl_gencode.tsv"

            file_resources["database_mappings/mirbase.tsv"] = "id_mapping/database_mappings/mirbase.tsv"

        super().__init__(path=path, file_resources=file_resources, col_rename=col_rename, index_col=index_col,

                         keys=keys, **kwargs)

    def load_dataframe(self, file_resources, blocksize=None):

        """

        Args:

            file_resources:

            blocksize:

        """

        transcripts_df = []

        # Concatenate transcripts ids by combining `database_mappings/` files from multiple RNAcentral databases

        for filename in (fname for fname in file_resources if "database_mappings" in fname):

            args = dict(low_memory=True, header=None,

                        names=["RNAcentral id", "database", "external id", "species_id", "RNA type", "gene symbol"],

                        dtype={'gene symbol': 'str',

                               'database': 'category', 'species_id': 'category', 'RNA type': 'category', })

            if blocksize:

                if filename.endswith('.tsv'):

                    id_mapping: dd.DataFrame = dd.read_table(

                        file_resources[filename], blocksize=None if isinstance(blocksize, bool) else blocksize, **args)

                elif filename.endswith('.parquet'):

                    id_mapping: dd.DataFrame = dd.read_parquet(

                        file_resources[filename], blocksize=None if isinstance(blocksize, bool) else blocksize, )

                else:

                    id_mapping = None

            else:

                if filename.endswith('.tsv'):

                    id_mapping = pd.read_table(file_resources[filename], **args)

                elif filename.endswith('.parquet'):

                    id_mapping = pd.read_parquet(file_resources[filename])

                else:

                    id_mapping = None

            if id_mapping is None:

                raise Exception("Must provide a file with 'database_mappings/(*).tsv' in file_resources")

            # Filter by species

            if isinstance(self.species_id, str):

                id_mapping = id_mapping.where(id_mapping["species_id"] == self.species_id)

            elif isinstance(self.species_id, Iterable):

                id_mapping = id_mapping.where(id_mapping["species_id"].isin(self.species_id))

            # Filter by index

            if self.keys and id_mapping.index.name == self.index_col:

                id_mapping = id_mapping.loc[id_mapping.index.isin(self.keys)]

            elif self.keys and id_mapping.index.name != self.index_col:

                id_mapping = id_mapping.loc[id_mapping[self.index_col].isin(self.keys)]

            # Add species_id prefix to index values to match the sequence ids

            if "RNAcentral id" in id_mapping.columns:

                id_mapping["RNAcentral id"] = id_mapping["RNAcentral id"] + "_" + id_mapping["species_id"].astype(str)

            elif "RNAcentral id" == id_mapping.index.name:

                id_mapping.index = id_mapping.index + "_" + id_mapping["species_id"].astype(str)

            # Add sequence column if a FASTA file provided for the database

            fasta_filename = f"{filename.split('/')[-1].split('.')[0]}.fasta"

            if fasta_filename in file_resources:

                seq_df = self.load_sequences(file_resources[fasta_filename])

                id_mapping = id_mapping.merge(seq_df, how='left',

                                              left_on="RNAcentral id",

                                              left_index=True if id_mapping.index.name == "RNAcentral id" else False,

                                              right_index=True)

            else:

                logger.info(f"{fasta_filename} not provided for `{filename}` so missing sequencing data")

            if self.remove_version_num and 'gene symbol' in id_mapping.columns:

                id_mapping["gene symbol"] = id_mapping["gene symbol"].str.replace("[.].\d*", "", regex=True)

            if self.remove_species_suffix:

                id_mapping["RNAcentral id"] = id_mapping["RNAcentral id"].str.replace("_(\d*)", '', regex=True)

            # Set index

            args = dict(sorted=True) if blocksize else {}

            id_mapping = id_mapping.set_index(self.index_col, **args)

            if isinstance(id_mapping, dd.DataFrame) and not id_mapping.known_divisions:

                id_mapping.divisions = id_mapping.compute_current_divisions()

            transcripts_df.append(id_mapping)

        # Concatenate multiple `database_mappings` files from different databases

        if blocksize:

            transcripts_df = dd.concat(transcripts_df, axis=0, interleave_partitions=True, join='outer')

        else:

            transcripts_df = pd.concat(transcripts_df, axis=0, join='outer')

        # Join go_id and Rfams annotations to each "RNAcentral id" from 'rnacentral_rfam_annotations.tsv'

        try:

            transcripts_df = self.add_rfam_annotation(transcripts_df, file_resources, blocksize)

        except Exception as e:

            logger.warning(f"Failed to add Rfam annotations to transcripts_df:")

            traceback.print_exc()

        return transcripts_df

    def add_rfam_annotation(self, transcripts_df: Union[pd.DataFrame, dd.DataFrame],

                            file_resources, blocksize=None) -> Union[pd.DataFrame, dd.DataFrame]:

        args = dict(low_memory=True, names=["RNAcentral id", "GO terms", "Rfams"])

        if blocksize:

            if 'rnacentral_rfam_annotations.tsv' in file_resources and isinstance(

                file_resources['rnacentral_rfam_annotations.tsv'], str):

                anns = dd.read_table(file_resources["rnacentral_rfam_annotations.tsv"], **args)

            else:

                anns = dd.read_table(file_resources["rnacentral_rfam_annotations.tsv.gz"], compression="gzip", **args)

            anns = anns.set_index("RNAcentral id", sorted=True)

            # Filter annotations by "RNAcentral id" in `transcripts_df`

            anns = anns.loc[anns.index.isin(transcripts_df.index.compute())]

            if not anns.known_divisions:

                anns.divisions = anns.compute_current_divisions()

            # Groupby on index

            anns_groupby: dd.DataFrame = anns \

                .groupby(by=lambda idx: idx) \

                .agg({col: get_agg_func('unique', use_dask=True) for col in ["GO terms", 'Rfams']})

        else:

            anns = pd.read_table(file_resources["rnacentral_rfam_annotations.tsv"], index_col='RNAcentral id', **args)

            idx = transcripts_df.index.compute() if isinstance(transcripts_df, dd.DataFrame) else transcripts_df.index

            anns = anns.loc[anns.index.isin(set(idx))]

            anns_groupby = anns.groupby("RNAcentral id").agg({col: 'unique' for col in ["GO terms", 'Rfams']})

        transcripts_df = transcripts_df.merge(anns_groupby, how='left', left_index=True, right_index=True)

        return transcripts_df

    def load_sequences(self, fasta_file: str, index=None, keys=None, blocksize=None):

        """

        Args:

            index ():

            fasta_file:

            keys ():

            blocksize:

        """

        fa = Fasta(fasta_file, as_raw=True)

        entries = []

        for key, record in tqdm.tqdm(fa.items(), desc=str(fasta_file)):

            id = re.sub("_(\d*)", '', key) if self.remove_species_suffix else key

            if keys is not None and self.index_col == 'RNAcentral id' and id not in keys:

                continue

            desc = record.long_name.split(" ", maxsplit=1)[-1]

            record_dict = {

                'RNAcentral id': key,

                'description': desc,

                SEQUENCE_COL: str(record),

            }

            entries.append(record_dict)

        df = pd.DataFrame(entries).set_index("RNAcentral id")

        return df

    def get_sequences(self,

                      index="RNAcentral id",

                      omic=None,

                      agg="all",

                      **kwargs):

        """

        Args:

            index:

            omic:

            agg:

            **kwargs:

        """

        dfs = []

        for filename in self.file_resources:

            if filename.endswith('.fa') or filename.endswith('.fasta'):

                seq_df = self.load_sequences(self.file_resources[filename])

                dfs.append(seq_df)

        seq_df = pd.concat(dfs, axis=0)

        seq_df = seq_df.groupby(index)[SEQUENCE_COL].agg(self.aggregator_fn(agg))

        return seq_df