"""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

import logging
from abc import abstractmethod
from collections.abc import Callable, Iterable, Iterator, Sequence
from datetime import datetime
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    Any,
    Concatenate,
    Generic,
    ParamSpec,
    Protocol,
    TypeVar,
    overload,
)
from typing_extensions import Self

from more_itertools import all_unique

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
    from amltk.types import Seed

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

logger = logging.getLogger(__name__)


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__()
        if not all_unique(metric.name for metric in metrics):
            raise ValueError(
                "All metrics must have unique names."
                f"Got {metrics} with names {[metric.name for metric in metrics]}",
            )

        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
        """

    @overload
    @abstractmethod
    def ask(self, n: int) -> Iterable[Trial[I]]:
        ...

    @overload
    @abstractmethod
    def ask(self, n: None = None) -> Trial[I]:
        ...

    @abstractmethod
    def ask(self, n: int | None = None) -> Trial[I] | Iterable[Trial[I]]:
        """Ask the optimizer for a trial to evaluate.

        Args:
            n: The number of trials to ask for. If `None`, ask for a single trial.

        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

    @classmethod
    @abstractmethod
    def create(
        cls,
        *,
        space: Node,
        metrics: Metric | Sequence[Metric],
        bucket: str | Path | PathBucket | None = None,
        seed: Seed | None = None,
    ) -> Self:
        """Create this optimizer.

        !!! note

            Subclasses should override this with more specific configuration
            but these arguments should be all that's necessary to create the optimizer.

        Args:
            space: The space to optimize over.
            bucket: The bucket for where to store things related to the trial.
            metrics: The metrics to optimize.
            seed: The seed to use for the optimizer.

        Returns:
            The optimizer.
        """

    class CreateSignature(Protocol):
        """A Protocol which defines the keywords required to create an
        optimizer with deterministic behavior at a desired location.

        This protocol matches the `Optimizer.create` classmethod, however we also
        allow any function which accepts the keyword arguments to create an
        Optimizer.
        """

        def __call__(
            self,
            *,
            space: Node,
            metrics: Metric | Sequence[Metric],
            bucket: PathBucket | None = None,
            seed: Seed | None = None,
        ) -> Optimizer:
            """A function which creates an optimizer for node.optimize should
            accept the following keyword arguments.

            Args:
                space: The node to optimize
                metrics: The metrics to optimize
                bucket: The bucket to store the results in
                seed: The seed to use for the optimization
            """
            ...

    @classmethod
    def _get_known_importable_optimizer_classes(cls) -> Iterator[type[Optimizer]]:
        """Get all developer known optimizer classes. This is used for defaults.

        Do not rely on this functionality and prefer to give concrete optimizers to
        functionality requiring one. This is intended for convenience of particular
        quickstart methods.
        """
        # NOTE: We can't use the `Optimizer.__subclasses__` method as the optimizers
        # are not imported by any other module initially and so they do no exist
        # until imported. Hence this manual iteration. For now, we be explicit and
        # only if the optimizer list grows should we consider dynamic importing.
        try:
            from amltk.optimization.optimizers.smac import SMACOptimizer

            yield SMACOptimizer
        except ImportError as e:
            logger.debug("Failed to import SMACOptimizer", exc_info=e)

        try:
            from amltk.optimization.optimizers.optuna import OptunaOptimizer

            yield OptunaOptimizer
        except ImportError as e:
            logger.debug("Failed to import OptunaOptimizer", exc_info=e)

        try:
            from amltk.optimization.optimizers.neps import NEPSOptimizer

            yield NEPSOptimizer
        except ImportError as e:
            logger.debug("Failed to import NEPSOptimizer", exc_info=e)
