import os
import time
from collections import Counter

import numpy as np
import pandas as pd
from citrination_client import (
    ChemicalFieldQuery,
    ChemicalFilter,
    CitrinationClient,
    DataQuery,
    DatasetQuery,
    FieldQuery,
    Filter,
    PifSystemQuery,
    PifSystemReturningQuery,
    PropertyQuery,
    ReferenceQuery,
)

try:
    from pandas.io.json import json_normalize
except ImportError:
    from pandas import json_normalize
from tqdm import tqdm

from matminer.data_retrieval.retrieve_base import BaseDataRetrieval

__author__ = [
    "Saurabh Bajaj <sbajaj@lbl.gov>",
    "Alireza Faghaninia <alireza.faghaninia@gmail.com>",
]


def parse_scalars(scalars):
    return get_value(scalars[0])


def get_value(dict_item):
    # TODO: deal with rest of formats in a scalar object
    if "value" in dict_item:
        return dict_item["value"]
    elif "minimum" in dict_item and "maximum" in dict_item:
        return "Minimum = {}, Maximum = {}".format(dict_item["minimum"], dict_item["maximum"])


class CitrineDataRetrieval(BaseDataRetrieval):
    """
    CitrineDataRetrieval is used to retrieve data from the Citrination database
    See API client docs at api_link below.
    """

    def __init__(self, api_key=None):
        """
        Args:
            api_key: (str) Your Citrine API key, or None if
                you've set the CITRINE_KEY environment variable
        """
        if api_key:
            api_key = api_key
        elif "CITRINATION_API_KEY" in os.environ:
            api_key = os.environ["CITRINATION_API_KEY"]
        elif "CITRINE_KEY" in os.environ:
            api_key = os.environ["CITRINE_KEY"]
        else:
            raise AttributeError(
                """Citrine API key not found.

            You need to get an API key from Citrination, and either supply it as an argument to
            this class or set it as the value of the CITRINATION_API_KEY environment variable

            See https://citrineinformatics.github.io/api-documentation/quickstart/index.html
            for details on how to get an API key"""
            )

        self.client = CitrinationClient(api_key, "https://citrination.com")

    def api_link(self):
        return "https://citrineinformatics.github.io/python-citrination-client/"

    def get_dataframe(
        self,
        criteria,
        properties=None,
        common_fields=None,
        secondary_fields=False,
        print_properties_options=True,
    ):
        """
        Gets a Pandas dataframe object from data retrieved from
        the Citrine API.

        Args:
            criteria (dict): see get_data method for supported keys except
                prop; prop should be included in properties.
            properties ([str]): requested properties/fields/columns.
                For example, ["Seebeck coefficient", "Band gap"]. If unsure
                about the exact words, capitalization, etc try something like
                ["gap"] and "max_results": 3 and print_properties_options=True
                to see the exact options for this field
            common_fields ([str]): fields that are common to all the requested
                properties. Common example can be "chemicalFormula". Look for
                suggested common fields after a quick query for more info
            secondary_fields (bool): if True, fields not included in properties
                may be added to the output (e.g. references). Recommended only
                if len(properties)==1
            print_properties_options (bool): whether to print available options
                for "properties" and "common_fields" arguments.
        Returns: (object) Pandas dataframe object containing the results

        """
        common_fields = common_fields or []
        properties = properties or [None]
        if criteria.get("prop"):
            properties.append(criteria.pop("prop"))
            properties = list(set(properties))
        all_fields = []
        for prop_counter, requested_prop in enumerate(properties):
            jsons = self.get_data(prop=requested_prop, **criteria)
            non_prop_df = pd.DataFrame()  # df w/o measurement column
            prop_df = pd.DataFrame()  # df containing only measurement column
            counter = 0  # variable to keep count of sample hit and set indexes
            for hit in tqdm(jsons):
                counter += 1
                if "system" in hit.keys():  # Check if 'system' key exists, else skip
                    system_value = hit["system"]
                    system_normdf = json_normalize(system_value)
                    non_prop_cols = [cols for cols in system_normdf.columns if "properties" not in cols]
                    non_prop_row = pd.DataFrame()
                    for col in non_prop_cols:
                        non_prop_row[col] = system_normdf[col]
                    non_prop_row.index = [counter] * len(system_normdf)
                    non_prop_df = non_prop_df.append(non_prop_row)
                    if "properties" in system_value:
                        p_df = pd.DataFrame()
                        # Rename duplicate property names in a record with progressive numbering
                        all_prop_names = [x["name"] for x in system_value["properties"]]
                        counts = {k: v for k, v in Counter(all_prop_names).items() if v > 1}
                        for i in reversed(range(len(all_prop_names))):
                            item = all_prop_names[i]
                            if item in counts and counts[item]:
                                all_prop_names[i] += "_" + str(counts[item])
                                counts[item] -= 1

                        # add each property, and its associated fields, as a new column
                        for p_idx, prop in enumerate(system_value["properties"]):
                            # Rename property name according to above duplicate numbering
                            prop["name"] = all_prop_names[p_idx]
                            if "scalars" in prop:
                                p_df.at[counter, prop["name"]] = parse_scalars(prop["scalars"])
                            elif "vectors" in prop:
                                p_df[prop["name"]] = prop["vectors"]
                            elif "matrices" in prop:
                                p_df[prop["name"]] = prop["matrices"]
                            # parse all keys in the Property object except 'name', 'scalars', 'vectors', and 'matrices'
                            for prop_key in prop:
                                if prop_key not in [
                                    "name",
                                    "scalars",
                                    "vectors",
                                    "matrices",
                                ]:
                                    # If value is a list of multiple items, set the cell to the entire list by first
                                    if isinstance(prop[prop_key], list) and len(prop[prop_key]) > 1:
                                        p_df[prop["name"] + "-" + prop_key] = np.nan
                                        p_df[prop["name"] + "-" + prop_key] = p_df[
                                            prop["name"] + "-" + prop_key
                                        ].astype(object)
                                    p_df.at[counter, prop["name"] + "-" + prop_key] = prop[prop_key]
                        p_df.index = [counter]
                        prop_df = prop_df.append(p_df)
            df_prop = pd.concat([non_prop_df, prop_df], axis=1)
            if prop_counter == 0:
                optcomcols = df_prop.columns.values
            else:
                optcomcols = list(set(optcomcols) & set(df_prop.columns))
            all_fields += list(df_prop.columns.values)
            if not secondary_fields:
                outcols = []
                for p in df_prop.columns.values:
                    if not requested_prop or requested_prop in p:
                        outcols.append(p)
                df_prop = df_prop[outcols + ["uid"] + common_fields]
            if prop_counter == 0:
                df = df_prop
            else:
                if "uid" in df and "uid" in df_prop:
                    uid = ["uid"]
                else:
                    uid = []
                try:
                    if len(uid + common_fields) > 0:
                        df = df.merge(df_prop, on=uid + common_fields, how="outer")
                    else:
                        df = df.join(df_prop, how="outer")
                except (TypeError, KeyError):
                    raise TypeError(
                        "Use scalar/string fields for common_fields" "common_fields among: {}".format(optcomcols)
                    )
        uninformative_columns = ["category", "uid"]
        optcomcols = [c for c in optcomcols if c not in uninformative_columns]
        for col in uninformative_columns:
            if col in df:
                df = df.drop(col, axis=1)
        if print_properties_options:
            print(f"all available fields:\n{list(set(all_fields))}")
            print(f"\nsuggested common fields:\n{optcomcols}")
        return df

    def get_data(
        self,
        formula=None,
        prop=None,
        data_type=None,
        reference=None,
        min_measurement=None,
        max_measurement=None,
        from_record=None,
        data_set_id=None,
        max_results=None,
    ):
        """
        Gets raw api data from Citrine in json format. See api_link for more
        information on input parameters

        Args:
            formula: (str) filter for the chemical formula field; only those
                results that have chemical formulas that contain this string
                will be returned
            prop: (str) name of the property to search for
            data_type: (str) 'EXPERIMENTAL'/'COMPUTATIONAL'/'MACHINE_LEARNING';
                filter for properties obtained from experimental work,
                computational methods, or machine learning.
            reference: (str) filter for the reference field; only those
                results that have contributors that contain this string
                will be returned
            min_measurement: (str/num) minimum of the property value range
            max_measurement: (str/num) maximum of the property value range
            from_record: (int) index of first record to return (indexed from 0)
            data_set_id: (int) id of the particular data set to search on
            max_results: (int) number of records to limit the results to

        Returns: (list) of jsons/pifs returned by Citrine's API
        """

        json_data = []
        start = from_record if from_record else 0
        per_page = 100
        refresh_time = 3  # seconds to wait between search calls

        # Construct all of the relevant queries from input args
        formula_query = ChemicalFieldQuery(filter=ChemicalFilter(equal=formula))
        prop_query = PropertyQuery(
            name=FieldQuery(filter=Filter(equal=prop)),
            value=FieldQuery(filter=Filter(min=min_measurement, max=max_measurement)),
            data_type=FieldQuery(filter=Filter(equal=data_type)),
        )
        ref_query = ReferenceQuery(doi=FieldQuery(filter=Filter(equal=reference)))

        system_query = PifSystemQuery(chemical_formula=formula_query, properties=prop_query, references=ref_query)
        dataset_query = DatasetQuery(id=Filter(equal=data_set_id))
        data_query = DataQuery(system=system_query, dataset=dataset_query)

        while True:
            # use per_page=max_results, eg: in case of max_results=68 < 100
            if max_results and max_results < per_page:
                pif_query = PifSystemReturningQuery(query=data_query, from_index=start, size=max_results)
            else:
                pif_query = PifSystemReturningQuery(query=data_query, from_index=start, size=per_page)

            # Check if any results found
            if "hits" not in self.client.search.pif_search(pif_query).as_dictionary():
                raise KeyError("No results found!")

            data = self.client.search.pif_search(pif_query).as_dictionary()["hits"]
            size = len(data)
            start += size
            json_data.extend(data)

            # check if limit is reached
            if max_results and len(json_data) > max_results:
                # get first multiple of 100 records
                json_data = json_data[:max_results]
                break
            if size < per_page:  # break out of last loop of results
                break
            time.sleep(refresh_time)
        return json_data

    def citations(self):
        return [
            "@article{OMara2016,"
            "doi = {10.1007/s11837-016-1984-0},"
            "url = {https://doi.org/10.1007/s11837-016-1984-0},"
            "year = {2016},"
            "month = jun,"
            "publisher = {Springer Science and Business Media {LLC}},"
            "volume = {68},"
            "number = {8},"
            "pages = {2031--2034},"
            "author = {Jordan O'Mara and Bryce Meredig and Kyle Michel},"
            "title = {Materials Data Infrastructure: A Case Study of the Citrination Platform to Examine Data Import,  Storage,  and Access},"
            "journal = {{JOM}}"
            "}"
        ]
