"""
This is inherited from pymatgen.core.structure
"""

from pymatgen.core import Structure
import os
import warnings
import numpy as np
from pymatgen.core.lattice import Lattice
from pymatgen.core.sites import PeriodicSite
from pymatgen.util.coord import lattice_points_in_supercell
from monty.io import zopen


class StructureKMCpy(Structure):
    @classmethod
    def from_cif(
        cls, filename, primitive=False, sort=False, merge_tol=0.0
    ) -> "StructureKMCpy":
        """this is a modified version of Istructure/Structure.from_file which only take cif input. What change here is that, I will add a tuple, (atom_Label,) in format (str,int) as site_properties. This is for the KMCPY usage. This from_cif will call a seperated cifparser function, which is get_labeled_structure

        Args:
            filename (str): filename path
            primitive (bool, optional): not implemented?. Defaults to False.
            sort (bool, optional): not implemented?. Defaults to False.
            merge_tol (float, optional): not implemented?. Defaults to 0.0.

        Returns:
            kmcpy.external.structure.StructureKMCpy: a structure generated by the cif with all information you need.
        """
        # modified version of from file, only read the cif
        filename = str(filename)

        from kmcpy.external.cif import CifParserKMCpy

        fname = os.path.basename(filename)
        with zopen(filename, "rt") as f:
            contents = f.read()
            parser = CifParserKMCpy.from_string(contents)
            s = parser.get_labeled_structures(primitive=primitive)[0]
        return cls.from_sites(s)

    @classmethod
    def from_sites(
        cls,
        sites: list[PeriodicSite],
        charge: float | None = None,
        validate_proximity: bool = False,
        to_unit_cell: bool = False,
        properties=None,
    ) -> "StructureKMCpy":
        """
        Convenience constructor to make a StructureKMCpy from a list of sites.

        Args:
            sites: Sequence of PeriodicSites. Sites must have the same
                lattice.
            charge: Charge of structure.
            validate_proximity (bool): Whether to check if there are sites
                that are less than 0.01 Ang apart. Defaults to False.
            to_unit_cell (bool): Whether to translate sites into the unit
                cell.

        Returns:
            (StructureKMCpy) Note that missing properties are set as None.
        """
        if len(sites) < 1:
            raise ValueError(f"You need at least one site to construct a {cls}")
        prop_keys: list[str] = []
        props = properties or {}
        lattice = sites[0].lattice
        for i, site in enumerate(sites):
            if site.lattice != lattice:
                raise ValueError("Sites must belong to the same lattice")
            for k, v in site.properties.items():
                if k not in prop_keys:
                    prop_keys.append(k)
                    props[k] = [None] * len(sites)
                props[k][i] = v
        for k, v in props.items():
            if any(vv is None for vv in v):
                warnings.warn(
                    f"Not all sites have property {k}. Missing values are set to None."
                )
        return cls(
            lattice,
            [site.species for site in sites],
            [site.frac_coords for site in sites],
            charge=charge,
            site_properties=props,
            validate_proximity=validate_proximity,
            to_unit_cell=to_unit_cell,
        )

    def find_site_by_wyckoff_sequence_and_label(self, wyckoff_sequence=0, label="Na2"):
        """Look for the site with specific wyckoff sequence and label. Return the First index of site

        KMCPY function


        Args:
            wyckoff_sequence (int, optional): wyckoff sequence as defined. Defaults to 0.
            label (str, optional): label of the site. Defaults to "Na2".

        Raises:
            ValueError: if wyckoff sequence not found, means that this is not generated by from_cif
            ValueError: _description_

        Returns:
            int: index of site
        """
        if "wyckoff_sequence" not in self.site_properties:
            raise ValueError(
                "wyckoff_sequence not found in site properties.make sure that this structure is generated by kmcpy.StructureKMCpy.from_cif() function!"
            )
        for i in self._sites:
            if (
                i.properties["wyckoff_sequence"] == wyckoff_sequence
                and i.properties["label"] == label
            ):
                return i
        raise ValueError(
            "find_site_by_wyckoff_sequence_and_label:not found. with input parameter: Wyckoff sequence=",
            wyckoff_sequence,
            " label=",
            label,
        )

    def find_site_by_wyckoff_sequence_label_and_supercell(
        self, wyckoff_sequence=0, label="Na2", supercell=(0, 1, 1), return_index=True
    ):
        """Look for the site with specific wyckoff sequence and label and supercell. Return the First index of site (this should be unique)

        KMCPY function

        Args:
            wyckoff_sequence (int, optional): wyckoff sequence as defined. Defaults to 0.
            label (str, optional): label of the site. Defaults to "Na2".
            supercell (tuple): tuple of the supercell. Must be tuple with integer element!
            return_index(bool): if True, return the index of site, if False, return the site object

        Raises:
            ValueError: if wyckoff sequence not found, means that this is not generated by from_cif
            ValueError: _description_

        Returns:
            int: index of site
            or
            pymatgen.core.site
        """
        if "wyckoff_sequence" not in self.site_properties:
            raise ValueError(
                "wyckoff_sequence not found in site properties.make sure that this structure is generated by kmcpy.external.structure.StructureKMCpy.from_cif() function!"
            )
        if return_index:
            for i in range(0, len(self._sites)):
                if (
                    self._sites[i].properties["wyckoff_sequence"] == wyckoff_sequence
                    and self._sites[i].properties["label"] == label
                    and self._sites[i].properties["supercell"] == supercell
                ):
                    return i
        else:
            for i in self._sites:
                if (
                    i.properties["wyckoff_sequence"] == wyckoff_sequence
                    and i.properties["label"] == label
                    and i.properties["supercell"] == supercell
                ):
                    return i
        raise ValueError(
            "find_site_by_wyckoff_sequence_label_and_supercell:not found. with input parameter: Wyckoff sequence=",
            wyckoff_sequence,
            " label=",
            label,
            " supercell",
            supercell,
        )

    def make_kmc_supercell(
        self, scaling_matrix=(1, 2, 3), to_unit_cell=True
    ) -> "StructureKMCpy":
        """kmcpy function, nothing strange, just add a property to show which supercell this site belongs to

        Args:
            scaling_matrix (tuple, optional): scaling matrix . Defaults to (1,2,3).

        Returns:
            Structure: kmcpy.external.structure.StructureKMCpy
        """

        scale_matrix = np.array(scaling_matrix, int)
        if scale_matrix.shape != (3, 3):
            scale_matrix = np.array(scale_matrix * np.eye(3), int)
        new_lattice = Lattice(np.dot(scale_matrix, self._lattice.matrix))

        fractional_translations_for_supercell = lattice_points_in_supercell(
            scale_matrix
        )

        # c_lat = new_lattice.get_cartesian_coords(f_lat)

        # print("flat",fractional_translations_for_supercell,tuple(np.dot(fractional_translations_for_supercell,scale_matrix).astype(int)))#debug

        new_sites = []

        for site in self:

            for (
                fractional_translation_for_supercell
            ) in fractional_translations_for_supercell:

                this_properties = site.properties.copy()
                this_properties["supercell"] = tuple(
                    np.rint(
                        np.dot(fractional_translation_for_supercell, scale_matrix)
                    ).astype(int)
                )
                # print(this_properties,new_lattice.get_cartesian_coords(fractional_translation_for_supercell))
                s = PeriodicSite(
                    site.species,
                    site.coords
                    + new_lattice.get_cartesian_coords(
                        fractional_translation_for_supercell
                    ),
                    new_lattice,
                    properties=this_properties,
                    coords_are_cartesian=True,
                    to_unit_cell=True,
                    skip_checks=True,
                )

                new_sites.append(s)

        new_charge = (
            self._charge * np.linalg.det(scale_matrix) if self._charge else None
        )

        supercell = StructureKMCpy.from_sites(
            new_sites, charge=new_charge, to_unit_cell=to_unit_cell
        )

        if to_unit_cell:
            for site in supercell:
                site.to_unit_cell(in_place=True)

        return supercell

    def kmc_info_to_tuple(self, supercell=(1, 2, 3), label="Na2", wyckoff_sequence=2):
        """
        this is also a KMC specified function. Convert the information of supercell, label, wyckoff sequence into a tuple .This tuple is working as the key to the dictionary from kmc_build_dict. Mainly for accelerating the searching progress

        Args:
            supercell (tuple, optional): belongs to which supercell. Defaults to (1,2,3).
            label (str, optional): label of the atom. Defaults to "Na2".
            wyckoff_sequence (int, optional): wyckoff sequence. Defaults to 2.

        Returns:
            tuple: len=5 tuple
        """
        tmp_list = [supercell[0], supercell[1], supercell[2], label, wyckoff_sequence]
        return tuple(tmp_list)

    def site_index_vector(self, supercell=(1, 2, 3), label="Na2", local_index=2):
        """
        this is also a KMC specified function. Convert the information of supercell, label, local_index into a tuple .This tuple is working as the key to the dictionary from kmc_build_dict. Mainly for accelerating the searching progress

        Used in the version3 of generate_events

        Args:
            supercell (tuple, optional): belongs to which supercell. Defaults to (1,2,3).
            label (str, optional): label of the atom. Defaults to "Na2".
            local_index (int, optional): local_index as identifier. Defaults to 2.

        Returns:
            tuple: len=5 tuple
        """
        tmp_list = [supercell[0], supercell[1], supercell[2], label, local_index]
        return tuple(tmp_list)

    def kmc_build_dict(self, skip_check=True):
        """build a dictionary. This is a KMC specified function_summary_

        Args:
            skip_check (bool, optional): whether to skip the check on whether there is duplicate sites. Defaults to True.

        Raises:
            KeyError: if structure is not generated by KMC compatible structure.from_cif
            KeyError: if there is duplicate sites (this shouldn't happen)

        Returns:
            dict: a dictionary. Key is the tuple with same format as class.kmc_info_to_tuple, Value is the global indices
        """
        for prerequisite in ["supercell", "label", "wyckoff_sequence"]:
            if prerequisite not in self.site_properties:
                raise KeyError(
                    "make sure the structure is generated by structure.from_cif().make_kmc_supercell!"
                )
        indices_dict_from_identifier = {}

        if skip_check:

            for site_index in range(0, len(self._sites)):
                indices_dict_from_identifier[
                    self.kmc_info_to_tuple(
                        supercell=self._sites[site_index].properties["supercell"],
                        label=self._sites[site_index].properties["label"],
                        wyckoff_sequence=self._sites[site_index].properties[
                            "wyckoff_sequence"
                        ],
                    )
                ] = site_index
        else:
            for site_index in range(0, len(self._sites)):
                tmp_key = self.kmc_info_to_tuple(
                    supercell=self._sites[site_index].properties["supercell"],
                    label=self._sites[site_index].properties["label"],
                    wyckoff_sequence=self._sites[site_index].properties[
                        "wyckoff_sequence"
                    ],
                )
                if tmp_key in indices_dict_from_identifier:
                    raise KeyError(
                        "duplicate sites identifies by kmc_build_dict function. This shouldn't happen. Please check this site: ",
                        self._sites.properties,
                    )
                indices_dict_from_identifier[
                    self.kmc_info_to_tuple(
                        supercell=self._sites[site_index].properties["supercell"],
                        label=self._sites[site_index].properties["label"],
                        wyckoff_sequence=self._sites[site_index].properties[
                            "wyckoff_sequence"
                        ],
                    )
                ] = site_index

        return indices_dict_from_identifier

    def kmc_build_dict3(self, skip_check=True):
        """build a dictionary. This is a KMC specified function

        Change from the original kmc_build_dict, use the local_index as the identifier

        Args:
            skip_check (bool, optional): whether to skip the check on whether there is duplicate sites. Defaults to True.

        Raises:
            KeyError: if structure is not generated by KMC compatible structure.from_cif
            KeyError: if there is duplicate sites (this shouldn't happen)

        Returns:
            dict: a dictionary. Key is the tuple with same format as class.kmc_info_to_tuple, Value is the global indices
        """
        for prerequisite in ["supercell", "label", "local_index"]:
            if prerequisite not in self.site_properties:
                raise KeyError(
                    "make sure the structure is generated by structure.from_cif().make_kmc_supercell!"
                )

        indices_dict_from_identifier = {}

        if skip_check:

            for site_index in range(0, len(self._sites)):
                indices_dict_from_identifier[
                    self.site_index_vector(
                        supercell=self._sites[site_index].properties["supercell"],
                        label=self._sites[site_index].properties["label"],
                        local_index=self._sites[site_index].properties["local_index"],
                    )
                ] = site_index
        else:
            for site_index in range(0, len(self._sites)):
                tmp_key = self.site_index_vector(
                    supercell=self._sites[site_index].properties["supercell"],
                    label=self._sites[site_index].properties["label"],
                    local_index=self._sites[site_index].properties["local_index"],
                )

                if tmp_key in indices_dict_from_identifier:
                    raise KeyError(
                        "duplicate sites identifies by kmc_build_dict function. This shouldn't happen. Please check this site: ",
                        self._sites.properties,
                    )

                indices_dict_from_identifier[
                    self.site_index_vector(
                        supercell=self._sites[site_index].properties["supercell"],
                        label=self._sites[site_index].properties["label"],
                        local_index=self._sites[site_index].properties["local_index"],
                    )
                ] = site_index

        return indices_dict_from_identifier
