from abc import abstractmethod
from frex.models import Explanation
from typing import Generator, Optional, Any
from frex.models import Candidate
from frex.pipeline_stages import PipelineStage


class CandidateGenerator(PipelineStage):
    """
    CandidateGenerator is a PipelineStage that is intended to generate new candidates to pass through a FREx Pipeline.
    This might take the form of applying some recommendation algorithm to choose new candidates
    (e.g., generate the top 100 items the user is most likely to purchase)
    or it may generate new candidates by modifying candidates that have been generated by an earlier stage.

    A new CandidateGenerator can be minimally defined by creating a new subclass of the CandidateGenerator class
    and defining the generate() function.
    However, it is likely that custom CandidateGenerator classes will need to
    also perform some access to data sources to produce the actual objects.
    """

    def __init__(self, *, generator_explanation: Explanation, **kwargs):
        """
        :param generator_explanation: The explanation to add to the Candidate when being yielded.
        """
        self.generator_explanation = generator_explanation

    @abstractmethod
    def generate(
        self,
        *,
        candidates: Optional[Generator[Candidate, None, None]] = None,
        context: Any
    ) -> Generator[Candidate, None, None]:
        """
        Apply a custom function to generate and yield new candidates to pass through a FREx Pipeline.

        :param candidates: A Generator yielding candidates. In the setup of a FREx Pipeline, this is typically another
            PipelineStage that is yielding candidates into the next stage.
        :param context: The current context being used to execute the Pipeline.
        :return: A Generator, yielding newly generated Candidate objects.
        """
        pass

    def __call__(
        self, *, candidates: Generator[Candidate, None, None] = None, context: Any
    ) -> Generator[Candidate, None, None]:
        return self.generate(candidates=candidates, context=context)
