from __future__ import annotations

from emmet.core.molecules.orbitals import OrbitalDoc
from emmet.core.mpid import MPculeID

from mp_api.client.core import BaseRester


class MoleculesOrbitalsRester(BaseRester[OrbitalDoc]):
    suffix = "molecules/orbitals"
    document_model = OrbitalDoc
    primary_key = "property_id"

    def search(
        self,
        molecule_ids: MPculeID | list[MPculeID] | None = None,
        property_ids: str | list[str] | None = None,
        charge: int | None = None,
        spin_multiplicity: int | None = None,
        level_of_theory: str | None = None,
        solvent: str | None = None,
        lot_solvent: str | None = None,
        formula: str | list[str] | None = None,
        elements: list[str] | None = None,
        exclude_elements: list[str] | None = None,
        chemsys: str | list[str] | None = None,
        electron_type_population: str | None = None,
        core_electrons: tuple[float, float] | None = None,
        valence_electrons: tuple[float, float] | None = None,
        rydberg_electrons: tuple[float, float] | None = None,
        total_electrons: tuple[float, float] | None = None,
        electron_type_lp: str | None = None,
        lp_type: str | None = None,
        s_character: tuple[float, float] | None = None,
        p_character: tuple[float, float] | None = None,
        d_character: tuple[float, float] | None = None,
        f_character: tuple[float, float] | None = None,
        lp_occupancy: tuple[float, float] | None = None,
        electron_type_bond: str | None = None,
        bond_type: str | None = None,
        s_character_atom1: tuple[float, float] | None = None,
        s_character_atom2: tuple[float, float] | None = None,
        p_character_atom1: tuple[float, float] | None = None,
        p_character_atom2: tuple[float, float] | None = None,
        d_character_atom1: tuple[float, float] | None = None,
        d_character_atom2: tuple[float, float] | None = None,
        f_character_atom1: tuple[float, float] | None = None,
        f_character_atom2: tuple[float, float] | None = None,
        polarization_atom1: tuple[float, float] | None = None,
        polarization_atom2: tuple[float, float] | None = None,
        bond_occupancy: tuple[float, float] | None = None,
        electron_type_interaction: str | None = None,
        donor_type: str | None = None,
        acceptor_type: str | None = None,
        perturbation_energy: tuple[float, float] | None = None,
        energy_difference: tuple[float, float] | None = None,
        fock_element: tuple[float, float] | None = None,
        num_chunks: int | None = None,
        chunk_size: int = 1000,
        all_fields: bool = True,
        fields: list[str] | None = None,
    ):
        """Query molecules redox docs using a variety of search criteria.

        Arguments:
            molecule_ids (MPculeID, List[MPculeID]): List of Materials Project Molecule IDs (MPculeIDs) to return data
                for.
            property_ids (str, List[str]): List of property IDs to return data for.
            charge (Tuple[int, int]): Minimum and maximum charge for the molecule.
            spin_multiplicity (Tuple[int, int]): Minimum and maximum spin for the molecule.
            level_of_theory (str): Desired level of theory (e.g. "wB97X-V/def2-TZVPPD/SMD")
            solvent (str): Desired solvent (e.g. "SOLVENT=WATER")
            lot_solvent (str): Desired combination of level of theory and solvent
                (e.g. "wB97X-V/def2-TZVPPD/SMD(SOLVENT=THF)")
            correction_level_of_theory (str): Desired correction level of theory (e.g. "wB97X-V/def2-TZVPPD/SMD")
            correction_solvent (str): Desired correction solvent (e.g. "SOLVENT=WATER")
            correction_lot_solvent (str): Desired correction combination of level of theory and solvent
                (e.g. "wB97X-V/def2-TZVPPD/SMD(SOLVENT=THF)")
            combined_lot_solvent (str): Desired combination of level of theory and solvent including both main
                thermo calculation and single-point energy correction
                (e.g. "wB97X-D/def2-SVPD/VACUUM//wB97X-V/def2-TZVPPD/SMD(SOLVENT=THF)")
            formula (str, List[str]): An alphabetical formula or list of formulas
                (e.g. "C2 Li2 O4", ["C2 H4", "C2 H6"]).
            elements (List[str]): A list of elements.
            exclude_elements (List(str)): List of elements to exclude.
            chemsys (str, List[str]): A chemical system, list of chemical systems
                (e.g., Li-C-O, [C-O-H-N, Li-N]).
            electron_type_population (str): Should alpha ('alpha'), beta ('beta'), or all electrons (None) be
                considered in a query of natural electron populations?
            core_electrons (Tuple[float, float]): Minimum and maximum number of core electrons in an atom
                in this molecule.
            valence_electrons (Tuple[float, float]): Minimum and maximum number of valence electrons in an atom
                in this molecule.
            rydberg_electrons (Tuple[float, float]): Minimum and maximum number of rydberg electrons in an atom
                in this molecule.
            total_electrons: (Tuple[float, float]): Minimum and maximum number of total electrons in an atom
                in this molecule.
            electron_type_lp (str): Should alpha ('alpha'), beta ('beta'), or all electrons (None) be
                considered in a query of lone pairs?
            lp_type (str): Type of orbital - "LP" for "lone pair" or "LV" for "lone vacant"
            s_character (Tuple[float, float]): Min and max percentage of the lone pair constituted by s atomic orbitals.
            p_character (Tuple[float, float]): Min and max percentage of the lone pair constituted by p atomic orbitals.
            d_character (Tuple[float, float]): Min and max percentage of the lone pair constituted by d atomic orbitals.
            f_character (Tuple[float, float]): Min and max percentage of the lone pair constituted by f atomic orbitals.
            lp_occupancy(Tuple[float, float]): Min and max number of electrons in the lone pair.
            electron_type_bond (str): Should alpha ('alpha'), beta ('beta'), or all electrons (None) be
                considered in a query of bonds?
            bond_type (str): Type of orbital, e.g. "BD" for bonding or "BD*" for antibonding
            s_character_atom1 (Tuple[float, float]): Min and max percentage of the bond constituted by s atomic orbitals
                on the first atom.
            s_character_atom2 (Tuple[float, float]): Min and max percentage of the bond constituted by s atomic orbitals
                on the second atom.
            p_character_atom1 (Tuple[float, float]): Min and max percentage of the bond constituted by p atomic orbitals
                on the first atom.
            p_character_atom2 (Tuple[float, float]): Min and max percentage of the bond constituted by p atomic orbitals
                on the second atom.
            d_character_atom1 (Tuple[float, float]): Min and max percentage of the bond constituted by d atomic orbitals
                on the first atom.
            d_character_atom2 (Tuple[float, float]): Min and max percentage of the bond constituted by d atomic orbitals
                on the second atom.
            f_character_atom1 (Tuple[float, float]): Min and max percentage of the bond constituted by f atomic orbitals
                on the first atom.
            f_character_atom2 (Tuple[float, float]): Min and max percentage of the bond constituted by f atomic orbitals
                on the second atom.
            polarization_atom1 (Tuple[float, float]): Min and max fraction of electrons in the bond donated by the
                first atom.
            polarization_atom2 (Tuple[float, float]): Min and max fraction of electrons in the bond donated by the
                second atom.
            bond_occupancy (Tuple[float, float]): Min and max number of electrons in the bond.
            electron_type_interaction (str): Should alpha ('alpha'), beta ('beta'), or all electrons (None) be
                considered in a query of orbital interactions?
            donor_type (str): Type of donor orbital, e.g. 'BD' for bonding or 'RY*' for anti-Rydberg
            acceptor_type (str): Type of acceptor orbital, e.g. 'BD' for bonding or 'RY*' for anti-Rydberg
            perturbation_energy (Tuple[float, float]): Min and max perturbation energy of the interaction
            energy_difference (Tuple[float, float]): Min and max energy difference between interacting orbitals
            fock_element (Tuple[float, float]): Min and max interaction Fock matrix element
            num_chunks (int): Maximum number of chunks of data to yield. None will yield all possible.
            chunk_size (int): Number of data entries per chunk.
            all_fields (bool): Whether to return all fields in the document. Defaults to True.
            fields (List[str]): List of fields in MoleculeDoc to return data for.
                Default is "molecule_id", "property_id", "solvent", "method", "last_updated"
                if all_fields is False.

        Returns:
            ([RedoxDoc]) List of molecule redox documents
        """
        query_params = {}  # type: dict

        min_max = [
            "core_electrons",
            "valence_electrons",
            "rydberg_electrons",
            "total_electrons",
            "s_character",
            "p_character",
            "d_character",
            "f_character",
            "lp_occupancy",
            "s_character_atom1",
            "s_character_atom2",
            "p_character_atom1",
            "p_character_atom2",
            "d_character_atom1",
            "d_character_atom2",
            "f_character_atom1",
            "f_character_atom2",
            "polarization_atom1",
            "polarization_atom2",
            "bond_occupancy",
            "perturbation_energy",
            "energy_difference",
            "fock_element",
        ]

        for param, value in locals().items():
            if param in min_max and value:
                if isinstance(value, (int, float)):
                    value = (value, value)
                query_params.update(
                    {
                        f"min_{param}": value[0],
                        f"max_{param}": value[1],
                    }
                )

        if molecule_ids:
            if isinstance(molecule_ids, str):
                molecule_ids = [molecule_ids]

            query_params.update({"molecule_ids": ",".join(molecule_ids)})

        if property_ids:
            if isinstance(property_ids, str):
                property_ids = [property_ids]

            query_params.update({"property_ids": ",".join(property_ids)})

        if charge:
            query_params.update({"charge": charge})

        if spin_multiplicity:
            query_params.update({"spin_multiplicity": spin_multiplicity})

        if level_of_theory:
            query_params.update({"level_of_theory": level_of_theory})

        if solvent:
            query_params.update({"solvent": solvent})

        if lot_solvent:
            query_params.update({"lot_solvent": lot_solvent})

        if formula:
            if isinstance(formula, str):
                formula = [formula]

            query_params.update({"formula": ",".join(formula)})

        if chemsys:
            if isinstance(chemsys, str):
                chemsys = [chemsys]

            query_params.update({"chemsys": ",".join(chemsys)})

        if elements:
            query_params.update({"elements": ",".join(elements)})

        if exclude_elements:
            query_params.update({"exclude_elements": ",".join(exclude_elements)})

        if electron_type_population:
            query_params.update({"electron_type_population": electron_type_population})

        if electron_type_lp:
            query_params.update({"electron_type_lp": electron_type_lp})

        if lp_type:
            query_params.update({"lp_type": lp_type})

        if electron_type_bond:
            query_params.update({"electron_type_bond": electron_type_bond})

        if bond_type:
            query_params.update({"bond_type": bond_type})

        if electron_type_interaction:
            query_params.update(
                {"electron_type_interaction": electron_type_interaction}
            )

        if donor_type:
            query_params.update({"donor_type": donor_type})

        if acceptor_type:
            query_params.update({"acceptor_type": acceptor_type})

        return super()._search(
            num_chunks=num_chunks,
            chunk_size=chunk_size,
            all_fields=all_fields,
            fields=fields,
            **query_params,
        )
