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


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class _LinkageDisequilibrium(object):
    """Perform LinkageDisequilibrium queries
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_ld(self, query_id, population_name, species, attribs=None,
               d_prime=None, r2=None, window_size=None,
               norm_func=brc.default_pandas):
        """
        Computes and returns LD values between the given variant and all other
        variants in a window centered around the given variant. The window size
        is set to 500 kb. GET method. Queries:
        `ld/:species/:id/:population_name <https://rest.ensembl.org/documentation/info/ld_id_get>`_.

        Parameters
        ----------
        query_id : str
            Variant id. Example values: 'rs56116432'.
        population_name : str
            Population for which to compute LD. Use GET
            /info/variation/populations/:species?filter=LD to retrieve a list
            of all populations with LD data. Example values:
            '1000GENOMES:phase_3:KHV'.
        species : str
            Species name/alias. Example values: 'homo_sapiens human'.
        attribs : bool or NoneType, optional, default: NoneType
            Add variation attributes for the variation which is used to compute
            LD data with the input variation: chr, start, end, strand,
            consequence_type, clinical_significance.
        d_prime : float or NoneType, optional, default: NoneType
            Measure of LD. If D' is provided only return pairs of variants
            whose D' value is equal to or greater than the value provided.
            Example values: '1.0'.
        r2 : float or NoneType, optional, default: NoneType
            Measure of LD. If r-squared is provided only return pairs of
            variants whose r-squared value is equal to or greater than the
            value provided. Example values: '0.85'.
        window_size : int or NoneType, optional, default: NoneType
            Window size in kb. The maximum allowed value for the window size is
            500 kb. LD is computed for the given variant and all variants that
            are located within the specified window. Example values: '500'.
        norm_func : obj:`func`, optional, default: default_pandas
            The normalisation function to use

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

        return self.rest_query(
            method_type,
            endpoint,
            attribs=attribs,
            d_prime=d_prime,
            r2=r2,
            window_size=window_size
        )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_ld_pairwise(self, id1, id2, species, d_prime=None,
                        population_name=None, r2=None,
                        norm_func=brc.default_pandas):
        """
        Computes and returns LD values between the given variants. GET method.
        Queries:
        `ld/:species/pairwise/:id1/:id2 <https://rest.ensembl.org/documentation/info/ld_pairwise_get>`_.

        Parameters
        ----------
        id1 : str
            Variant id1. Example values: 'rs6792369'.
        id2 : str
            Variant id2. Example values: 'rs1042779'.
        species : str
            Species name/alias. Example values: 'homo_sapiens human'.
        d_prime : float or NoneType, optional, default: NoneType
            Measure of LD. If D' is provided only return pairs of variants
            whose D' value is equal to or greater than the value provided.
            Example values: '1.0'.
        population_name : str or NoneType, optional, default: NoneType
            Only compute LD for this population. Use GET
            /info/variation/populations/:species?filter=LD to retrieve a list
            of all populations with LD data. Example values:
            '1000GENOMES:phase_3:KHV'.
        r2 : float or NoneType, optional, default: NoneType
            Measure of LD. If r-squared is provided only return pairs of
            variants whose r-squared value is equal to or greater than the
            value provided. Example values: '0.85'.
        norm_func : obj:`func`, optional, default: default_pandas
            The normalisation function to use

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

        return self.rest_query(
            method_type,
            endpoint,
            d_prime=d_prime,
            population_name=population_name,
            r2=r2
        )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_ld_region(self, population_name, region, species, d_prime=None,
                      r2=None, norm_func=brc.default_pandas):
        """
        Computes and returns LD values between all pairs of variants in the
        defined region. GET method. Queries:
        `ld/:species/region/:region/:population_name <https://rest.ensembl.org/documentation/info/ld_region_get>`_.

        Parameters
        ----------
        population_name : str
            Population for which to compute LD. Use GET
            /info/variation/populations/:species?filter=LD to retrieve a list
            of all populations with LD data. Example values:
            '1000GENOMES:phase_3:KHV'.
        region : str
            Query region. A maximum region size of 500 kb is allowed. If the
            query region overlaps the MHC region only a maximum region size of
            10 kb is allowed. Example values: '6:25837556..25843455'.
        species : str
            Species name/alias. Example values: 'homo_sapiens human'.
        d_prime : float or NoneType, optional, default: NoneType
            Measure of LD. If D' is provided only return pairs of variants
            whose D' value is equal to or greater than the value provided.
            Example values: '1.0'.
        r2 : float or NoneType, optional, default: NoneType
            Measure of LD. If r-squared is provided only return pairs of
            variants whose r-squared value is equal to or greater than the
            value provided. Example values: '0.85'.
        norm_func : obj:`func`, optional, default: default_pandas
            The normalisation function to use

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

        return self.rest_query(
            method_type,
            endpoint,
            d_prime=d_prime,
            r2=r2
        )


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class LinkageDisequilibrium(brc._BaseRestClient, _LinkageDisequilibrium):
    """Perform LinkageDisequilibrium queries
    """
    pass
