from collections import defaultdict
from math import ceil

from maggma.core import Builder, Store
from maggma.utils import grouper
from pymatgen.analysis.structure_matcher import ElementComparator, StructureMatcher
from pymatgen.core.structure import Structure

from emmet.builders.settings import EmmetBuildSettings
from emmet.core.provenance import ProvenanceDoc, SNLDict
from emmet.core.utils import get_sg, jsanitize, utcnow

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Iterable


class ProvenanceBuilder(Builder):
    def __init__(
        self,
        materials: Store,
        provenance: Store,
        source_snls: list[Store],
        settings: EmmetBuildSettings | None = None,
        query: dict | None = None,
        **kwargs,
    ):
        """
        Creates provenance from source SNLs and materials

        Args:
            materials: Store of materials docs to tag with SNLs
            provenance: Store to update with provenance data
            source_snls: List of locations to grab SNLs
            query : query on materials to limit search
        """
        self.materials = materials
        self.provenance = provenance
        self.source_snls = source_snls
        self.settings = EmmetBuildSettings.autoload(settings)
        self.query = query or {}
        self.kwargs = kwargs

        materials.key = "material_id"
        provenance.key = "material_id"
        for s in source_snls:
            s.key = "snl_id"

        super().__init__(
            sources=[materials, *source_snls], targets=[provenance], **kwargs
        )

    def ensure_indicies(self):
        self.materials.ensure_index("material_id", unique=True)
        self.materials.ensure_index("formula_pretty")

        self.provenance.ensure_index("material_id", unique=True)
        self.provenance.ensure_index("formula_pretty")

        for s in self.source_snls:
            s.ensure_index("snl_id")
            s.ensure_index("formula_pretty")

    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
        self.ensure_indicies()

        # Find all formulas for materials that have been updated since this
        # builder was last ran
        q = self.query
        updated_materials = self.provenance.newer_in(
            self.materials, criteria=q, exhaustive=True
        )
        forms_to_update = set(
            self.materials.distinct(
                "formula_pretty", {"material_id": {"$in": updated_materials}}
            )
        )

        # Find all new SNL formulas since the builder was last run
        for source in self.source_snls:
            new_snls = self.provenance.newer_in(source)
            forms_to_update |= set(
                source.distinct("formula_pretty", {source.key: {"$in": new_snls}})
            )

        # Now reduce to the set of formulas we actually have
        forms_avail = set(self.materials.distinct("formula_pretty", self.query))
        forms_to_update = forms_to_update & forms_avail

        mat_ids = set(
            self.materials.distinct(
                "material_id", {"formula_pretty": {"$in": list(forms_to_update)}}
            )
        ) & set(updated_materials)

        N = ceil(len(mat_ids) / number_splits)

        self.logger.info(
            f"Found {len(mat_ids)} new/updated systems to distribute to workers "
            f"in {N} chunks."
        )

        for chunk in grouper(mat_ids, N):
            yield {"query": {"material_id": {"$in": chunk}}}

    def get_items(self) -> tuple[list[dict], list[dict]]:  # type: ignore
        """
        Gets all materials to assocaite with SNLs
        Returns:
            generator of materials and SNLs that could match
        """
        self.logger.info("Provenance Builder Started")

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

        # Find all formulas for materials that have been updated since this
        # builder was last ran
        q = self.query
        updated_materials = self.provenance.newer_in(
            self.materials, criteria=q, exhaustive=True
        )
        forms_to_update = set(
            self.materials.distinct(
                "formula_pretty", {"material_id": {"$in": updated_materials}}
            )
        )

        # Find all new SNL formulas since the builder was last run
        for source in self.source_snls:
            new_snls = self.provenance.newer_in(source)
            forms_to_update |= set(
                source.distinct("formula_pretty", {source.key: {"$in": new_snls}})
            )

        # Now reduce to the set of formulas we actually have
        forms_avail = set(self.materials.distinct("formula_pretty", self.query))
        forms_to_update = forms_to_update & forms_avail

        mat_ids = set(
            self.materials.distinct(
                "material_id", {"formula_pretty": {"$in": list(forms_to_update)}}
            )
        ) & set(updated_materials)

        self.total = len(mat_ids)

        self.logger.info(f"Found {self.total} new/updated systems to process")

        for mat_id in mat_ids:
            mat = self.materials.query_one(
                properties=[
                    "material_id",
                    "last_updated",
                    "structure",
                    "initial_structures",
                    "formula_pretty",
                    "deprecated",
                ],
                criteria={"material_id": mat_id},
            )

            snls = []  # type: list
            for source in self.source_snls:
                snls.extend(
                    source.query(criteria={"formula_pretty": mat["formula_pretty"]})
                )

            snl_groups = defaultdict(list)
            for snl in snls:
                struc = Structure.from_dict(snl)
                snl_sg = get_sg(struc)
                struc.snl = SNLDict(**snl)  # type: ignore[attr-defined]
                snl_groups[snl_sg].append(struc)

            mat_sg = get_sg(Structure.from_dict(mat["structure"]))

            snl_structs = snl_groups[mat_sg]

            self.logger.debug(f"Found {len(snl_structs)} potential snls for {mat_id}")
            yield mat, snl_structs

    def process_item(self, item) -> dict:
        """
        Matches SNLS and Materials
        Args:
            item (tuple): a tuple of materials and snls
        Returns:
            list(dict): a list of collected snls with material ids
        """
        mat, snl_structs = item
        formula_pretty = mat["formula_pretty"]
        snl_doc = None
        self.logger.debug(f"Finding Provenance {formula_pretty}")

        # Match up SNLS with materials

        matched_snls = self.match(snl_structs, mat)

        if len(matched_snls) > 0:
            doc = ProvenanceDoc.from_SNLs(
                material_id=mat["material_id"],
                structure=Structure.from_dict(mat["structure"]),
                snls=matched_snls,
                deprecated=mat["deprecated"],
            )
        else:
            doc = ProvenanceDoc(  # type: ignore[call-arg]
                material_id=mat["material_id"],
                structure=Structure.from_dict(mat["structure"]),
                deprecated=mat["deprecated"],
                created_at=utcnow(),
            )

        doc.authors.append(self.settings.DEFAULT_AUTHOR)
        doc.history.append(self.settings.DEFAULT_HISTORY)
        doc.references.append(self.settings.DEFAULT_REFERENCE)

        snl_doc = jsanitize(doc.dict(exclude_none=False), allow_bson=True)

        return snl_doc

    def match(self, snl_structs, mat):
        """
        Finds a material doc that matches with the given snl
        Args:
            snl_structs ([dict]): the snls struct list
            mat (dict): a materials doc
        Returns:
            generator of materials doc keys
        """

        m_strucs = [Structure.from_dict(mat["structure"])] + [
            Structure.from_dict(init_struc) for init_struc in mat["initial_structures"]
        ]

        sm = StructureMatcher(
            ltol=self.settings.LTOL,
            stol=self.settings.STOL,
            angle_tol=self.settings.ANGLE_TOL,
            primitive_cell=True,
            scale=True,
            attempt_supercell=False,
            allow_subset=False,
            comparator=ElementComparator(),
        )

        snls = []

        for s in m_strucs:
            for snl_struc in snl_structs:
                if sm.fit(s, snl_struc):
                    if snl_struc.snl not in snls:
                        snls.append(snl_struc.snl)

        self.logger.debug(f"Found {len(snls)} SNLs for {mat['material_id']}")
        return snls

    def update_targets(self, items):
        """
        Inserts the new SNL docs into the SNL collection
        """
        snls = list(filter(None, items))

        if len(snls) > 0:
            self.logger.info(f"Found {len(snls)} SNLs to update")
            self.provenance.update(snls)
        else:
            self.logger.info("No items to update")
