__license__ = """
Copyright 2020 Potsdam-Institut für Klimafolgenforschung e.V.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import itertools
import logging
import typing

import pandas as pd
import requests
import treelib


class UNFCCCApiReader:
    """Provides simplified unified access to the Flexible Query API of the UNFCCC data
    access for all parties.

    Essentially encapsulates https://di.unfccc.int/flex_non_annex1 and
    https://di.unfccc.int/flex_annex1 .

    Attributes
    ----------
    parties : pandas.DataFrame
        All parties, with their ID, code, and full name.
    gases : pandas.DataFrame
        The available gases and their IDs.
    annex_one_reader : UNFCCCSingleCategoryApiReader
        The API reader object for Annex I parties.
    non_annex_one_reader : UNFCCCSingleCategoryApiReader
        The API reader object for non-Annex I parties.
    """

    def __init__(self, *, base_url="https://di.unfccc.int/api/"):
        """
        Parameters
        ----------
        base_url : str
            Location of the UNFCCC api.
        """
        self.annex_one_reader = UNFCCCSingleCategoryApiReader(
            party_category="annexOne", base_url=base_url
        )
        self.non_annex_one_reader = UNFCCCSingleCategoryApiReader(
            party_category="nonAnnexOne", base_url=base_url
        )

        self.parties = pd.concat(
            [self.annex_one_reader.parties, self.non_annex_one_reader.parties]
        ).sort_index()
        self.gases = pd.concat(
            [self.annex_one_reader.gases, self.non_annex_one_reader.gases]
        ).sort_index()
        # drop duplicated gases
        self.gases = self.gases[~self.gases.index.duplicated(keep="first")]

    def query(
        self,
        *,
        party_code: str,
        gases: typing.Union[typing.List[str], None] = None,
        progress: bool = False,
    ) -> pd.DataFrame:
        """Query the UNFCCC for data.

        Parameters
        ----------
        party_code : str
            ISO code of a party for which to query. For possible values, see
            :py:attr:`~UNFCCCApiReader.parties`.
        gases : list of str, optional
            Limit the query to these gases. For possible values, see
            :py:attr:`~UNFCCCApiReader.gases`. Default: query for all gases.
        progress : bool
            Display a progress bar. Requires the :py:mod:`tqdm` library. Default: false.

        Returns
        -------
        pandas.DataFrame

        Notes
        -----
        If you need more fine-grained control over which variables to query for,
        including restricting the query to specific measures, categories, or
        classifications or to query for multiple parties at once, please see the
        corresponding methods :py:meth:`UNFCCCApiReader.annex_one_reader.query` and
        :py:meth:`UNFCCCApiReader.non_annex_one_reader.query`.
        """
        if party_code in self.annex_one_reader.parties["code"].values:
            reader = self.annex_one_reader
        elif party_code in self.non_annex_one_reader.parties["code"].values:
            reader = self.non_annex_one_reader
        else:
            raise KeyError(party_code)

        return reader.query(party_codes=[party_code], gases=gases, progress=progress)


class UNFCCCSingleCategoryApiReader:
    """Provides access to the Flexible Query API of the UNFCCC data access for a single
    category, either annexOne or nonAnnexOne.

    Use this class if you want to do fine-grained queries for specific measures,
    categories, years, or classifications.

    Essentially encapsulates https://di.unfccc.int/flex_non_annex1
    or https://di.unfccc.int/flex_annex1 .

    Attributes
    ----------
    parties : pandas.DataFrame
        All parties in this category, with their ID, code, and full name.
    years : pandas.DataFrame
        All years for which data is available, mapping the ID to the year.
    category_tree : treelib.Tree
        The available categories and their relationships. Use
        :py:meth:`~UNFCCCSingleCategoryApiReader.show_category_hierarchy` for displaying
        the category tree.
    classifications : pandas.DataFrame
        All classifications and their IDs.
    measure_tree : treelib.Tree
        The available measures and their relationsips. Use
        :py:meth:`~UNFCCCSingleCategoryApiReader.show_measure_hierarchy` for displaying
        the measure tree.
    gases : pandas.DataFrame
        The available gases and their IDs.
    units : pandas.DataFrame
        The available units and their IDs.
    conversion_factors : pandas.DataFrame
        Conversion factors between units for the specified gases.
    variables : pandas.DataFrame
        The available variables with the corresponding category, classification,
        measure, gas, and unit.
    """

    def __init__(self, *, party_category: str, base_url="https://di.unfccc.int/api/"):
        """
        Parameters
        ----------
        party_category : str
           Either ``nonAnnexOne`` or ``annexOne``.
        base_url : str
           Location of the UNFCCC api.
        """
        self.base_url = base_url

        parties_raw = self._get(f"parties/{party_category}")
        parties_entries = []
        for entry in parties_raw:
            if entry["categoryCode"] == party_category and entry["name"] != "Groups":
                parties_entries.append(entry["parties"])
        if not parties_entries:
            raise ValueError(
                f"Could not find parties for the party_category {party_category!r}."
            )

        self.parties = (
            pd.DataFrame(itertools.chain(*parties_entries))
            .set_index("id")
            .sort_index()
            .drop_duplicates()
        )
        self._parties_dict = dict(self.parties["code"])
        self.years = (
            pd.DataFrame(self._get("years/single")[party_category])
            .set_index("id")
            .sort_index()
        )
        self._years_dict = dict(self.years["name"])
        for i in self._years_dict:
            if self._years_dict[i].startswith("Last Inventory Year"):
                self._years_dict[i] = self._years_dict[i][-5:-1]

        # note that category names are not unique!
        category_hierarchy = self._get("dimension-instances/category")[party_category][
            0
        ]
        self.category_tree = self._walk(category_hierarchy)

        self.classifications = (
            pd.DataFrame(
                self._get("dimension-instances/classification")[party_category]
            )
            .set_index("id")
            .sort_index()
        )
        self._classifications_dict = dict(self.classifications["name"])

        measure_hierarchy = self._get("dimension-instances/measure")[party_category]
        self.measure_tree = treelib.Tree()
        sr = self.measure_tree.create_node("__root__")
        for i in range(len(measure_hierarchy)):
            self._walk(measure_hierarchy[i], tree=self.measure_tree, parent=sr)

        self.gases = (
            pd.DataFrame(self._get("dimension-instances/gas")[party_category])
            .set_index("id")
            .sort_index()
        )
        self._gases_dict = dict(self.gases["name"])

        unit_info = self._get("conversion/fq")
        self.units = pd.DataFrame(unit_info["units"]).set_index("id").sort_index()
        self._units_dict = dict(self.units["name"])
        self.conversion_factors = pd.DataFrame(unit_info[party_category])

        # variable IDs are not unique, because category names are not unique
        # just give up and delete the duplicated ones
        variables_raw = self._get(f"variables/fq/{party_category}")
        self.variables = (
            pd.DataFrame(variables_raw).set_index("variableId").sort_index()
        )
        self.variables = self.variables[~self.variables.index.duplicated(keep="first")]
        self._variables_dict = {x["variableId"]: x for x in variables_raw}

    def _flexible_query(
        self,
        *,
        variable_ids: typing.List[int],
        party_ids: typing.List[int],
        year_ids: typing.List[int],
    ) -> typing.List[dict]:

        if len(variable_ids) > 3000:
            logging.warning(
                "Your query parameters lead to a lot of variables selected at once. "
                "If the query fails, try restricting your query more."
            )

        return self._post(
            "records/flexible-queries",
            json={
                "variableIds": variable_ids,
                "partyIds": party_ids,
                "yearIds": year_ids,
            },
        )

    def query(
        self,
        *,
        party_codes: typing.List[str],
        category_ids: typing.Union[None, typing.List[int]] = None,
        classifications: typing.Union[None, typing.List[str]] = None,
        measure_ids: typing.Union[None, typing.List[int]] = None,
        gases: typing.Union[None, typing.List[str]] = None,
        batch_size: int = 1000,
        progress: bool = False,
    ) -> pd.DataFrame:
        """Query the UNFCCC for data.

        Parameters
        ----------
        party_codes : list of str
            List of ISO codes of parties for which to query. For possible values, see
            :py:attr:`~UNFCCCSingleCategoryApiReader.parties`.
        category_ids : list of int, optional
            List of category IDs to query. For possible values, see
            :py:meth:`~UNFCCCSingleCategoryApiReader.show_category_hierarchy()`.
            Default: query for all categories.
        classifications : list of str, optional
            List of classifications to query. For possible values, see
            :py:attr:`~UNFCCCSingleCategoryApiReader.classifications`.
            Default: query for all classifications.
        measure_ids : list of int, optional
            List of measure IDs to query. For possible values, see
            :py:meth:`~UNFCCCSingleCategoryApiReader.show_measure_hierarchy()`.
            Default: query for all measures.
        gases : list of str, optional
            Limit the query to these gases. For possible values, see
            :py:attr:`~UNFCCCApiReader.gases`. Default: query for all gases.
        batch_size : int, optional
            Number of variables to query in a single API query in the same batch to
            avoid internal server errors. Larger queries are split automatically.
            The default is 1000, which seems to work fine.
        progress : bool
            Display a progress bar. Requires the :py:mod:`tqdm` library. Default: false.


        Returns
        -------
        pandas.DataFrame

        Notes
        -----
        Further documentation about the meaning of parties, categories, classifications,
        measures and gases is available at the `UNFCCC documentation`_.

        .. _UNFCCC documentation: https://unfccc.int/process-and-meetings/\
transparency-and-reporting/greenhouse-gas-data/data-interface-help#eq-7
        """
        party_ids = []
        for code in party_codes:
            party_ids.append(self._name_id(self.parties, code, key="code"))

        # always query all years
        year_ids = list(self.years.index)

        variable_ids = self._select_variable_ids(
            classifications, category_ids, measure_ids, gases
        )

        i = 0
        raw_response = []
        if progress:
            import tqdm

            pbar = tqdm.tqdm(total=len(variable_ids))

        while i < len(variable_ids):
            batched_variable_ids = variable_ids[i : i + batch_size]
            i += batch_size

            batched_response = self._flexible_query(
                variable_ids=batched_variable_ids,
                party_ids=party_ids,
                year_ids=year_ids,
            )
            raw_response += batched_response
            if progress:
                pbar.update(len(batched_variable_ids))

        if progress:
            pbar.close()

        return self._parse_raw_answer(raw_response)

    def _parse_raw_answer(self, raw: typing.List[dict]) -> pd.DataFrame:
        data = []
        for dp in raw:
            variable = self._variables_dict[dp["variableId"]]

            try:
                category = self.category_tree[variable["categoryId"]].tag
            except treelib.tree.NodeIDAbsentError:
                category = f'unknown category nr. {variable["categoryId"]}'

            row = {
                "party": self._parties_dict[dp["partyId"]],
                "category": category,
                "classification": self._classifications_dict[
                    variable["classificationId"]
                ],
                "measure": self.measure_tree[variable["measureId"]].tag,
                "gas": self._gases_dict[variable["gasId"]],
                "unit": self._units_dict[variable["unitId"]],
                "year": self._years_dict[dp["yearId"]],
                "numberValue": dp["numberValue"],
                "stringValue": dp["stringValue"],
            }
            data.append(row)

        df = pd.DataFrame(data)
        df.sort_values(
            ["party", "category", "classification", "measure", "gas", "unit", "year"],
            inplace=True,
        )
        df.drop_duplicates(inplace=True)
        df.reset_index(inplace=True, drop=True)

        return df

    def _select_variable_ids(
        self, classifications, category_ids, measure_ids, gases
    ) -> typing.List[int]:
        # select variables from classification
        if classifications is None:
            classification_mask = pd.Series(
                data=[True] * len(self.variables), index=self.variables.index
            )
        else:
            classification_mask = pd.Series(
                data=[False] * len(self.variables), index=self.variables.index
            )
            for classification in classifications:
                cid = self._name_id(self.classifications, classification)
                classification_mask[self.variables["classificationId"] == cid] = True

        # select variables from categories
        if category_ids is None:
            category_mask = pd.Series(
                data=[True] * len(self.variables), index=self.variables.index
            )
        else:
            category_mask = pd.Series(
                data=[False] * len(self.variables), index=self.variables.index
            )
            for cid in category_ids:
                category_mask[self.variables["categoryId"] == cid] = True

        # select variables from measures
        if measure_ids is None:
            measure_mask = pd.Series(
                data=[True] * len(self.variables), index=self.variables.index
            )
        else:
            measure_mask = pd.Series(
                data=[False] * len(self.variables), index=self.variables.index
            )
            for mid in measure_ids:
                measure_mask[self.variables["measureId"] == mid] = True

        # select variables from gases
        if gases is None:
            gas_mask = pd.Series(
                data=[True] * len(self.variables), index=self.variables.index
            )
        else:
            gas_mask = pd.Series(
                data=[False] * len(self.variables), index=self.variables.index
            )
            for gas in gases:
                gid = self._name_id(self.gases, gas)
                gas_mask[self.variables["gasId"] == gid] = True

        selected_variables = self.variables[
            classification_mask & category_mask & measure_mask & gas_mask
        ]
        return [int(x) for x in selected_variables.index]

    @staticmethod
    def _name_id(df, name, key="name"):
        try:
            return int(df[df[key] == name].index[0])
        except IndexError:
            raise KeyError(name)

    def show_category_hierarchy(self) -> None:
        """Print the hierarchy of categories and their IDs."""
        return self.category_tree.show(idhidden=False)

    def show_measure_hierarchy(self) -> None:
        """Print the hierarchy of measures and their IDs."""
        return self.measure_tree.show(idhidden=False)

    @classmethod
    def _walk(cls, node: dict, tree: treelib.Tree = None, parent=None) -> treelib.Tree:
        if tree is None:
            tree = treelib.Tree()

        tree.create_node(tag=node["name"], identifier=node["id"], parent=parent)

        if "children" in node:
            for child in node["children"]:
                cls._walk(child, tree=tree, parent=node["id"])

        return tree

    def _get(self, component: str) -> typing.Union[dict, list]:
        resp = requests.get(self.base_url + component)
        resp.raise_for_status()
        return resp.json()

    def _post(self, component: str, json: dict) -> typing.List[dict]:
        resp = requests.post(self.base_url + component, json=json)
        resp.raise_for_status()
        return resp.json()
