from __future__ import annotations

import copy
from hashlib import blake2b
from typing import TYPE_CHECKING, Any

import networkx as nx
from pydantic import Field
from pymatgen.analysis.graphs import MoleculeGraph
from pymatgen.core.structure import Molecule

from emmet.core.molecules import MolPropertyOrigin
from emmet.core.molecules.molecule_property import PropertyDoc
from emmet.core.mpid import MPculeID
from emmet.core.qchem.task import TaskDocument
from emmet.core.types.pymatgen_types.structure_graph_adapter import MoleculeGraphType
from emmet.core.utils import make_mol_graph

if TYPE_CHECKING:
    from typing import Any


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


metals = [
    "Li",
    "Be",
    "Na",
    "Mg",
    "Al",
    "K",
    "Ca",
    "Sc",
    "Ti",
    "V",
    "Cr",
    "Mn",
    "Fe",
    "Co",
    "Ni",
    "Cu",
    "Zn",
    "Ga",
    "Rb",
    "Sr",
    "Y",
    "Zr",
    "Nb",
    "Mo",
    "Tc",
    "Ru",
    "Rh",
    "Pd",
    "Ag",
    "Cd",
    "In",
    "Sn",
    "Cs",
    "Ba",
    "Hf",
    "Ta",
    "W",
    "Re",
    "Os",
    "Ir",
    "Pt",
    "Au",
    "Hg",
    "Tl",
    "Pb",
    "Bi",
]

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 _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:
                bond_type = "covalent"
            elif 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[int | None, list[int | None]],
    energy_cutoff: float,
    metal_indices: list[int],
):
    """
    Extract bonds from "perturbation_energy" NBO output
    """

    bonds = set()  # type: ignore

    # No metals, so don't need to use perturbation analysis to get bonds
    if len(metal_indices) == 0:
        return bonds

    if len(nbo["perturbation_energy"]) > index:
        for inter_ind in nbo["perturbation_energy"][index].get("donor type", list()):
            coord = False
            m_ind: int | None = None
            x_ind: int | None = 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.from_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 str(e) in metals]

    poss_coord: dict[int | None, list[int | None]] = 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)

    try:
        if not nx.is_connected(mg_copy.graph.to_undirected()):
            warnings.add("Metal-centered complex")
    except nx.exception.NetworkXPointlessConcept:
        if len(mg.molecule) == 1:
            warnings.add("Single-atom; no bonds")

    return (mg, list(warnings))


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

    property_name: str = "bonding"

    molecule_graph: MoleculeGraphType = 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: MPculeID,
        preferred_methods: list[str],
        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: MPculeID
        :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
                and (
                    task.orig["rem"].get("run_nbo6", False)
                    or task.orig["rem"].get("nbo_external", False)
                )
            ):
                method = "nbo"
                mg, warnings = nbo_molecule_graph(mol, task.output.nbo)
                mg_made = True

            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)
                mg_made = True

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

        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)

        id_string = f"bonding-{molecule_id}-{task.task_id}-{task.lot_solvent}-{method}"
        h = blake2b()
        h.update(id_string.encode("utf-8"))
        property_id = h.hexdigest()

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