import logging
from typing import Dict, List, Any

import numpy as np
from pydantic import Field
from pymatgen.analysis.graphs import StructureGraph
from pymatgen.analysis.local_env import NearNeighbors

from emmet.core.material_property import PropertyDoc

AVAILABLE_METHODS = {nn.__name__: nn for nn in NearNeighbors.__subclasses__()}


class BondingDoc(PropertyDoc):
    """Structure graphs representing chemical bonds calculated from structure
    using near neighbor strategies as defined in pymatgen."""

    property_name: str = "bonding"

    structure_graph: StructureGraph = Field(
        description="Structure graph",
    )

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

    bond_types: Dict[str, List[float]] = Field(
        description="Dictionary of bond types to their length, e.g. a Fe-O to "
        "a list of the lengths of Fe-O bonds in Angstrom."
    )
    bond_length_stats: Dict[str, Any] = Field(
        description="Dictionary of statistics of bonds in structure "
        "with keys all_weights, min, max, mean and variance."
    )
    coordination_envs: List[str] = Field(
        description="List of co-ordination environments, e.g. ['Mo-S(6)', 'S-Mo(3)']."
    )
    coordination_envs_anonymous: List[str] = Field(
        description="List of co-ordination environments without elements "
        "present, e.g. ['A-B(6)', 'A-B(3)']."
    )

    @classmethod
    def from_structure(
        cls,
        structure,
        material_id,
        preferred_methods=(
            "CrystalNN",
            "MinimumDistanceNN",
        ),
        **kwargs
    ):
        """
        Calculate

        :param structure: ideally an oxidation state-decorated structure
        :param material_id: mpid
        :param preferred_methods: list of strings of NearNeighbor classes or NearNeighbor instances
        :param deprecated: whether source material is or is not deprecated
        :param kwargs: to pass to PropertyDoc
        :return:
        """

        bonding_info = None
        preferred_methods = [  # type: ignore
            AVAILABLE_METHODS[method]() if isinstance(method, str) else method
            for method in preferred_methods
        ]

        for method in preferred_methods:
            try:
                sg = StructureGraph.with_local_env_strategy(structure, method)

                # ensure edge weights are specifically bond lengths
                edge_weights = []
                for u, v, d in sg.graph.edges(data=True):
                    jimage = np.array(d["to_jimage"])
                    dist = sg.structure.get_distance(u, v, jimage=jimage)
                    edge_weights.append((u, v, d["to_jimage"], dist))
                for u, v, to_jimage, dist in edge_weights:
                    sg.alter_edge(u, v, to_jimage=to_jimage, new_weight=dist)

                bonding_info = {
                    "method": method.__class__.__name__,
                    "structure_graph": sg,
                    "bond_types": sg.types_and_weights_of_connections,
                    "bond_length_stats": sg.weight_statistics,
                    "coordination_envs": sg.types_of_coordination_environments(),
                    "coordination_envs_anonymous": sg.types_of_coordination_environments(
                        anonymous=True
                    ),
                }

                break

            except Exception as e:
                logging.warning(
                    "Failed to calculate bonding: {} {} {}".format(
                        material_id, method, e
                    )
                )

        if bonding_info:
            return super().from_structure(
                meta_structure=structure,
                material_id=material_id,
                **bonding_info,
                **kwargs
            )