from typing import Dict, List, Any, Optional, Tuple
import copy

from pydantic import Field
import networkx as nx

from pymatgen.core.structure import Molecule
from pymatgen.analysis.graphs import MoleculeGraph
from pymatgen.analysis.local_env import OpenBabelNN, metal_edge_extender


from emmet.core.mpid import MPID
from emmet.core.qchem.task import TaskDocument
from emmet.core.material import PropertyOrigin
from emmet.core.molecules.molecule_property import PropertyDoc


__author__ = "Evan Spotte-Smith <ewcspottesmith@lbl.gov>"


metals = ["Li", "Mg", "Ca", "Zn", "Al"]

BOND_METHODS = ["nbo", "critic2", "OpenBabelNN + metal_edge_extender"]


def fix_C_Li_bonds(critic: Dict) -> Dict:
    """
    Adjust C-Li coordinate bonding for Critic2 calculations.

    :param critic: Critic2 output dictionary

    :return:
        critic: modified Critic2 output dictionary

    """
    for key in critic["bonding"]:
        if critic["bonding"][key]["atoms"] == ["Li", "C"] or critic["bonding"][key][
            "atoms"
        ] == ["C", "Li"]:
            if (
                critic["bonding"][key]["field"] <= 0.02
                and critic["bonding"][key]["field"] > 0.012
                and critic["bonding"][key]["distance"] < 2.5
            ):
                critic["processed"]["bonds"].append(
                    [int(entry) - 1 for entry in critic["bonding"][key]["atom_ids"]]
                )
    return critic


def make_mol_graph(
    mol: Molecule, critic_bonds: Optional[List[List[int]]] = None
) -> MoleculeGraph:
    """
    Construct a MoleculeGraph using OpenBabelNN with metal_edge_extender and
    (optionally) Critic2 bonding information.

    This bonding scheme was used to define bonding for the Lithium-Ion Battery
    Electrolyte (LIBE) dataset (DOI: 10.1038/s41597-021-00986-9)

    :param mol: Molecule to be converted to MoleculeGraph
    :param critic_bonds: (optional) List of lists [a, b], where a and b are
        atom indices (0-indexed)

    :return: mol_graph, a MoleculeGraph
    """
    mol_graph = MoleculeGraph.with_local_env_strategy(mol, OpenBabelNN())
    mol_graph = metal_edge_extender(mol_graph)
    if critic_bonds:
        mg_edges = mol_graph.graph.edges()
        for bond in critic_bonds:
            bond.sort()
            if bond[0] != bond[1]:
                bond_tup = (bond[0], bond[1])
                if bond_tup not in mg_edges:
                    mol_graph.add_edge(bond_tup[0], bond_tup[1])
    return mol_graph


def _bonds_hybridization(nbo: Dict[str, Any], index: int):
    """
    Extract bonds from "hybridization_character" NBO output
    """

    bonds = set()
    warnings = set()

    if len(nbo["hybridization_character"]) > index:
        for bond_ind in nbo["hybridization_character"][index].get("type", list()):
            if nbo["hybridization_character"][index]["type"][bond_ind] != "BD":
                continue

            from_ind = (
                int(nbo["hybridization_character"][index]["atom 1 number"][bond_ind])
                - 1
            )
            to_ind = (
                int(nbo["hybridization_character"][index]["atom 2 number"][bond_ind])
                - 1
            )

            if (
                nbo["hybridization_character"][index]["atom 1 symbol"][bond_ind]
                in metals
            ):
                m_contrib = float(
                    nbo["hybridization_character"][index]["atom 1 polarization"][
                        bond_ind
                    ]
                )
            elif (
                nbo["hybridization_character"][index]["atom 2 symbol"][bond_ind]
                in metals
            ):
                m_contrib = float(
                    nbo["hybridization_character"][index]["atom 2 polarization"][
                        bond_ind
                    ]
                )
            else:
                m_contrib = None

            if m_contrib is None or m_contrib >= 30.0:
                bond_type = "covalent"
                warnings.add("Contains covalent bond with metal atom")
            else:
                bond_type = "electrostatic"

            bonds.add((from_ind, to_ind, bond_type))

    return bonds, warnings


def _bonds_peturbation(
    nbo: Dict[str, Any],
    index: int,
    poss_coord: Dict[Optional[int], List[Optional[int]]],
    energy_cutoff: float,
    metal_indices: List[int],
):
    """
    Extract bonds from "perturbation_energy" NBO output
    """

    bonds = set()
    if len(nbo["perturbation_energy"]) > index:
        for inter_ind in nbo["perturbation_energy"][index].get("donor type", list()):
            coord = False
            m_ind: Optional[int] = None
            x_ind: Optional[int] = None
            if (
                int(
                    nbo["perturbation_energy"][index]["acceptor atom 1 number"][
                        inter_ind
                    ]
                )
                - 1
                in metal_indices
            ):
                if (
                    nbo["perturbation_energy"][index]["donor type"][inter_ind] == "LP"
                    and nbo["perturbation_energy"][index]["acceptor type"][inter_ind]
                    == "LV"
                ):
                    coord = True
                    m_ind = (
                        int(
                            nbo["perturbation_energy"][index]["acceptor atom 1 number"][
                                inter_ind
                            ]
                        )
                        - 1
                    )
                    x_ind = (
                        int(
                            nbo["perturbation_energy"][index]["donor atom 1 number"][
                                inter_ind
                            ]
                        )
                        - 1
                    )
                elif (
                    nbo["perturbation_energy"][index]["donor type"][inter_ind] == "LP"
                    and nbo["perturbation_energy"][index]["acceptor type"][inter_ind]
                    == "RY*"
                ):
                    coord = True
                    m_ind = (
                        int(
                            nbo["perturbation_energy"][index]["acceptor atom 1 number"][
                                inter_ind
                            ]
                        )
                        - 1
                    )
                    x_ind = (
                        int(
                            nbo["perturbation_energy"][index]["donor atom 1 number"][
                                inter_ind
                            ]
                        )
                        - 1
                    )
            elif (
                nbo["perturbation_energy"][index]["donor atom 1 number"][inter_ind] - 1
                in metal_indices
            ):
                if (
                    nbo["perturbation_energy"][index]["donor type"][inter_ind] == "LP"
                    and nbo["perturbation_energy"][index]["acceptor type"][inter_ind]
                    == "LV"
                ):
                    coord = True
                    m_ind = (
                        int(
                            nbo["perturbation_energy"][index]["donor atom 1 number"][
                                inter_ind
                            ]
                        )
                        - 1
                    )
                    x_ind = (
                        int(
                            nbo["perturbation_energy"][index]["acceptor atom 1 number"][
                                inter_ind
                            ]
                        )
                        - 1
                    )

            if not coord:
                continue
            elif x_ind not in poss_coord[m_ind]:
                continue

            energy = float(
                nbo["perturbation_energy"][index]["perturbation energy"][inter_ind]
            )
            if energy >= energy_cutoff:
                bonds.add((x_ind, m_ind, "electrostatic"))
    return bonds


def nbo_molecule_graph(mol: Molecule, nbo: Dict[str, Any]):
    """
    Construct a molecule graph from NBO data.

    :param mol: molecule to be analyzed
    :param nbo: Output from NBO7
    :return:
    """

    mg = MoleculeGraph.with_empty_graph(mol)

    alpha_bonds, warnings = _bonds_hybridization(nbo, 1)
    beta_bonds, new_warnings = _bonds_hybridization(nbo, 3)
    warnings = warnings.union(new_warnings)

    distance_cutoff = 3.0
    energy_cutoff = 3.0
    metal_indices = [i for i, e in enumerate(mol.species) if e in metals]

    poss_coord: Dict[Optional[int], List[Optional[int]]] = dict()
    dist_mat = mol.distance_matrix
    for i in metal_indices:
        poss_coord[i] = list()
        row = dist_mat[i]
        for j, val in enumerate(row):
            if i != j and val < distance_cutoff:
                poss_coord[i].append(j)

    new_alpha_bonds = _bonds_peturbation(
        nbo, 0, poss_coord, energy_cutoff, metal_indices
    )
    alpha_bonds = alpha_bonds.union(new_alpha_bonds)

    if mol.spin_multiplicity != 1:
        new_beta_bonds = _bonds_peturbation(
            nbo, 1, poss_coord, energy_cutoff, metal_indices
        )
        beta_bonds = beta_bonds.union(new_beta_bonds)

    sorted_alpha = set([tuple(sorted([a[0], a[1]])) for a in alpha_bonds])
    sorted_beta = set([tuple(sorted([b[0], b[1]])) for b in beta_bonds])

    if sorted_alpha != sorted_beta:
        warnings.add("Difference in bonding between alpha and beta electrons")

    for bond in alpha_bonds.union(beta_bonds):
        if (bond[0], bond[1]) not in mg.graph.edges() and (
            bond[1],
            bond[0],
        ) not in mg.graph.edges():
            if bond[0] < bond[1]:
                mg.add_edge(bond[0], bond[1], edge_properties={"type": bond[2]})
            else:
                mg.add_edge(bond[1], bond[0], edge_properties={"type": bond[2]})

    mg_copy = copy.deepcopy(mg)
    mg_copy.remove_nodes(metal_indices)

    if not nx.is_connected(mg_copy.graph.to_undirected()):
        warnings.add("Metal-centered complex")

    return (mg, list(warnings))


class BondingDoc(PropertyDoc):
    """Representation of molecular bonding."""

    property_name = "bonding"

    molecule_graph: MoleculeGraph = Field(..., description="Molecule graph")

    method: str = Field(..., description="Method used to compute molecule graph")

    bond_types: Dict[str, List[float]] = Field(
        dict(),
        description="Dictionary of bond types to their length, e.g. C-O to "
        "a list of the lengths of C-O bonds in Angstrom.",
    )

    bonds: List[Tuple[int, int]] = Field(
        [],
        description="List of bonds in the form (a, b), where a and b are 0-indexed atom indices",
    )

    bonds_nometal: List[Tuple[int, int]] = Field(
        [],
        description="List of bonds in the form (a, b), where a and b are 0-indexed atom indices, "
        "with all metal ions removed",
    )

    @classmethod
    def from_task(
        cls,
        task: TaskDocument,
        molecule_id: MPID,
        preferred_methods: List,
        deprecated: bool = False,
        **kwargs,
    ):  # type: ignore[override]
        """
        Determine bonding from a task document

        Method preferences are as follows:
        - NBO7
        - OpenBabelNN + metal_edge_extender in pymatgen
        - Critic2 (really OpenBabelNN + metal_edge_extender + Critic2)

        :param task: task document from which bonding properties can be extracted
        :param molecule_id: mpid
        :param preferred_methods: list of methods; by default, NBO7, Critic2, and the combination
            of OpenBabelNN and metal_edge_extender in pymatgen, in that order
        :param kwargs: to pass to PropertyDoc
        :return:
        """

        mg_made = False
        method = None
        warnings = list()

        if task.output.optimized_molecule is not None:
            mol = task.output.optimized_molecule
        else:
            mol = task.output.initial_molecule

        for m in preferred_methods:
            if mg_made:
                break

            if m == "nbo" and task.output.nbo is not None:
                if task.orig["rem"].get("run_nbo6", False):
                    method = "nbo"
                    mg, warnings = nbo_molecule_graph(mol, task.output.nbo)
                else:
                    # Cannot make NBO molecule graph with NBO5
                    continue

            elif m == "critic2" and task.critic2 is not None:
                method = "critic2"
                critic = fix_C_Li_bonds(task.critic2)
                critic_bonds = critic["processed"]["bonds"]
                mg = make_mol_graph(mol, critic_bonds=critic_bonds)

            else:
                method = "OpenBabelNN + metal_edge_extender"
                mg = make_mol_graph(mol)

        bonds = list()
        for bond in mg.graph.edges():
            bonds.append(sorted([bond[0], bond[1]]))

        # Calculate bond lengths
        bond_types = dict()
        for u, v in mg.graph.edges():
            species_u = str(mg.molecule.species[u])
            species_v = str(mg.molecule.species[v])
            if species_u < species_v:
                species = f"{species_u}-{species_v}"
            else:
                species = f"{species_v}-{species_u}"
            dist = mg.molecule.get_distance(u, v)
            if species not in bond_types:
                bond_types[species] = [dist]
            else:
                bond_types[species].append(dist)

        m_inds = [e for e in range(len(mol)) if str(mol.species[e]) in metals]

        bonds_nometal = list()
        for bond in bonds:
            if not any([m in bond for m in m_inds]):
                bonds_nometal.append(bond)

        return super().from_molecule(
            meta_molecule=mol,
            molecule_id=molecule_id,
            method=method,
            warnings=warnings,
            molecule_graph=mg,
            bond_types=bond_types,
            bonds=bonds,
            bonds_nometal=bonds_nometal,
            deprecated=deprecated,
            origins=[PropertyOrigin(name="bonding", task_id=task.task_id)],
            **kwargs,
        )
