import warnings
from itertools import chain
from typing import Dict, Iterator, List, Optional, Set

from maggma.core import Builder, Store
from maggma.stores import S3Store
from monty.json import MontyDecoder
from pymatgen.analysis.phase_diagram import PhaseDiagramError
from pymatgen.entries.computed_entries import ComputedStructureEntry

from emmet.builders.utils import HiddenPrints
from emmet.core.thermo import ThermoDoc, PhaseDiagramDoc
from emmet.core.utils import jsanitize


class ThermoBuilder(Builder):
    def __init__(
        self,
        thermo: Store,
        corrected_entries: Store,
        phase_diagram: Optional[Store] = None,
        query: Optional[Dict] = None,
        num_phase_diagram_eles: Optional[int] = None,
        chunk_size: int = 1000,
        **kwargs,
    ):
        """
        Calculates thermodynamic quantities for materials from phase
        diagram constructions

        Args:
            thermo (Store): Store of thermodynamic data such as formation
                energy and decomposition pathway
            corrected_entries (Store): Store of corrected entry data to use in thermo data and phase diagram
                construction. This is required and should be built with the CorrectedEntriesBuilder.
            phase_diagram (Store): Store of phase diagram data for each unique chemical system
            query (dict): dictionary to limit materials to be analyzed
            num_phase_diagram_eles (int): Maximum number of elements to use in phase diagram construction
                for data within the separate phase_diagram collection
            chunk_size (int): Size of chemsys chunks to process at any one time.
        """

        self.thermo = thermo
        self.query = query if query else {}
        self.corrected_entries = corrected_entries
        self.phase_diagram = phase_diagram
        self.num_phase_diagram_eles = num_phase_diagram_eles
        self.chunk_size = chunk_size
        self._completed_tasks: Set[str] = set()

        if self.thermo.key != "thermo_id":
            warnings.warn(
                f"Key for the thermo store is incorrect and has been changed from {self.thermo.key} to thermo_id!"
            )
            self.thermo.key = "thermo_id"

        if self.corrected_entries.key != "chemsys":
            warnings.warn(
                "Key for the corrected entries store is incorrect and has been changed "
                f"from {self.corrected_entries.key} to chemsys!"
            )
            self.corrected_entries.key = "chemsys"

        sources = [corrected_entries]
        targets = [thermo]

        if self.phase_diagram is not None:

            if self.phase_diagram.key != "phase_diagram_id":
                warnings.warn(
                    f"Key for the phase diagram store is incorrect and has been changed from {self.phase_diagram.key} to phase_diagram_id!"  # noqa: E501
                )
                self.phase_diagram.key = "phase_diagram_id"

            targets.append(phase_diagram)  # type: ignore

        super().__init__(sources=sources, targets=targets, chunk_size=chunk_size, **kwargs)

    def ensure_indexes(self):
        """
        Ensures indicies on the tasks and materials collections
        """

        # Search index for corrected_entries
        self.corrected_entries.ensure_index("chemsys")
        self.corrected_entries.ensure_index("last_updated")

        # Search index for thermo
        self.thermo.ensure_index("material_id")
        self.thermo.ensure_index("thermo_id")
        self.thermo.ensure_index("thermo_type")
        self.thermo.ensure_index("last_updated")

        # Search index for thermo
        self.thermo.ensure_index("material_id")
        self.thermo.ensure_index("thermo_id")
        self.thermo.ensure_index("thermo_type")
        self.thermo.ensure_index("last_updated")

        # Search index for phase_diagram
        if self.phase_diagram:
            coll = self.phase_diagram

            if isinstance(self.phase_diagram, S3Store):
                coll = self.phase_diagram.index

            coll.ensure_index("chemsys")
            coll.ensure_index("phase_diagram_id")

    def get_items(self) -> Iterator[List[Dict]]:
        """
        Gets whole chemical systems of entries to process
        """

        self.logger.info("Thermo Builder Started")

        self.logger.info("Setting indexes")
        self.ensure_indexes()

        to_process_chemsys = self._get_chemsys_to_process()

        self.logger.info(f"Found {len(to_process_chemsys)} chemical systems with new/updated materials to process")
        self.total = len(to_process_chemsys)

        # Yield the chemical systems in order of increasing size
        # Will build them in a similar manner to fast Pourbaix
        for chemsys in sorted(to_process_chemsys, key=lambda x: len(x.split("-")), reverse=False):
            corrected_entries = self.corrected_entries.query_one({"chemsys": chemsys})
            yield corrected_entries

    def process_item(self, item):

        if not item:
            return None

        for thermo_type, entry_list in item["entries"].items():

            entries = [ComputedStructureEntry.from_dict(entry) for entry in entry_list]
            chemsys = item["chemsys"]
            elements = chemsys.split("-")

            self.logger.debug(f"Processing {len(entries)} entries for {chemsys} and thermo type {thermo_type}")

            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                with HiddenPrints():

                    pd_thermo_doc_pair = self._produce_pair(entries, thermo_type, elements)

        return pd_thermo_doc_pair

    def _produce_pair(self, pd_entries, thermo_type, elements):
        # Produce thermo and phase diagram pair

        try:
            # Obtain phase diagram
            pd = ThermoDoc.construct_phase_diagram(pd_entries)

            # Iterate through entry material IDs and construct list of thermo docs to update
            docs = ThermoDoc.from_entries(pd_entries, thermo_type, pd, deprecated=False)

            pd_docs = [None]

            if self.phase_diagram:
                if self.num_phase_diagram_eles is None or len(elements) <= self.num_phase_diagram_eles:
                    chemsys = "-".join(sorted(set([e.symbol for e in pd.elements])))
                    pd_id = "{}_{}".format(chemsys, str(thermo_type))
                    pd_doc = PhaseDiagramDoc(
                        phase_diagram_id=pd_id, chemsys=chemsys, phase_diagram=pd, thermo_type=thermo_type,
                    )

                    pd_data = jsanitize(pd_doc.dict(), allow_bson=True)

                    pd_docs = [pd_data]

            docs_pd_pair = (
                jsanitize([d.dict() for d in docs], allow_bson=True),
                pd_docs,
            )

            return docs_pd_pair

        except PhaseDiagramError as p:
            elsyms = []
            for e in pd_entries:
                elsyms.extend([el.symbol for el in e.composition.elements])

            self.logger.error(f"Phase diagram error in chemsys {'-'.join(sorted(set(elsyms)))}: {p}")
            return (None, None)

    def update_targets(self, items):
        """
        Inserts the thermo and phase diagram docs into the thermo collection
        Args:
            items ([tuple(List[dict],List[dict])]): a list of thermo dictionaries to update
        """

        thermo_docs = [pair[0] for pair in items]
        phase_diagram_docs = [pair[1] for pair in items]

        # flatten out lists
        thermo_docs = list(filter(None, chain.from_iterable(thermo_docs)))
        phase_diagram_docs = list(filter(None, chain.from_iterable(phase_diagram_docs)))

        # Check if already updated this run
        thermo_docs = [i for i in thermo_docs if i["thermo_id"] not in self._completed_tasks]

        self._completed_tasks |= {i["thermo_id"] for i in thermo_docs}

        for item in thermo_docs:
            if isinstance(item["last_updated"], dict):
                item["last_updated"] = MontyDecoder().process_decoded(item["last_updated"])

        if self.phase_diagram:
            self.phase_diagram.update(phase_diagram_docs)

        if len(thermo_docs) > 0:
            self.logger.info(f"Updating {len(thermo_docs)} thermo documents")
            self.thermo.update(docs=thermo_docs, key=["thermo_id"])
        else:
            self.logger.info("No thermo items to update")

    def _get_chemsys_to_process(self):
        # Use last-updated to find new chemsys
        corrected_entries_chemsys_dates = {
            d[self.corrected_entries.key]: d[self.corrected_entries.last_updated_field]
            for d in self.corrected_entries.query(
                {}, properties=[self.corrected_entries.key, self.corrected_entries.last_updated_field]
            )
        }

        thermo_chemsys_dates = {}
        for d in self.thermo.query(
            {"deprecated": False}, properties=[self.corrected_entries.key, self.thermo.last_updated_field]
        ):

            entry = thermo_chemsys_dates.get(d[self.corrected_entries.key], None)
            if entry is None or d[self.thermo.last_updated_field] < entry:
                thermo_chemsys_dates[d[self.corrected_entries.key]] = d[self.thermo.last_updated_field]

        to_process_chemsys = [
            chemsys
            for chemsys in corrected_entries_chemsys_dates
            if (chemsys not in thermo_chemsys_dates)
            or (thermo_chemsys_dates[chemsys] < corrected_entries_chemsys_dates[chemsys])
        ]

        return to_process_chemsys
