import warnings
from abc import ABC, abstractmethod
from copy import deepcopy
from typing import Dict, List, Optional
from uuid import uuid4

import numpy as np
from ase.data import covalent_radii
from ase.geometry import get_distances

from agox.candidates import Candidate, StandardCandidate
from agox.environments import Environment
from agox.helpers.confinement import Confinement
from agox.main import State
from agox.observer import Observer
from agox.samplers import Sampler

dimensionality_angles = {
    3: {"theta": [0, 2 * np.pi], "phi": [0, np.pi]},
    2: {"theta": [0, 2 * np.pi], "phi": [np.pi / 2, np.pi / 2]},
    1: {"theta": [0, 0], "phi": [np.pi / 2, np.pi / 2]},
}


class GeneratorBaseClass(ABC, Observer):
    def __init__(
        self,
        confinement: Optional[Confinement] = None,
        confinement_cell: Optional[np.ndarray] = None,
        confinement_corner: Optional[np.ndarray] = None,    
        c1: float = 0.75,
        c2: float = 1.25,
        use_mic: bool = True,
        environment: Sampler = None,
        sampler: Environment = None,
        gets: dict[str, str] | None = None,
        sets: dict[str, str] | None = None,
        order: int = 2,
        replace: bool = True,
        surname: Optional[str] = None,
        **kwargs,
    ) -> None:
        """ 
        Base class for all generators.

        Parameters:
        -----------
        confinement: Confinement object
            Confinement object that the generator is confined to.
        confinement_cell: np.ndarray
            The cell that the generator is confined to. Only used if confinement is None.
        confinement_corner: np.ndarray
            Lower left corner of the confinement cell. Only used if confinement is None.
        c1: float
            Lower limit on bond lengths.
        c2: float
            Upper limit on bond lengths.
        use_mic: bool
            If the minimum image convention should be used.
        environment: Environment object
            The environment object used in the search.
        sampler: Sampler object
            Sampler object used to create and get sample members from. Only 
            relevant if the generator is used without a collector.
        gets: Dict
            Dictionary of get methods.
        sets: Dict
            Dictionary of set methods.
        order: int
            Order of the generator.
        replace: bool
            If the sampler should replace the parents.
        """
        if sets is None:
            sets = {"set_key": "candidates"}
        Observer.__init__(self, sets=sets, gets=gets, order=order, surname=surname, **kwargs)

        self.c1 = c1  # Lower limit on bond lengths.
        self.c2 = c2  # Upper limit on bond lengths.
        self.use_mic = use_mic
        self.environment = environment
        self.sampler = sampler
        self.replace = replace
        self.candidate_instanstiator = StandardCandidate

        if confinement is not None:
            self.confinement = confinement
        else:
            self.confinement = Confinement(cell=confinement_cell, corner=confinement_corner)

        if self.environment is not None:
            self.plot_confinement(environment)

        self.add_observer_method(
            self.generate, sets=self.sets[0], gets=self.gets[0], order=self.order[0], handler_identifier="AGOX"
        )

    @abstractmethod
    def _get_candidates(
        self, candidate: Candidate, parents: List[Candidate], environment: Environment
    ) -> List:  # pragma: no cover
        """
        Method that returns a list of generated candidates
        """
        return []

    def get_candidates(self, sampler: Sampler, environment: Environment) -> List[Candidate]:
        """
        Method used to get new candidates.
        A generator method must have 'get_number_of_parents' method implemented.

        Parameters:
        -----------
        sampler: Sampler object
            Sampler object used to create and get sample members from.
        environment: Environment object
            The environment object used in the search.

        Returns: list
            List of candidates generated by the generator.
        """
        number_of_parents = self.get_number_of_parents(sampler)

        # If no parents are needed, the search template is given to generator.
        if number_of_parents == 0:
            parents = []
            candidate = environment.get_template().copy()
        # Check that the returned number of parents equal the required number.
        else:
            parents = sampler.get_parents(number_of_parents, replace=self.replace)
            if len(parents) != number_of_parents:
                return []
            candidate = deepcopy(parents[0])  # Use the first parent as a template for the candidate.
            candidate.info = dict() # Info should not be used but this is a safe guard against propagating info from parents if it has been set.
            candidate.reset_meta_information()
            candidate.set_constraint()

        # Generates new candidates and add relevant meta information.
        candidates = self._get_candidates(candidate, parents, environment)
        for candidate in candidates:
            candidate.add_meta_information("generator", self.name)
            candidate.add_meta_information("uuid", str(uuid4()))
            candidate.add_meta_information(
                "parent_search_iterations", [p.get_meta_information("search_iteration") for p in parents]
            )
            candidate.add_meta_information("parent_uuids", [p.get_meta_information("uuid") for p in parents])

            if parents:
                spawn_uuids = []
                for parent in parents:
                    try:
                        spawn_uuids += parent.get_meta_information("spawn_uuids")
                    except TypeError:
                        spawn_uuid = parent.get_meta_information("spawn_uuids")
                        raise ValueError(
                            f"Parent candidate does not have 'spawn_uuids' meta information correctly set. {spawn_uuid}"
                        )
            else:
                spawn_uuids = [candidate.get_meta_information("uuid")]

            candidate.add_meta_information("spawn_uuids", spawn_uuids)

        return candidates

    @property
    @abstractmethod
    def name(self) -> str:  # pragma: no cover
        pass

    @abstractmethod
    def get_number_of_parents(self, sampler: Sampler) -> int:  # pragma: no cover
        """
        Method that returns the number of parents needed to generate a candidate.

        Parameters:
        -----------
        sampler: Sampler object
            Sampler object used to create and get sample members from.

        Returns: Int
            Number of parents needed to generate a candidate.
        """
        pass

    def __call__(self, sampler: Sampler, environment: Environment) -> List[Candidate]:
        return self.get_candidates(sampler, environment)

    ####################################################################################################################
    # Convenience methods:
    ####################################################################################################################

    def check_new_position(
        self,
        candidate: Candidate,
        new_position: np.ndarray,
        number: int,
        skipped_indices: Optional[List] = None,
    ) -> bool:
        """ """

        if skipped_indices is None:
            skipped_indices = []

        state = True
        succesful = False
        for i in range(len(candidate)):
            if i in skipped_indices:
                continue

            covalent_dist_ij = covalent_radii[candidate[i].number] + covalent_radii[number]
            rmin = self.c1 * covalent_dist_ij
            rmax = self.c2 * covalent_dist_ij

            if self.use_mic:
                vector, distance = get_distances(
                    new_position, candidate.positions[i], cell=candidate.cell, pbc=candidate.pbc
                )
            else:
                distance = np.linalg.norm(candidate.positions[i] - new_position)
            if distance < rmin:  # If a position is too close we should just give up.
                return False
            elif not distance > rmax:  # If at least a single position is not too far we have a bond.
                succesful = True
        return succesful * state

    def get_sphere_vector(self, atomic_number_i: int, atomic_number_j: int) -> np.ndarray:
        """
        Get a random vector on the sphere of appropriate radii.

        Behavior changes based on self.dimensionality:
        3: Vector on sphere.
        2: Vector on circle (in xy)
        1: Vector on line (x)
        """
        dim = self.confinement.dimensionality        
        covalent_dist_ij = covalent_radii[atomic_number_i] + covalent_radii[atomic_number_j]
        rmin = self.c1 * covalent_dist_ij
        rmax = self.c2 * covalent_dist_ij
        r = np.random.uniform(rmin**dim, rmax**dim) ** (1 / dim)
        return self.get_displacement_vector(r)

    def get_displacement_vector(self, radius: float) -> np.ndarray:
        dim = self.confinement.dimensionality        
        theta = np.random.uniform(*dimensionality_angles[dim]["theta"])
        phi = np.random.uniform(*dimensionality_angles[dim]["phi"])
        displacement = radius * np.array([np.cos(theta) * np.sin(phi), np.sin(theta) * np.sin(phi), np.cos(phi)])
        return displacement

    def get_box_vector(self) -> np.ndarray:
        return self.confinement._get_box_vector()  # From confinement-class.
    
    def check_confinement(self, position: np.ndarray) -> bool:
        return self.confinement.check_confinement(position)
    
    def get_dimensionality(self) -> int:
        return self.confinement.dimensionality

    ####################################################################################################################
    # Observer functions
    ####################################################################################################################

    def start_candidate(self) -> List[StandardCandidate]:
        """
        This method generates a candidate using the start generator, which allows other generators
        to kick-start the sampler.
        """
        from agox.generators import RandomGenerator

        return RandomGenerator(
            confinement=self.confinement.copy(),
            c1=self.c1,
            c2=self.c2,
            use_mic=self.use_mic,
        )(self.sampler, self.environment)

    @Observer.observer_method
    def generate(self, state: State) -> None:
        candidates = self.get_candidates(self.sampler, self.environment)
        if len(candidates) == 0 and self.sampler is not None and not self.sampler.initialized:
            candidates = self.start_candidate()
            self.writer("Fall-back to start generator, generated {} candidate ".format(len(candidates)))
        state.add_to_cache(self, self.set_key, candidates, mode="a")

    def plot_confinement(self, environment: Environment) -> None:
        if not self.confinement.confined:
            return

        import matplotlib.pyplot as plt

        from agox.utils.matplotlib_utils import use_agox_mpl_backend
        from agox.utils.plot import plot_atoms, plot_cell

        use_agox_mpl_backend()

        atoms = environment.get_template()
        atoms.set_constraint(environment.get_constraints())

        fig, axs = plt.subplots(1, 3, figsize=(15, 5))

        for ax, plane in zip(axs, ["xy", "xz", "yz"]):
            plot_atoms(ax, atoms, plane=plane, plot_constraint=True)
            plot_cell(ax, atoms.cell, plane=plane, collection_kwargs=dict(edgecolors="black", linestyles="dotted"))
            plot_cell(
                ax,
                self.confinement.cell,
                plane=plane,
                offset=self.confinement.corner,
                collection_kwargs=dict(edgecolors="red", linestyles="dashed"),
            )

        fig.savefig(f"confinement_plot_{self.name}.png")
        plt.close()
