"""Configuration class for gradients."""

from __future__ import annotations

from typing import TYPE_CHECKING, Self

import numpy as np
from pydantic import ConfigDict, PositiveInt, model_validator

from ropt.config.utils import ImmutableBaseModel, 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(ImmutableBaseModel):
    """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.

    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`).
    """

    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

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

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

    @model_validator(mode="after")
    def _check(self) -> Self:
        check_enum_values(self.perturbation_types, PerturbationType)
        check_enum_values(self.boundary_types, BoundaryType)
        return self

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

        try:
            magnitudes = np.broadcast_to(magnitudes, (variable_count,))
        except ValueError as err:
            msg = (
                "the perturbation magnitudes cannot be broadcasted "
                f"to a length of {variable_count}"
            )
            raise ValueError(msg) from err

        if boundary_types.size == 1:
            boundary_types = np.broadcast_to(
                immutable_array(boundary_types),
                (variable_count,),
            )
        elif boundary_types.size == variable_count:
            boundary_types = immutable_array(boundary_types)
        else:
            msg = f"perturbation boundary_types must have {variable_count} items"
            raise ValueError(msg)

        if types.size == 1:
            types = np.broadcast_to(immutable_array(types), (variable_count,))
        elif types.size == variable_count:
            types = immutable_array(types)
        else:
            msg = f"perturbation types must have {variable_count} items"
            raise ValueError(msg)

        relative = types == PerturbationType.RELATIVE
        if not np.all(
            np.logical_and(
                np.isfinite(variables.lower_bounds[relative]),
                np.isfinite(variables.upper_bounds[relative]),
            ),
        ):
            msg = "The variable bounds must be finite to use relative perturbations"
            raise ValueError(msg)
        magnitudes = np.where(
            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,
            }
        )
