"""Perform queries against the lookup endpoints
"""
from ensembl_rest_client.base import base_rest_client as brc

try:
    # v2
    from pandas import json_normalize
except ImportError:
    # v 1.X
    from pandas.io.json import json_normalize

import pandas as pd
import itertools


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def NORM_POST_LOOKUP_ID_GENE(data, **kwargs):
    """
    When issuing a POST lookup ID query for a gene IDs, format returned
    data into a reasonable `pandas.DataFrame`. Not that missing IDs are not
    returned. Also, at present this will only work for 'gene' features

    Parameters
    ----------
    data : dict of dict
        The keys of the dictionary are GENE IDs and the values are dicts
        with nested data (if expand is True)
    **kwargs
        Arguments to the POST query that are required to process the
        different formats
    """
    try:
        # Determine the format
        data_format = kwargs['format']
    except KeyError:
        data_format = 'condensed'

    # Generate a list of dict and remove empty entries
    defined_data = [v for v in data.values() if v is not None]

    # These *_cols lists define the columns (as named after de-normalising)
    # and the column order we want in the final output. Not that as defined
    # here they are for the 'condensed' data_format
    exon_cols = ['trans_id', 'exon_id']
    gene_cols = ['gene_id', 'gene_version']
    trans_cols = ['trans_id', 'trans_version']
    prot_cols = ['trans_Translation_id']

    # The *_meta columns are the columns we want from the un-nested level
    # of the decoded JSON that we want to keep after de-normalising
    exon_meta = ['id']
    trans_meta = ['id', 'version']

    # The full data format has more info such as co-ordinates, as such we
    # adjust the columns we want to keep
    if data_format == 'full':
        gene_cols.extend(
            ['gene_display_name', 'gene_seq_region_name', 'gene_start',
             'gene_end', 'gene_strand', 'gene_biotype',
             'gene_description', 'gene_source']
        )
        trans_cols.extend(
            ['trans_display_name', 'trans_is_canonical',
             'trans_seq_region_name', 'trans_start', 'trans_end',
             'trans_strand', 'trans_biotype', 'trans_source']
        )
        prot_cols.extend(
            ['trans_Translation_start', 'trans_Translation_end',
             'trans_Translation_length']
        )
        exon_cols.extend(
            ['exon_version', 'exon_seq_region_name', 'exon_start',
             'exon_end', 'exon_strand']
        )
        trans_meta.extend(
            ['display_name', 'seq_region_name', 'start', 'end', 'strand',
             'biotype', 'description', 'source']
        )
    try:
        # If we have expanded results then we will have a lot lore nested
        # data on transcripts and exons and translations, this is expanded
        # using JSON normalise
        if kwargs['expand'] is True:
            # Extract a list of all the transcipts, this is done to get the
            # transcripts at the top level so I can do a separate
            # extraction of gene->transcript and transcript->exon and then
            # merge. I can't seem to noralize all in one go
            trans = itertools.chain.from_iterable(
                v['Transcript'] for v in defined_data
            )
            # Normalise exons/transcript and extract columns
            exons = json_normalize(
                trans,
                meta_prefix='trans_',
                record_prefix='exon_',
                meta=exon_meta,
                record_path=['Exon']
            )[exon_cols]
            trans = json_normalize(
                defined_data,
                meta_prefix='gene_',
                record_prefix='trans_',
                meta=trans_meta,
                record_path="Transcript",
                max_level=1,
                errors='ignore',
                sep='_'
            )[gene_cols + trans_cols + prot_cols]
            trans.columns = map(str.lower, trans.columns)
            return trans.merge(
                exons, left_on='trans_id', right_on='trans_id'
            ).set_index(['gene_id', 'trans_id', 'exon_id'])
    except KeyError:
        pass

    # Here we are processing unexpanded data but we still want to have some
    # sensible column naming and ordering
    df = pd.DataFrame(defined_data)
    df.columns = pd.Index(['gene_{0}'.format(i) for i in df.columns])
    return df[gene_cols].set_index('gene_id')


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class _Lookup(object):
    """Perform Lookup queries
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_lookup_id(self, query_id, db_type=None, expand=None,
                      data_format=None, phenotypes=None, species=None,
                      utr=None, norm_func=brc.default_pandas):
        """
        Find the species and database for a single identifier e.g. gene,
        transcript, protein. GET method. Queries:
        `lookup/id/:id <https://rest.ensembl.org/documentation/info/lookup>`_.

        Parameters
        ----------
        query_id : str
            An Ensembl stable ID. Example values: 'ENSG00000157764'.
        db_type : str or NoneType, optional, default: NoneType
            Restrict the search to a database other than the default. Useful if
            you need to use a DB other than core. Example values: 'core
            otherfeatures'.
        expand : bool or NoneType, optional, default: NoneType
            Expands the search to include any connected features. e.g. If the
            object is a gene, its transcripts, translations and exons will be
            returned as well.
        data_format : str or NoneType, optional, default: NoneType
            Specify the formats to emit from this endpoint. Allowed values:
            full,condensed
        phenotypes : bool or NoneType, optional, default: NoneType
            Include phenotypes. Only available for gene objects.
        species : str or NoneType, optional, default: NoneType
            Species name/alias. Example values: 'homo_sapiens human'.
        utr : bool or NoneType, optional, default: NoneType
            Include 5' and 3' UTR features. Only available if the expand option
            is used.
        norm_func : obj:`func`, optional, default: default_pandas
            The normalisation function to use when obj.pandas is True

        Returns
        -------
        data : list or dict or :obj:`pandas.DataFrame`
            The rest query data
        """
        method_type = 'GET'
        result_format = 'dict'
        endpoint = 'lookup/id/{id}'.format(
            id=query_id
        )

        return self.rest_query(
            method_type,
            endpoint,
            db_type=db_type,
            expand=expand,
            format=data_format,
            phenotypes=phenotypes,
            species=species,
            utr=utr,
            norm_func=norm_func
        )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def post_lookup_id(self, ids, db_type=None, expand=None, data_format=None,
                       object_type=None, species=None,
                       norm_func=NORM_POST_LOOKUP_ID_GENE):
        """
        Find the species and database for several identifiers. IDs that are not
        found are returned with no data. POST method. Queries:
        `lookup/id <https://rest.ensembl.org/documentation/info/lookup_post>`_.
        Max post size: 1000.

        Parameters
        ----------
        ids : list
            Post option not defined in required/optional args.
        db_type : str or NoneType, optional, default: NoneType
            Restrict the search to a database other than the default. Useful if
            you need to use a DB other than core. Example values: 'core
            otherfeatures'.
        expand : bool or NoneType, optional, default: NoneType
            Expands the search to include any connected features. e.g. If the
            object is a gene, its transcripts, translations and exons will be
            returned as well.
        data_format : str or NoneType, optional, default: NoneType
            Specify the formats to emit from this endpoint. Allowed values:
            full,condensed
        object_type : str or NoneType, optional, default: NoneType
            Filter by feature type. Example values: 'gene','transcript'.
        species : str or NoneType, optional, default: NoneType
            Species name/alias. Causes problems if the species doesn't match
            the identifiers in the POST body. Example values: 'homo_sapiens
            human'.
        norm_func : obj:`func`, optional, default: NORM_POST_LOOKUP_ID_GENE
            The normalisation function to use when obj.pandas is True

        Returns
        -------
        data : list or dict or :obj:`pandas.DataFrame`
            The rest query data
        """
        method_type = 'POST'
        result_format = 'dict'
        maximum_post_size = 1000
        endpoint = 'lookup/id'

        return self.rest_query(
            method_type,
            endpoint,
            maximum_post_size=maximum_post_size,
            ids=ids,
            db_type=db_type,
            expand=expand,
            format=data_format,
            object_type=object_type,
            species=species,
            norm_func=norm_func
        )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_lookup_symbol(self, species, symbol, expand=None,
                          data_format=None, norm_func=brc.default_pandas):
        """
        Find the species and database for a symbol in a linked external
        database. GET method. Queries:
        `lookup/symbol/:species/:symbol <https://rest.ensembl.org/documentation/info/symbol_lookup>`_.

        Parameters
        ----------
        species : str
            Species name/alias. Example values: 'homo_sapiens human'.
        symbol : str
            A name or symbol from an annotation source has been linked to a
            genetic feature. Example values: 'BRCA2'.
        expand : bool or NoneType, optional, default: NoneType
            Expands the search to include any connected features. e.g. If the
            object is a gene, its transcripts, translations and exons will be
            returned as well.
        norm_func : obj:`func`, optional, default: default_pandas
            The normalisation function to use when obj.pandas is True

        Returns
        -------
        data : list or dict or :obj:`pandas.DataFrame`
            The rest query data
        """
        method_type = 'GET'
        endpoint = 'lookup/symbol/{species}/{symbol}'.format(
            species=species,
            symbol=symbol
        )

        return self.rest_query(
            method_type,
            endpoint,
            expand=expand,
            format=data_format,
            norm_func=norm_func
        )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def post_lookup_symbol(self, species, symbols, expand=None,
                           data_format=None, norm_func=brc.default_pandas):
        """
        Find the species and database for a set of symbols in a linked external
        database. Unknown symbols are omitted from the response. POST method.
        Queries:
        `lookup/symbol/:species/:symbol <https://rest.ensembl.org/documentation/info/symbol_post>`_.
        Max post size: 1000.

        Parameters
        ----------
        species : str
            Species name/alias for the whole batch of symbols. Example values:
            'homo_sapiens human'.
        symbols : list
            Post option not defined in required/optional args.
        expand : bool or NoneType, optional, default: NoneType
            Expands the search to include any connected features. e.g. If the
            object is a gene, its transcripts, translations and exons will be
            returned as well.
        data_format : str or NoneType, optional, default: NoneType
            Specify the layout of the response. Allowed values: full,condensed
        norm_func : obj:`func`, optional, default: default_pandas
            The normalisation function to use when obj.pandas is True

        Returns
        -------
        data : list or dict or :obj:`pandas.DataFrame`
            The rest query data
        """
        raise NotImplementedError(
            'I think this method call is broken actually in the rest API'
        )
        method_type = 'POST'
        maximum_post_size = 1000
        endpoint = 'lookup/symbol/{species}/{symbol}'.format(
            species=species,
            symbol=symbols
        )

        return self.rest_query(
            method_type,
            endpoint,
            maximum_post_size=maximum_post_size,
            expand=expand,
            format=data_format,
            norm_func=norm_func
        )


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class Lookup(brc._BaseRestClient, _Lookup):
    """Perform Lookup queries
    """
    pass
