from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Generic, Sequence, TypeVar

from typing_extensions import Protocol

from .interval import Interval
from .sample import Sample

CostT = TypeVar("CostT", covariant=True)


class ObjectiveFn(Protocol[CostT]):
    """Representation of a function that can be optimized by an optimizer.

    An objective function evaluates samples generated by the optimizer and returns a scalar cost
    value for each sample. The objective function is capable of evaluating one or more samples at
    once.
    """

    def eval_sample(self, sample: Sample) -> CostT:
        """Evaluate a single sample.

        Arguments:
            sample: The sample to evaluate

        Returns:
            The cost of the sample
        """
        ...

    def eval_samples(self, samples: Sequence[Sample]) -> Sequence[CostT]:
        """Evaluate several samples sequentially.

        Arguments:
            samples: A sequence of samples to evaluate

        Returns:
            The cost of each sample in the same order they were given
        """
        ...

    def eval_samples_parallel(self, samples: Sequence[Sample], processes: int) -> Sequence[CostT]:
        """Evaluate several samples parallelizing according to the number of processes.

        Arguments:
            samples: A sequence of samples to evaluate
            processes: The number of processes to use to evaluate the samples

        Returns:
            The cost of each sample in the sample order they were given
        """
        ...


ResultT = TypeVar("ResultT")


class Optimizer(Generic[CostT, ResultT], ABC):
    """An optimizer selects samples to be evaluated by the cost function."""

    @abstractmethod
    def optimize(
        self, func: ObjectiveFn[CostT], bounds: Sequence[Interval], budget: int, seed: int
    ) -> ResultT:
        """Evaluate samples and use the result to select more samples until the budget is reached.

        The optimize method is responsible for generating samples from the
        """
        ...
