"""The base [`Optimizer`][amltk.optimization.optimizer.Optimizer] class,
defines the API we require optimizers to implement.

* [`ask()`][amltk.optimization.optimizer.Optimizer.ask] - Ask the optimizer for a
    new [`Trial`][amltk.optimization.trial.Trial] to evaluate.
* [`tell()`][amltk.optimization.optimizer.Optimizer.tell] - Tell the optimizer
    the result of the sampled config. This comes in the form of a
    [`Trial.Report`][amltk.optimization.trial.Trial.Report].

Additionally, to aid users from switching between optimizers, the
[`preferred_parser()`][amltk.optimization.optimizer.Optimizer.preferred_parser]
method should return either a `parser` function or a string that can be used
with [`node.search_space(parser=..._)`][amltk.pipeline.Node.search_space] to
extract the search space for the optimizer.
"""
from __future__ import annotations

from abc import abstractmethod
from collections.abc import Callable, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Concatenate, Generic, ParamSpec, TypeVar

from amltk.store.paths.path_bucket import PathBucket

if TYPE_CHECKING:
    from amltk.optimization.metric import Metric
    from amltk.optimization.trial import Trial
    from amltk.pipeline import Node

I = TypeVar("I")  # noqa: E741
P = ParamSpec("P")
ParserOutput = TypeVar("ParserOutput")


class Optimizer(Generic[I]):
    """An optimizer protocol.

    An optimizer is an object that can be asked for a trail using `ask` and a
    `tell` to inform the optimizer of the report from that trial.
    """

    metrics: Sequence[Metric]
    """The metrics to optimize."""

    bucket: PathBucket
    """The bucket to give to trials generated by this optimizer."""

    def __init__(
        self,
        metrics: Sequence[Metric],
        bucket: PathBucket | None = None,
    ) -> None:
        """Initialize the optimizer.

        Args:
            metrics: The metrics to optimize.
            bucket: The bucket to store results of individual trials from this
                optimizer.
        """
        super().__init__()
        self.metrics = metrics
        self.bucket = (
            bucket
            if bucket is not None
            else PathBucket(f"{self.__class__.__name__}-{datetime.now().isoformat()}")
        )

    @abstractmethod
    def tell(self, report: Trial.Report[I]) -> None:
        """Tell the optimizer the report for an asked trial.

        Args:
            report: The report for a trial
        """

    @abstractmethod
    def ask(self) -> Trial[I]:
        """Ask the optimizer for a trial to evaluate.

        Returns:
            A config to sample.
        """
        ...

    @classmethod
    def preferred_parser(
        cls,
    ) -> str | Callable[Concatenate[Node, ...], Any] | Callable[[Node], Any] | None:
        """The preferred parser for this optimizer.

        !!! note

            Subclasses should override this as required.

        """
        return None
