import numpy as np
from abc import ABC, abstractmethod
from agox.module import Module, register_modules
from agox.observer import Observer
from agox.writer import Writer, agox_writer

class CollectorBaseClass(ABC, Observer, Writer):

    """
    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.
    verbose : bool, optional
        If True, print information to the console, by default True.
    use_counter : bool, optional
        If True the writer will not print until the Observer has finished all its 
        action in a given iteration.
    prefix : str, optional
        Prefix for the writer, by default ''.
    surname : str, optional
        Surname for the writer, by default ''.
    """

    def __init__(self, generators=None, sampler=None, environment=None, order=2,
        sets={'set_key':'candidates'}, gets={}, verbose=True, use_counter=True,
        prefix='', surname=''):
        """
        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. 
        """
        Observer.__init__(self, sets=sets, gets=gets, order=order, surname=surname)
        Writer.__init__(self, verbose=verbose, use_counter=use_counter, prefix=prefix)

        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)

        # Register the generators as attributes of the collector, so that they can be tracked.
        register_modules(self, self.generators, 'generator')

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

        Returns
        -------
        list 
           List of candidate objects. 
        """
        return candidates
    
    @abstractmethod
    def get_number_of_candidates(self, iteration):
        """
        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.
        return list_of_integers

    @agox_writer
    @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