"""Configuration class for gradients."""

from __future__ import annotations

from typing import TYPE_CHECKING, Literal, Self

import numpy as np
from pydantic import (
    BaseModel,
    ConfigDict,
    PositiveInt,
    field_validator,
    model_validator,
)

from ropt.config.utils import check_enum_values, immutable_array
from ropt.config.validated_types import (  # noqa: TC001
    Array1D,
    Array1DInt,
    ArrayEnum,
    ItemOrTuple,
)
from ropt.enums import BoundaryType, PerturbationType

from .constants import (
    DEFAULT_NUMBER_OF_PERTURBATIONS,
    DEFAULT_PERTURBATION_BOUNDARY_TYPE,
    DEFAULT_PERTURBATION_MAGNITUDE,
    DEFAULT_PERTURBATION_TYPE,
    DEFAULT_SEED,
)

if TYPE_CHECKING:
    from ropt.config.enopt import VariablesConfig
    from ropt.transforms import OptModelTransforms


class GradientConfig(BaseModel):
    """Configuration class for gradient calculations.

    This class, `GradientConfig`, defines the configuration for gradient
    calculations used in an [`EnOptConfig`][ropt.config.enopt.EnOptConfig]
    object.

    Gradients are estimated using function values calculated from perturbed
    variables and the unperturbed variables. The `number_of_perturbations` field
    determines the number of perturbed variables used, which must be at least
    one.

    If function evaluations for some perturbed variables fail, the gradient may
    still be estimated as long as a minimum number of evaluations succeed. The
    `perturbation_min_success` field specifies this minimum. By default, it
    equals `number_of_perturbations`.

    Perturbations are generated by sampler objects configured in the parent
    [`EnOptConfig`][ropt.config.enopt.EnOptConfig] object. The `samplers` field
    specifies, for each variable, the index of the sampler to use. A random
    number generator is created to support samplers that require random numbers.

    The generated perturbation values are scaled by the
    `perturbation_magnitudes` and can be modified based on the
    `perturbation_types`. See [`PerturbationType`][ropt.enums.PerturbationType]
    for details on available perturbation types.

    Perturbed variables may violate the defined variable bounds. The
    `boundary_types` field specifies how to handle such violations. See
    [`BoundaryType`][ropt.enums.BoundaryType] for details on available boundary
    handling methods.

    The `perturbation_types` and `boundary_types` fields use values from the
    [`PerturbationType`][ropt.enums.PerturbationType] and
    [`BoundaryType`][ropt.enums.BoundaryType] enumerations, respectively.

    Gradients are calculated for each realization individually and then combined
    into a total gradient. If `number_of_perturbations` is low, or even just
    one, individual gradient calculations may be unreliable. In this case,
    setting `merge_realizations` to `True` directs the optimizer to combine the
    results of all realizations directly into a single gradient estimate.

    The `evaluation_policy` option controls how and when objective functions and
    gradients are calculated. It accepts one of three string values:

    - `"speculative"`: Evaluate the gradient whenever the objective function
        is requested, even if the optimizer hasn't explicitly asked for the
        gradient at that point. This approach can potentially improve load
        balancing on HPC clusters by initiating gradient work earlier, though
        its effectiveness depends on whether the optimizer can utilize these
        speculatively computed gradients.
    - `"separate"`: Always launch function and gradient evaluations as
        distinct operations, even if the optimizer requests both simultaneously.
        This is particularly useful when employing realization filters (see
        [`RealizationFilterConfig`][ropt.config.enopt.RealizationFilterConfig])
        that might disable certain realizations, as it can potentially reduce
        the number of gradient evaluations needed.
    - `"auto"`: Evaluate functions and/or gradients strictly according to the
        optimizer's requests. Calculations are performed only when the
        optimization algorithm explicitly requires them.

    Note: Seed for Samplers
        The `seed` value ensures consistent results across repeated runs with
        the same configuration. To obtain unique results for each optimization
        run, modify the seed. A common approach is to use a tuple with a unique
        ID as the first element, ensuring reproducibility across nested and
        parallel plan evaluations.

    Attributes:
        number_of_perturbations:  Number of perturbations (default:
            [`DEFAULT_NUMBER_OF_PERTURBATIONS`][ropt.config.enopt.constants.DEFAULT_NUMBER_OF_PERTURBATIONS]).
        perturbation_min_success: Minimum number of successful function evaluations
            for perturbed variables (default: equal to `number_of_perturbations`).
        perturbation_magnitudes:  Magnitudes of the perturbations for each variable
            (default:
            [`DEFAULT_PERTURBATION_MAGNITUDE`][ropt.config.enopt.constants.DEFAULT_PERTURBATION_MAGNITUDE]).
        perturbation_types:       Type of perturbation for each variable (see
            [`PerturbationType`][ropt.enums.PerturbationType], default:
            [`DEFAULT_PERTURBATION_TYPE`][ropt.config.enopt.constants.DEFAULT_PERTURBATION_TYPE]).
        boundary_types:           How to handle perturbations that violate boundary
            conditions (see [`BoundaryType`][ropt.enums.BoundaryType], default:
            [`DEFAULT_PERTURBATION_BOUNDARY_TYPE`][ropt.config.enopt.constants.DEFAULT_PERTURBATION_BOUNDARY_TYPE]).
        samplers:                 Indices of the samplers to use for each variable.
        seed:                     Seed for the random number generator used by the samplers.
        merge_realizations:       Merge all realizations for the final gradient
            calculation (default: `False`).
        evaluation_policy:        How to evaluate functions and gradients.
    """

    number_of_perturbations: PositiveInt = DEFAULT_NUMBER_OF_PERTURBATIONS
    perturbation_min_success: PositiveInt | None = None
    perturbation_magnitudes: Array1D = np.array(DEFAULT_PERTURBATION_MAGNITUDE)
    perturbation_types: ArrayEnum = np.array(DEFAULT_PERTURBATION_TYPE)
    boundary_types: ArrayEnum = np.array(DEFAULT_PERTURBATION_BOUNDARY_TYPE)
    samplers: Array1DInt | None = None
    seed: ItemOrTuple[int] = (DEFAULT_SEED,)
    merge_realizations: bool = False
    evaluation_policy: Literal["speculative", "separate", "auto"] = "auto"

    model_config = ConfigDict(
        arbitrary_types_allowed=True,
        extra="forbid",
        validate_default=True,
    )

    @model_validator(mode="after")
    def _check_perturbation_min_success(self) -> Self:
        if (
            self.perturbation_min_success is None
            or self.perturbation_min_success > self.number_of_perturbations
        ):
            self.perturbation_min_success = self.number_of_perturbations
        return self

    @field_validator("perturbation_types", mode="after")
    @classmethod
    def _check_perturbation_types(cls, value: ArrayEnum) -> ArrayEnum:
        check_enum_values(value, PerturbationType)
        return value

    @field_validator("boundary_types", mode="after")
    @classmethod
    def _check_boundary_types(cls, value: ArrayEnum) -> ArrayEnum:
        check_enum_values(value, BoundaryType)
        return value

    def fix_perturbations(
        self, variables: VariablesConfig, transforms: OptModelTransforms | None
    ) -> GradientConfig:
        variable_shape = (variables.initial_values.size,)
        magnitudes = np.broadcast_to(self.perturbation_magnitudes, variable_shape)
        boundary_types = np.broadcast_to(
            immutable_array(self.boundary_types), variable_shape
        )
        types = np.broadcast_to(
            immutable_array(self.perturbation_types), variable_shape
        )

        magnitudes = np.where(
            types == PerturbationType.RELATIVE,
            (variables.upper_bounds - variables.lower_bounds) * magnitudes,
            magnitudes,
        )

        if transforms is not None and transforms.variables is not None:
            absolute = types == PerturbationType.ABSOLUTE
            transformed = transforms.variables.magnitudes_to_optimizer(magnitudes)
            magnitudes[absolute] = transformed[absolute]

        return self.model_copy(
            update={
                "perturbation_magnitudes": magnitudes,
                "boundary_types": boundary_types,
                "perturbation_types": types,
            }
        )
