from abc import ABC, abstractmethod  # noqa: N999

from agox.environments import Environment
from agox.generators import Generator
from agox.observer import Observer
from agox.samplers import Sampler


class CollectorBaseClass(ABC, Observer):
    """
    Base class for all collectors.

    Collectors produce several candidates using a set of generators.

    Parameters
    ----------
    generators : list
        List of AGOX generator instances.
    sampler : Sampler instance
        An instance of an AGOX sampler
    environment : Environment object.
        An instance of an AGOX environment object.
    order : int, optional
        Order of the collector, by default 2
    sets : dict, optional
        Dictionary of set keys, e.g. {'set_key':'candidates'}. Used to select
        in which entry in the agox.main.State cache the collector should store the candidates.
    gets : dict, optional
        Dictionary of get keys, e.g. {'get_key':'candidates'}. Used to select
        from which entry in the agox.main.State cache the collector should get candidates.
    prefix : str, optional
        Prefix for the writer, by default ''.
    surname : str, optional
        Surname for the writer, by default ''.
    """

    def __init__(
        self,
        generators: list[Generator] = None,
        sampler: Sampler = None,
        environment: Environment = None,
        order: int | float = 2,
        sets: dict[str, str] | None = None,
        gets: dict[str, str] | None = None,
        surname: str | None = None,
        **kwargs,
    ) -> None:
        """
        Quite simple, except for lines 50-54 that remove generators as observers,
        as they are now part of a Collector. This may now yield wanted behaviour
        for generators that need observers_methods that are not just for
        generation, if that becomes the case an update to those lines are
        probably necesary.

        Parameters
        ----------
        generators : list
            List of AGOX generator instances.
        sampler : Sampler instance
            An instance of an AGOX sampler
        environment : Environment object.
            An instance of an AGOX environment object.
        """
        if sets is None:
            sets = {"set_key": "candidates"}

        Observer.__init__(self, sets=sets, gets=gets, order=order, surname=surname, **kwargs)

        assert generators is not None
        assert environment is not None

        self.generators = generators
        self.sampler = sampler
        self.environment = environment
        self.candidates = []
        self.plot_confinement()

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

        for generator in self.generators:
            observer_methods = [observer_method for observer_method in generator.observer_methods.values()]
            for observer_method in observer_methods:
                if observer_method.method_name == "generate":
                    generator.remove_observer_method(observer_method)

    @abstractmethod
    def make_candidates(self):  # pragma: no cover
        """
        Method that calls generators to produce a collection of candidates.

        Returns
        -------
        list
           List of candidate objects.
        """
        ...

    @abstractmethod
    def get_number_of_candidates(self, iteration):  # pragma: no cover
        """
        Method  that computes the number of candidates for the given iteration.

        Parameters
        ----------
        iteration : int
            Iteration count.

        Returns
        -------
        list
            List of integers corresponding to the self.generators list dictating
            how many candidates are generated (or really how many times each
            generator is called.)
        """
        # Number of candidates to generate.
        ...

    @Observer.observer_method
    def generate_candidates(self, state):
        # Make the candidates - this is the method that differs between versions of the class.
        if self.do_check():
            candidates = self.make_candidates()

        iteration = self.get_iteration_counter()
        for candidate in candidates:
            candidate.add_meta_information("search_iteration", iteration)

        # Add to the iteration_cache:
        state.add_to_cache(self, self.set_key, candidates, "a")
        self.writer("Number of candidates this iteration: {}".format(len(candidates)))

    ########################################################################################################################
    # Other methods
    ########################################################################################################################

    def attach(self, main):
        super().attach(main)
        for generator in self.generators:
            generator.attach(main)

    def plot_confinement(self):
        for generator in self.generators:
            generator.plot_confinement(self.environment)

    def get_fallback_generator(self):
        """
        When sampler is empty but initialized, use a RandomGenerator
        to make new candidates.
        """
        from agox.generators import RandomGenerator

        fallback_generator = RandomGenerator(**self.environment.get_confinement())
        return fallback_generator
