from __future__ import annotations

import datetime
import math
from typing import Any
from typing import overload
from typing import TYPE_CHECKING

from optuna import distributions
from optuna._deprecated import deprecated_func
from optuna._warnings import optuna_warn
from optuna.distributions import CategoricalDistribution
from optuna.distributions import FloatDistribution
from optuna.distributions import IntDistribution
from optuna.study._constrained_optimization import _CONSTRAINTS_KEY
from optuna.study._constrained_optimization import _get_constraints_from_system_attrs
from optuna.trial._base import BaseTrial


if TYPE_CHECKING:
    from collections.abc import Sequence

    from optuna.distributions import BaseDistribution
    from optuna.distributions import CategoricalChoiceType

_suggest_deprecated_msg = "Use suggest_float{args} instead."


class FixedTrial(BaseTrial):
    """A trial class which suggests a fixed value for each parameter.

    This object has the same methods as :class:`~optuna.trial.Trial`, and it suggests pre-defined
    parameter values. The parameter values can be determined at the construction of the
    :class:`~optuna.trial.FixedTrial` object. In contrast to :class:`~optuna.trial.Trial`,
    :class:`~optuna.trial.FixedTrial` does not depend on :class:`~optuna.study.Study`, and it is
    useful for deploying optimization results.

    Example:

        Evaluate an objective function with parameter values given by a user.

        .. testcode::

            import optuna


            def objective(trial):
                x = trial.suggest_float("x", -100, 100)
                y = trial.suggest_categorical("y", [-1, 0, 1])
                return x**2 + y


            assert objective(optuna.trial.FixedTrial({"x": 1, "y": 0})) == 1


    .. note::
        Please refer to :class:`~optuna.trial.Trial` for details of methods and properties.

    Args:
        params:
            A dictionary containing all parameters.
        number:
            A trial number. Defaults to ``0``.

    """

    def __init__(self, params: dict[str, Any], number: int = 0) -> None:
        self._params = params
        self._suggested_params: dict[str, Any] = {}
        self._distributions: dict[str, BaseDistribution] = {}
        self._user_attrs: dict[str, Any] = {}
        self._system_attrs: dict[str, Any] = {}
        self._datetime_start = datetime.datetime.now()
        self._number = number

    def suggest_float(
        self,
        name: str,
        low: float,
        high: float,
        *,
        step: float | None = None,
        log: bool = False,
    ) -> float:
        return self._suggest(name, FloatDistribution(low, high, log=log, step=step))

    @deprecated_func("3.0.0", "6.0.0", text=_suggest_deprecated_msg.format(args=""))
    def suggest_uniform(self, name: str, low: float, high: float) -> float:
        return self.suggest_float(name, low, high)

    @deprecated_func("3.0.0", "6.0.0", text=_suggest_deprecated_msg.format(args="(..., log=True)"))
    def suggest_loguniform(self, name: str, low: float, high: float) -> float:
        return self.suggest_float(name, low, high, log=True)

    @deprecated_func("3.0.0", "6.0.0", text=_suggest_deprecated_msg.format(args="(..., step=...)"))
    def suggest_discrete_uniform(self, name: str, low: float, high: float, q: float) -> float:
        return self.suggest_float(name, low, high, step=q)

    def suggest_int(
        self, name: str, low: int, high: int, *, step: int = 1, log: bool = False
    ) -> int:
        return int(self._suggest(name, IntDistribution(low, high, log=log, step=step)))

    @overload
    def suggest_categorical(self, name: str, choices: Sequence[None]) -> None: ...

    @overload
    def suggest_categorical(self, name: str, choices: Sequence[bool]) -> bool: ...

    @overload
    def suggest_categorical(self, name: str, choices: Sequence[int]) -> int: ...

    @overload
    def suggest_categorical(self, name: str, choices: Sequence[float]) -> float: ...

    @overload
    def suggest_categorical(self, name: str, choices: Sequence[str]) -> str: ...

    @overload
    def suggest_categorical(
        self, name: str, choices: Sequence[CategoricalChoiceType]
    ) -> CategoricalChoiceType: ...

    def suggest_categorical(
        self, name: str, choices: Sequence[CategoricalChoiceType]
    ) -> CategoricalChoiceType:
        return self._suggest(name, CategoricalDistribution(choices=choices))

    def report(self, value: float, step: int) -> None:
        pass

    def should_prune(self) -> bool:
        return False

    def set_user_attr(self, key: str, value: Any) -> None:
        self._user_attrs[key] = value

    def _suggest(self, name: str, distribution: BaseDistribution) -> Any:
        if name not in self._params:
            raise ValueError(
                f"The value of the parameter '{name}' is not found. Please set it at "
                "the construction of the FixedTrial object."
            )

        value = self._params[name]
        param_value_in_internal_repr = distribution.to_internal_repr(value)
        if not distribution._contains(param_value_in_internal_repr):
            optuna_warn(
                f"The value {value} of the parameter '{name}' is out of "
                f"the range of the distribution {distribution}."
            )

        if name in self._distributions:
            distributions.check_distribution_compatibility(self._distributions[name], distribution)

        self._suggested_params[name] = value
        self._distributions[name] = distribution

        return value

    @property
    def params(self) -> dict[str, Any]:
        return self._suggested_params

    @property
    def distributions(self) -> dict[str, BaseDistribution]:
        return self._distributions

    @property
    def user_attrs(self) -> dict[str, Any]:
        return self._user_attrs

    @property
    def system_attrs(self) -> dict[str, Any]:
        return self._system_attrs

    @property
    def datetime_start(self) -> datetime.datetime | None:
        return self._datetime_start

    @property
    def number(self) -> int:
        return self._number

    @property
    def constraints(self) -> dict[str, float]:
        """Returns constraint values.

        The trial is considered feasible when all constraint values are zero or less.

        Returns:
            constraint values of trial.
        """

        return _get_constraints_from_system_attrs(self.system_attrs)

    def set_constraint(self, key: str, value: float) -> None:
        """Set a constraint value for the trial.

        Args:
            key:
                A constraint name.
            value:
                A constraint value. The trial is considered feasible when all constraint values
                are zero or less.
        """

        try:
            # For convenience, we allow users to set a value that can be cast to `float`.
            value = float(value)
        except (TypeError, ValueError):
            message = (
                f"The `value` argument is of type '{type(value)}' but supposed to be a float."
            )
            raise TypeError(message) from None

        if math.isnan(value):
            raise ValueError(f"Attempted to set a constraint for {key!r}, but NaN is not allowed.")

        constraint_key = f"{_CONSTRAINTS_KEY}:{key}"

        if constraint_key in self._system_attrs:
            # Do nothing if already set.
            optuna_warn(
                f"The constraint value is ignored because this constraint `{key=}` is already set."
            )
            return

        self._system_attrs[constraint_key] = value