"""The [`SMACOptimizer`][amltk.optimization.optimizers.smac.SMACOptimizer],
is a wrapper around the [`smac`](https://github.com/automl/smac3) optimizer.

!!! tip "Requirements"

    This requires `smac` which can be installed with:

    ```bash
    pip install amltk[smac]

    # Or directly
    pip install smac
    ```

This uses `ConfigSpace` as its [`search_space()`][amltk.pipeline.Node.search_space] to
optimize.

Users should report results using
[`trial.success(cost=...)`][amltk.optimization.Trial.success]
where `cost=` is a scaler value or a list of scaler values for the objective
to **minimize**.

Visit their documentation for what you can pass to
[`SMACOptimizer.create()`][amltk.optimization.optimizers.smac.SMACOptimizer.create].

The below example shows how you can use SMAC to optimize an sklearn pipeline.

```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.smac import SMACOptimizer
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)
        loss = 1 - accuracy
        return trial.success(loss=loss, accuracy=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=SMACOptimizer.preferred_parser())
optimizer = SMACOptimizer.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())
"""  # noqa: E501
from __future__ import annotations

import logging
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Literal
from typing_extensions import override

import numpy as np
from pynisher import MemoryLimitException, TimeoutException
from smac import HyperparameterOptimizationFacade, MultiFidelityFacade, Scenario
from smac.runhistory import (
    StatusType,
    TrialInfo as SMACTrialInfo,
    TrialValue as SMACTrialValue,
)

from amltk.optimization import Optimizer, Trial
from amltk.randomness import as_int

if TYPE_CHECKING:
    from pathlib import Path
    from typing_extensions import Self

    from ConfigSpace import ConfigurationSpace
    from smac.facade import AbstractFacade

    from amltk.store import Bucket
    from amltk.types import FidT, Seed


logger = logging.getLogger(__name__)


class SMACOptimizer(Optimizer[SMACTrialInfo]):
    """An optimizer that uses SMAC to optimize a config space."""

    def __init__(
        self,
        *,
        facade: AbstractFacade,
        bucket: Bucket | None = None,
        fidelities: Mapping[str, FidT] | None = None,
    ) -> None:
        """Initialize the optimizer.

        Args:
            facade: The SMAC facade to use.
            bucket: The bucket given to trials generated by this optimizer.
            fidelities: The fidelities to use, if any.
        """
        super().__init__(bucket=bucket)
        self.facade = facade
        self.fidelities = fidelities

    @classmethod
    def create(
        cls,
        *,
        space: ConfigurationSpace,
        bucket: Bucket | None = None,
        seed: Seed | None = None,
        fidelities: Mapping[str, FidT] | None = None,
        continue_from_last_run: bool = False,
        logging_level: int | Path | Literal[False] | None = False,
    ) -> Self:
        """Create a new SMAC optimizer using either the HPO facade or
        a mutli-fidelity facade.

        Args:
            space: The config space to optimize.
            bucket: The bucket given to trials generated by this optimizer.
            seed: The seed to use for the optimizer.
            fidelities: The fidelities to use, if any.
            continue_from_last_run: Whether to continue from a previous run.
            logging_level: The logging level to use.
                This argument is passed forward to SMAC, use False to disable
                SMAC's handling of logging.
        """
        seed = as_int(seed)

        facade_cls: type[AbstractFacade]
        if fidelities:
            if len(fidelities) == 1:
                v = next(iter(fidelities.values()))
                min_budget, max_budget = v
            else:
                min_budget, max_budget = 1.0, 100.0

            scenario = Scenario(
                configspace=space,
                seed=seed,
                min_budget=min_budget,
                max_budget=max_budget,
            )
            facade_cls = MultiFidelityFacade

        else:
            scenario = Scenario(configspace=space, seed=seed)
            facade_cls = HyperparameterOptimizationFacade

        facade = facade_cls(
            scenario=scenario,
            target_function="dummy",  # NOTE: https://github.com/automl/SMAC3/issues/946
            overwrite=not continue_from_last_run,
            logging_level=logging_level,
        )
        return cls(facade=facade, fidelities=fidelities, bucket=bucket)

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

        Returns:
            The trial info for the new config.
        """
        smac_trial_info = self.facade.ask()
        config = smac_trial_info.config
        budget = smac_trial_info.budget
        instance = smac_trial_info.instance
        seed = smac_trial_info.seed

        if self.fidelities and budget:
            if len(self.fidelities) == 1:
                k, _ = next(iter(self.fidelities.items()))
                trial_fids = {k: budget}
            else:
                trial_fids = {"budget": budget}
        else:
            trial_fids = None

        config_id = self.facade.runhistory.config_ids[config]
        unique_name = f"{config_id=}_{seed=}_{budget=}_{instance=}"
        trial: Trial[SMACTrialInfo] = Trial(
            name=unique_name,
            config=dict(config),
            info=smac_trial_info,
            seed=seed,
            fidelities=trial_fids,
            bucket=self.bucket,
        )
        logger.debug(f"Asked for trial {trial.name}")
        return trial

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

        Args:
            report: The report of the trial.
        """
        assert report.trial.info is not None
        logger.debug(f"Telling report for trial {report.trial.name}")
        # If we're successful, get the cost and times and report them
        if report.status is Trial.Status.SUCCESS:
            if "cost" not in report.results:
                raise ValueError(
                    f"Report must have 'cost' if successful but got {report}."
                    " Use `trial.success(cost=...)` to set the results of the trial.",
                )

            reported_costs = report.results["cost"]
            if isinstance(reported_costs, np.number | int | float):
                reported_costs = float(reported_costs)
            elif isinstance(reported_costs, Sequence):
                reported_costs = [float(c) for c in reported_costs]
            else:
                raise ValueError(
                    f"Cost must be a float or sequence of floats, got {reported_costs}."
                    " Use `trial.success(cost=...)` to set the results of the trial.",
                )

            trial_value = SMACTrialValue(
                time=report.time.duration,
                starttime=report.time.start,
                endtime=report.time.end,
                cost=reported_costs,
                status=StatusType.SUCCESS,
                additional_info=report.results.get("additional_info", {}),
            )
            self.facade.tell(info=report.trial.info, value=trial_value, save=True)
            return

        if report.status is Trial.Status.FAIL:
            duration = report.time.duration
            start = report.time.start
            end = report.time.end
            reported_cost = report.results.get("cost", None)
            additional_info = report.results.get("additional_info", {})
        else:
            duration = 0
            start = 0
            end = 0
            reported_cost = None
            additional_info = {}

        # We got either a fail or a crash, time to deal with it
        status_types: dict[type, StatusType] = {
            MemoryLimitException: StatusType.MEMORYOUT,
            TimeoutException: StatusType.TIMEOUT,
        }
        status_type = StatusType.CRASHED

        if report.exception is not None:
            status_type = status_types.get(type(report.exception), StatusType.CRASHED)
            additional_info["exception"] = str(report.exception)
            additional_info["traceback"] = report.traceback

        # If we have no reported costs, we need to ensure that we have a
        # valid crash_cost based on the number of objectives
        crash_cost = self.facade.scenario.crash_cost
        objectives = self.facade.scenario.objectives

        cost: float | list[float]

        if reported_cost is not None:
            cost = reported_cost
        elif isinstance(crash_cost, float):
            cost = crash_cost
        elif isinstance(crash_cost, float):
            cost = [crash_cost for _ in range(len(objectives))]
        elif isinstance(crash_cost, Sequence):
            cost = list(crash_cost)
        else:
            raise ValueError(
                f"Multiple crash cost reported ({crash_cost}) for only a single"
                f" objective in `Scenario({objectives=}, ...)",
            )

        if isinstance(cost, Sequence) and (len(cost) != len(objectives)):
            raise ValueError(
                f"Length of crash cost ({len(cost)}) and objectives "
                f"({len(objectives)}) must be equal",
            )

        trial_value = SMACTrialValue(
            time=duration,
            starttime=start,
            endtime=end,
            cost=cost,
            status=status_type,
            additional_info=additional_info,
        )
        self.facade.tell(info=report.trial.info, value=trial_value, save=True)

    @override
    @classmethod
    def preferred_parser(cls) -> Literal["configspace"]:
        """The preferred parser for this optimizer."""
        return "configspace"
