"""[Optuna](https://optuna.org/) is an automatic hyperparameter optimization
software framework, particularly designed for machine learning.

!!! tip "Requirements"

    This requires `Optuna` which can be installed with:

    ```bash
    pip install amltk[optuna]

    # Or directly
    pip install optuna
    ```

We provide a thin wrapper called
[`OptunaOptimizer`][amltk.optimization.optimizers.optuna.OptunaOptimizer] from which
you can integrate `Optuna` into your workflow.

This uses an Optuna-like [`search_space()`][amltk.pipeline.Node.search_space] for
its optimization.

Users should report results using
[`trial.success()`][amltk.optimization.Trial.success]
with either `cost=` or `values=` depending on any optimization directions
given to the underyling optimizer created. Please see their documentation
for more.

Visit their documentation for what you can pass to
[`OptunaOptimizer.create()`][amltk.optimization.optimizers.optuna.OptunaOptimizer.create],
which is forward to [`optun.create_study()`][optuna.create_study].

```python exec="True" source="material-block" result="python"
from __future__ import annotations

import logging

from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split

from amltk.optimization.optimizers.optuna import OptunaOptimizer
from amltk.scheduling import Scheduler
from amltk.optimization import History, Trial
from amltk.pipeline import Component

logging.basicConfig(level=logging.INFO)


def target_function(trial: Trial, pipeline: Pipeline) -> Trial.Report:
    X, y = load_iris(return_X_y=True)
    X_train, X_test, y_train, y_test = train_test_split(X, y)
    clf = pipeline.configure(trial.config).build("sklearn")

    with trial.begin():
        clf.fit(X_train, y_train)
        y_pred = clf.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)

        trial.summary["accuracy"] = accuracy
        return trial.success(cost=1-accuracy)

    return trial.fail()
from amltk._doc import make_picklable; make_picklable(target_function)  # markdown-exec: hide


pipeline = Component(RandomForestClassifier, space={"n_estimators": (10, 100)})
space = pipeline.search_space(parser=OptunaOptimizer.preferred_parser())
optimizer = OptunaOptimizer.create(space=space)

N_WORKERS = 2
scheduler = Scheduler.with_processes(N_WORKERS)
task = scheduler.task(target_function)

history = History()

@scheduler.on_start(repeat=N_WORKERS)
def on_start():
    trial = optimizer.ask()
    task.submit(trial, pipeline)

@task.on_result
def tell_and_launch_trial(_, report: Trial.Report):
    if scheduler.running():
        optimizer.tell(report)
        trial = optimizer.ask()
        task.submit(trial, pipeline)


@task.on_result
def add_to_history(_, report: Trial.Report):
    history.add(report)

scheduler.run(timeout=3, wait=False)

print(history.df())
```

!!! todo "Some more documentation"

    Sorry!

"""  # noqa: E501
from __future__ import annotations

from collections.abc import Sequence
from typing import TYPE_CHECKING, Any
from typing_extensions import Self, override

import optuna
from optuna.study import Study, StudyDirection
from optuna.trial import (
    Trial as OptunaTrial,
    TrialState,
)

from amltk.optimization import Optimizer, Trial
from amltk.pipeline.parsers.optuna import parser

if TYPE_CHECKING:
    from typing import Protocol

    from amltk.pipeline import Node
    from amltk.pipeline.parsers.optuna import OptunaSearchSpace
    from amltk.store import Bucket

    class OptunaParser(Protocol):
        """A protocol for Optuna search space parser."""

        def __call__(
            self,
            node: Node,
            *,
            flat: bool = False,
            delim: str = ":",
        ) -> OptunaSearchSpace:
            """See [`optuna_parser`][amltk.pipeline.parsers.optuna.parser]."""
            ...


class OptunaOptimizer(Optimizer[OptunaTrial]):
    """An optimizer that uses Optuna to optimize a search space."""

    @override
    def __init__(
        self,
        *,
        study: Study,
        bucket: Bucket | None = None,
        space: OptunaSearchSpace,
    ) -> None:
        """Initialize the optimizer.

        Args:
            study: The Optuna Study to use.
            bucket: The bucket given to trials generated by this optimizer.
            space: Defines the current search space.
        """
        super().__init__(bucket=bucket)
        self.study = study
        self.space = space

    @classmethod
    def create(
        cls,
        *,
        space: OptunaSearchSpace,
        bucket: Bucket | None = None,
        **kwargs: Any,
    ) -> Self:
        """Create a new Optuna optimizer. For more information, check Optuna
            documentation
            [here](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.create_study.html#).

        Args:
            space: Defines the current search space.
            bucket: The bucket given to trials generated by this optimizer.
            **kwargs: Additional arguments to pass to
                [`optuna.create_study`][optuna.create_study].

        Returns:
            Self: The newly created optimizer.
        """
        study = optuna.create_study(**kwargs)
        return cls(study=study, space=space, bucket=bucket)

    @override
    def ask(self) -> Trial[OptunaTrial]:
        """Ask the optimizer for a new config.

        Returns:
            The trial info for the new config.
        """
        optuna_trial = self.study.ask(self.space)
        config = optuna_trial.params
        trial_number = optuna_trial.number
        unique_name = f"{trial_number=}"
        return Trial(
            name=unique_name,
            config=config,
            info=optuna_trial,
            bucket=self.bucket,
        )

    @override
    def tell(self, report: Trial.Report[OptunaTrial]) -> None:
        """Tell the optimizer the result of the sampled config.

        Args:
            report: The report of the trial.
        """
        trial = report.trial.info
        assert trial is not None

        if report.status is Trial.Status.SUCCESS:
            trial_state = TrialState.COMPLETE
            values = self._verify_success_report_values(report)
        else:
            trial_state = TrialState.FAIL
            values = None

        self.study.tell(trial=trial, values=values, state=trial_state)

    def _verify_success_report_values(
        self,
        report: Trial.Report[OptunaTrial],
    ) -> float | Sequence[float]:
        """Verify that the report is valid.

        Args:
            report: The report to check.

        Raises:
            ValueError: If both "cost" and "values" reported or
                if the study direction is not "minimize" and "cost" is reported.
        """
        if "cost" in report.results and "values" in report.results:
            raise ValueError(
                "Both 'cost' and 'values' were provided in the report. "
                "Only one of them should be provided.",
            )

        if "cost" not in report.results and "values" not in report.results:
            raise ValueError(
                "Neither 'cost' nor 'values' were provided in the report. "
                "At least one of them should be provided.",
            )

        directions = self.study.directions

        values = None
        if "cost" in report.results:
            if not all(direct == StudyDirection.MINIMIZE for direct in directions):
                raise ValueError(
                    "The study direction is not 'minimize',"
                    " but 'cost' was provided in the report.",
                )
            values = report.results["cost"]
        else:
            values = report.results["values"]

        if not (
            isinstance(values, float | int)
            or (
                isinstance(values, Sequence)
                and all(isinstance(value, float | int) for value in values)
            )
        ):
            raise ValueError(
                f"Reported {values=} should be float or a sequence of floats",
            )

        return values

    @override
    @classmethod
    def preferred_parser(cls) -> OptunaParser:
        return parser
