from abc import ABC
from dataclasses import field
from typing import ClassVar

import torch

try:
    import bitsandbytes as bnb
except Exception:
    bnb = None
import ludwig.schema.utils as schema_utils
from ludwig.api_annotations import DeveloperAPI
from ludwig.error import ConfigValidationError
from ludwig.schema.metadata import OPTIMIZER_METADATA
from ludwig.schema.metadata.parameter_metadata import convert_metadata_to_json, ParameterMetadata
from ludwig.schema.utils import ludwig_dataclass
from ludwig.utils.registry import Registry

optimizer_registry = Registry()


@DeveloperAPI
def register_optimizer(name: str):
    def wrap(optimizer_config: BaseOptimizerConfig):
        optimizer_registry[name] = (optimizer_config.optimizer_class, optimizer_config)
        return optimizer_config

    return wrap


@DeveloperAPI
def get_optimizer_cls(name: str):
    """Get the optimizer schema class from the optimizer schema class registry."""
    return optimizer_registry[name][1]


@DeveloperAPI
@ludwig_dataclass
class BaseOptimizerConfig(schema_utils.BaseMarshmallowConfig, ABC):
    """Base class for optimizers. Not meant to be used directly.

    The dataclass format prevents arbitrary properties from being set. Consequently, in child classes, all properties
    from the corresponding `torch.optim.Optimizer` class are copied over: check each class to check which attributes are
    different from the torch-specified defaults.
    """

    optimizer_class: ClassVar[torch.optim.Optimizer | None] = None
    "Class variable pointing to the corresponding `torch.optim.Optimizer` class."

    type: str
    """Name corresponding to an optimizer `ludwig.modules.optimization_modules.optimizer_registry`.

    Technically mutable, but attempting to load a derived optimizer with `type` set to a mismatched value will result in
    a `ValidationError`.
    """

    @property
    def is_paged(self) -> bool:
        """Returns True if the optimizer is a Paged optimizer."""
        return False

    @property
    def is_8bit(self) -> bool:
        """Returns True if the optimizer is an 8-bit optimizer."""
        return False


@DeveloperAPI
@register_optimizer(name="sgd")
@ludwig_dataclass
class SGDOptimizerConfig(BaseOptimizerConfig):
    """Parameters for stochastic gradient descent."""

    optimizer_class: ClassVar[torch.optim.Optimizer] = torch.optim.SGD
    """Points to `torch.optim.SGD`."""

    type: str = schema_utils.ProtectedString("sgd")
    """Must be 'sgd' - corresponds to name in `ludwig.modules.optimization_modules.optimizer_registry` (default:
       'sgd')"""

    # Defaults taken from https://pytorch.org/docs/stable/generated/torch.optim.SGD.html#torch.optim.SGD :
    momentum: float = schema_utils.NonNegativeFloat(
        default=0.0,
        description="Momentum factor.",
        parameter_metadata=OPTIMIZER_METADATA["momentum"],
    )

    weight_decay: float = schema_utils.NonNegativeFloat(
        default=0.0,
        description="Weight decay ($L2$ penalty).",
        parameter_metadata=OPTIMIZER_METADATA["weight_decay"],
    )

    dampening: float = schema_utils.NonNegativeFloat(
        default=0.0,
        description="Dampening for momentum.",
        parameter_metadata=OPTIMIZER_METADATA["dampening"],
    )

    nesterov: bool = schema_utils.Boolean(
        default=False,
        description="Enables Nesterov momentum.",
        parameter_metadata=OPTIMIZER_METADATA["nesterov"],
    )


if bnb is not None:

    @DeveloperAPI
    @register_optimizer(name="sgd_8bit")
    @ludwig_dataclass
    class SGD8BitOptimizerConfig(SGDOptimizerConfig):
        """Parameters for stochastic gradient descent."""

        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.SGD8bit

        type: str = schema_utils.ProtectedString("sgd_8bit")

        block_wise: bool = schema_utils.Boolean(
            default=False,
            description="Whether to use block wise update.",
        )

        percentile_clipping: int = schema_utils.IntegerRange(
            default=100,
            min=0,
            max=100,
            description="Percentile clipping.",
        )

        @property
        def is_8bit(self) -> bool:
            return True


@DeveloperAPI
@register_optimizer(name="lbfgs")
@ludwig_dataclass
class LBFGSOptimizerConfig(BaseOptimizerConfig):
    """Parameters for stochastic gradient descent."""

    optimizer_class: ClassVar[torch.optim.Optimizer] = torch.optim.LBFGS
    """Points to `torch.optim.LBFGS`."""

    type: str = schema_utils.ProtectedString("lbfgs")
    """Must be 'lbfgs' - corresponds to name in `ludwig.modules.optimization_modules.optimizer_registry` (default:
       'lbfgs')"""

    # Defaults taken from https://pytorch.org/docs/stable/generated/torch.optim.LBFGS.html#torch.optim.LBFGS
    max_iter: int = schema_utils.Integer(
        default=20,
        description="Maximum number of iterations per optimization step.",
        parameter_metadata=OPTIMIZER_METADATA["max_iter"],
    )

    max_eval: int = schema_utils.Integer(
        default=None,
        allow_none=True,
        description="Maximum number of function evaluations per optimization step. Default: `max_iter` * 1.25.",
        parameter_metadata=OPTIMIZER_METADATA["max_eval"],
    )

    tolerance_grad: float = schema_utils.NonNegativeFloat(
        default=1e-07,
        description="Termination tolerance on first order optimality.",
        parameter_metadata=OPTIMIZER_METADATA["tolerance_grad"],
    )

    tolerance_change: float = schema_utils.NonNegativeFloat(
        default=1e-09,
        description="Termination tolerance on function value/parameter changes.",
        parameter_metadata=OPTIMIZER_METADATA["tolerance_change"],
    )

    history_size: int = schema_utils.Integer(
        default=100, description="Update history size.", parameter_metadata=OPTIMIZER_METADATA["history_size"]
    )

    line_search_fn: str = schema_utils.StringOptions(
        ["strong_wolfe"],
        default=None,
        allow_none=True,
        description="Line search function to use.",
        parameter_metadata=OPTIMIZER_METADATA["line_search_fn"],
    )


@DeveloperAPI
@register_optimizer(name="adam")
@ludwig_dataclass
class AdamOptimizerConfig(BaseOptimizerConfig):
    """Parameters for adam optimization."""

    optimizer_class: ClassVar[torch.optim.Optimizer] = torch.optim.Adam
    """Points to `torch.optim.Adam`."""

    type: str = schema_utils.ProtectedString("adam")
    """Must be 'adam' - corresponds to name in `ludwig.modules.optimization_modules.optimizer_registry`
       (default: 'adam')"""

    # Defaults taken from https://pytorch.org/docs/stable/generated/torch.optim.Adam.html#torch.optim.Adam :
    betas: tuple[float, float] = schema_utils.FloatRangeTupleDataclassField(
        default=(0.9, 0.999),
        description="Coefficients used for computing running averages of gradient and its square.",
        parameter_metadata=OPTIMIZER_METADATA["betas"],
    )

    eps: float = schema_utils.NonNegativeFloat(
        default=1e-08,
        description="Term added to the denominator to improve numerical stability.",
        parameter_metadata=OPTIMIZER_METADATA["eps"],
    )

    weight_decay: float = schema_utils.NonNegativeFloat(
        default=0.0, description="Weight decay (L2 penalty).", parameter_metadata=OPTIMIZER_METADATA["weight_decay"]
    )

    amsgrad: bool = schema_utils.Boolean(
        default=False,
        description="Whether to use the AMSGrad variant of this algorithm from the paper 'On the Convergence of Adam "
        "and Beyond'.",
        parameter_metadata=OPTIMIZER_METADATA["amsgrad"],
    )


if bnb is not None:

    @DeveloperAPI
    @register_optimizer(name="adam_8bit")
    @ludwig_dataclass
    class Adam8BitOptimizerConfig(AdamOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.Adam8bit

        type: str = schema_utils.ProtectedString("adam_8bit")

        block_wise: bool = schema_utils.Boolean(
            default=True,
            description="Whether to use block wise update.",
        )

        percentile_clipping: int = schema_utils.IntegerRange(
            default=100,
            min=0,
            max=100,
            description="Percentile clipping.",
        )

        @property
        def is_8bit(self) -> bool:
            return True

    @DeveloperAPI
    @register_optimizer(name="paged_adam")
    @ludwig_dataclass
    class PagedAdamOptimizerConfig(Adam8BitOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.PagedAdam

        type: str = schema_utils.ProtectedString("paged_adam")

        @property
        def is_paged(self) -> bool:
            return True

        @property
        def is_8bit(self) -> bool:
            return False

    @DeveloperAPI
    @register_optimizer(name="paged_adam_8bit")
    @ludwig_dataclass
    class PagedAdam8BitOptimizerConfig(PagedAdamOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.PagedAdam8bit

        type: str = schema_utils.ProtectedString("paged_adam_8bit")

        @property
        def is_8bit(self) -> bool:
            return True


@DeveloperAPI
@register_optimizer(name="adamw")
@ludwig_dataclass
class AdamWOptimizerConfig(BaseOptimizerConfig):
    """Parameters for adamw optimization."""

    optimizer_class: ClassVar[torch.optim.Optimizer] = torch.optim.AdamW
    """Points to `torch.optim.AdamW`."""

    type: str = schema_utils.ProtectedString("adamw")
    """Must be 'adamw' - corresponds to name in `ludwig.modules.optimization_modules.optimizer_registry`
       (default: 'adamw')"""

    # Defaults taken from https://pytorch.org/docs/stable/generated/torch.optim.Adam.html#torch.optim.Adam :
    betas: tuple[float, float] = schema_utils.FloatRangeTupleDataclassField(
        default=(0.9, 0.999),
        description="Coefficients used for computing running averages of gradient and its square.",
        parameter_metadata=OPTIMIZER_METADATA["betas"],
    )

    eps: float = schema_utils.NonNegativeFloat(
        default=1e-08,
        description="Term added to the denominator to improve numerical stability.",
        parameter_metadata=OPTIMIZER_METADATA["eps"],
    )

    weight_decay: float = schema_utils.NonNegativeFloat(
        default=0.0, description="Weight decay ($L2$ penalty).", parameter_metadata=OPTIMIZER_METADATA["weight_decay"]
    )

    amsgrad: bool = schema_utils.Boolean(
        default=False,
        description="Whether to use the AMSGrad variant of this algorithm from the paper 'On the Convergence of Adam "
        "and Beyond'. ",
        parameter_metadata=OPTIMIZER_METADATA["amsgrad"],
    )


if bnb is not None:

    @DeveloperAPI
    @register_optimizer(name="adamw_8bit")
    @ludwig_dataclass
    class AdamW8BitOptimizerConfig(AdamWOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.AdamW8bit

        type: str = schema_utils.ProtectedString("adamw_8bit")

        block_wise: bool = schema_utils.Boolean(
            default=True,
            description="Whether to use block wise update.",
        )

        percentile_clipping: int = schema_utils.IntegerRange(
            default=100,
            min=0,
            max=100,
            description="Percentile clipping.",
        )

        @property
        def is_8bit(self) -> bool:
            return True

    @DeveloperAPI
    @register_optimizer(name="paged_adamw")
    @ludwig_dataclass
    class PagedAdamWOptimizerConfig(AdamW8BitOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.PagedAdamW

        type: str = schema_utils.ProtectedString("paged_adamw")

        @property
        def is_paged(self) -> bool:
            return True

        @property
        def is_8bit(self) -> bool:
            return False

    @DeveloperAPI
    @register_optimizer(name="paged_adamw_8bit")
    @ludwig_dataclass
    class PagedAdamW8BitOptimizerConfig(PagedAdamWOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.PagedAdamW8bit

        type: str = schema_utils.ProtectedString("paged_adamw_8bit")

        @property
        def is_8bit(self) -> bool:
            return True


@DeveloperAPI
@register_optimizer(name="adadelta")
@ludwig_dataclass
class AdadeltaOptimizerConfig(BaseOptimizerConfig):
    """Parameters for adadelta optimization."""

    optimizer_class: ClassVar[torch.optim.Optimizer] = torch.optim.Adadelta
    """Points to `torch.optim.Adadelta`."""

    type: str = schema_utils.ProtectedString("adadelta")
    """Must be 'adadelta' - corresponds to name in `ludwig.modules.optimization_modules.optimizer_registry`
       (default: 'adadelta')"""

    # Defaults taken from https://pytorch.org/docs/stable/generated/torch.optim.Adadelta.html#torch.optim.Adadelta :
    rho: float = schema_utils.FloatRange(
        default=0.9,
        min=0,
        max=1,
        description="Coefficient used for computing a running average of squared gradients.",
        parameter_metadata=OPTIMIZER_METADATA["rho"],
    )

    eps: float = schema_utils.NonNegativeFloat(
        default=1e-06,
        description="Term added to the denominator to improve numerical stability.",
        parameter_metadata=OPTIMIZER_METADATA["eps"],
    )

    weight_decay: float = schema_utils.NonNegativeFloat(
        default=0.0, description="Weight decay ($L2$ penalty).", parameter_metadata=OPTIMIZER_METADATA["weight_decay"]
    )


@DeveloperAPI
@register_optimizer(name="adagrad")
@ludwig_dataclass
class AdagradOptimizerConfig(BaseOptimizerConfig):
    """Parameters for adagrad optimization."""

    # Example docstring
    optimizer_class: ClassVar[torch.optim.Optimizer] = torch.optim.Adagrad
    """Points to `torch.optim.Adagrad`."""

    type: str = schema_utils.ProtectedString("adagrad")
    """Must be 'adagrad' - corresponds to name in `ludwig.modules.optimization_modules.optimizer_registry`
       (default: 'adagrad')"""

    # Defaults taken from https://pytorch.org/docs/stable/generated/torch.optim.Adagrad.html#torch.optim.Adagrad :
    initial_accumulator_value: float = schema_utils.NonNegativeFloat(
        default=0, description="", parameter_metadata=OPTIMIZER_METADATA["initial_accumulator_value"]
    )

    lr_decay: float = schema_utils.FloatRange(
        default=0, description="Learning rate decay.", parameter_metadata=OPTIMIZER_METADATA["lr_decay"]
    )

    weight_decay: float = schema_utils.FloatRange(
        default=0, description="Weight decay ($L2$ penalty).", parameter_metadata=OPTIMIZER_METADATA["weight_decay"]
    )

    eps: float = schema_utils.FloatRange(
        default=1e-10,
        description="Term added to the denominator to improve numerical stability.",
        parameter_metadata=OPTIMIZER_METADATA["eps"],
    )


if bnb is not None:

    @DeveloperAPI
    @register_optimizer(name="adagrad_8bit")
    @ludwig_dataclass
    class Adagrad8BitOptimizerConfig(AdagradOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.Adagrad8bit

        type: str = schema_utils.ProtectedString("adagrad_8bit")

        block_wise: bool = schema_utils.Boolean(
            default=True,
            description="Whether to use block wise update.",
        )

        percentile_clipping: int = schema_utils.IntegerRange(
            default=100,
            min=0,
            max=100,
            description="Percentile clipping.",
        )

        @property
        def is_8bit(self) -> bool:
            return True


@DeveloperAPI
@register_optimizer(name="adamax")
@ludwig_dataclass
class AdamaxOptimizerConfig(BaseOptimizerConfig):
    """Parameters for adamax optimization."""

    optimizer_class: ClassVar[torch.optim.Optimizer] = torch.optim.Adamax
    """Points to `torch.optim.Adamax`."""

    type: str = schema_utils.ProtectedString("adamax")
    """Must be 'adamax' - corresponds to name in `ludwig.modules.optimization_modules.optimizer_registry`
       (default: 'adamax')"""

    # Defaults taken from https://pytorch.org/docs/stable/generated/torch.optim.Adamax.html#torch.optim.Adamax :
    betas: tuple[float, float] = schema_utils.FloatRangeTupleDataclassField(
        default=(0.9, 0.999),
        description="Coefficients used for computing running averages of gradient and its square.",
        parameter_metadata=OPTIMIZER_METADATA["betas"],
    )

    eps: float = schema_utils.NonNegativeFloat(
        default=1e-08,
        description="Term added to the denominator to improve numerical stability.",
        parameter_metadata=OPTIMIZER_METADATA["eps"],
    )

    weight_decay: float = schema_utils.NonNegativeFloat(
        default=0.0, description="Weight decay ($L2$ penalty).", parameter_metadata=OPTIMIZER_METADATA["weight_decay"]
    )


# NOTE: keep ftrl and nadam optimizers out of registry:
# @register_optimizer(name="ftrl")
@DeveloperAPI
@ludwig_dataclass
class FtrlOptimizerConfig(BaseOptimizerConfig):
    # optimizer_class: ClassVar[torch.optim.Optimizer] = torch.optim.Ftrl
    type: str = schema_utils.ProtectedString("ftrl")

    learning_rate_power: float = schema_utils.FloatRange(
        default=-0.5, max=0, parameter_metadata=OPTIMIZER_METADATA["learning_rate_power"]
    )

    initial_accumulator_value: float = schema_utils.NonNegativeFloat(
        default=0.1, parameter_metadata=OPTIMIZER_METADATA["initial_accumulator_value"]
    )

    l1_regularization_strength: float = schema_utils.NonNegativeFloat(
        default=0.0, parameter_metadata=OPTIMIZER_METADATA["l1_regularization_strength"]
    )

    l2_regularization_strength: float = schema_utils.NonNegativeFloat(
        default=0.0, parameter_metadata=OPTIMIZER_METADATA["l2_regularization_strength"]
    )


@DeveloperAPI
@register_optimizer(name="nadam")
@ludwig_dataclass
class NadamOptimizerConfig(BaseOptimizerConfig):
    optimizer_class: ClassVar[torch.optim.Optimizer] = torch.optim.NAdam
    """Points to `torch.optim.NAdam`."""

    type: str = schema_utils.ProtectedString("nadam")

    # Defaults taken from https://pytorch.org/docs/stable/generated/torch.optim.NAdam.html#torch.optim.NAdam :

    betas: tuple[float, float] = schema_utils.FloatRangeTupleDataclassField(
        default=(0.9, 0.999),
        description="Coefficients used for computing running averages of gradient and its square.",
        parameter_metadata=OPTIMIZER_METADATA["betas"],
    )

    eps: float = schema_utils.NonNegativeFloat(
        default=1e-08,
        description="Term added to the denominator to improve numerical stability.",
        parameter_metadata=OPTIMIZER_METADATA["eps"],
    )

    weight_decay: float = schema_utils.NonNegativeFloat(
        default=0.0, description="Weight decay ($L2$ penalty).", parameter_metadata=OPTIMIZER_METADATA["weight_decay"]
    )

    momentum_decay: float = schema_utils.NonNegativeFloat(
        default=4e-3, description="Momentum decay.", parameter_metadata=OPTIMIZER_METADATA["momentum_decay"]
    )


@DeveloperAPI
@register_optimizer(name="rmsprop")
@ludwig_dataclass
class RMSPropOptimizerConfig(BaseOptimizerConfig):
    """Parameters for rmsprop optimization."""

    optimizer_class: ClassVar[torch.optim.Optimizer] = torch.optim.RMSprop
    """Points to `torch.optim.RMSprop`."""

    type: str = schema_utils.ProtectedString("rmsprop")
    """Must be 'rmsprop' - corresponds to name in `ludwig.modules.optimization_modules.optimizer_registry`
       (default: 'rmsprop')"""

    # Defaults taken from https://pytorch.org/docs/stable/generated/torch.optim.RMSprop.html#torch.optim.RMSprop:
    momentum: float = schema_utils.NonNegativeFloat(
        default=0.0,
        description="Momentum factor.",
        parameter_metadata=OPTIMIZER_METADATA["momentum"],
    )

    alpha: float = schema_utils.NonNegativeFloat(
        default=0.99,
        description="Smoothing constant.",
        parameter_metadata=OPTIMIZER_METADATA["alpha"],
    )

    eps: float = schema_utils.NonNegativeFloat(
        default=1e-08,
        description="Term added to the denominator to improve numerical stability.",
        parameter_metadata=OPTIMIZER_METADATA["eps"],
    )

    centered: bool = schema_utils.Boolean(
        default=False,
        description="If True, computes the centered RMSProp, and the gradient is normalized by an estimation of its "
        "variance.",
        parameter_metadata=OPTIMIZER_METADATA["centered"],
    )

    weight_decay: float = schema_utils.NonNegativeFloat(default=0.0, description="Weight decay ($L2$ penalty).")


if bnb is not None:

    @DeveloperAPI
    @register_optimizer(name="rmsprop_8bit")
    @ludwig_dataclass
    class RMSProp8BitOptimizerConfig(RMSPropOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.RMSprop8bit

        type: str = schema_utils.ProtectedString("rmsprop_8bit")

        block_wise: bool = schema_utils.Boolean(
            default=True,
            description="Whether to use block wise update.",
        )

        percentile_clipping: int = schema_utils.IntegerRange(
            default=100,
            min=0,
            max=100,
            description="Percentile clipping.",
        )

        @property
        def is_8bit(self) -> bool:
            return True


if bnb is not None:

    @DeveloperAPI
    @register_optimizer(name="lamb")
    @ludwig_dataclass
    class LAMBOptimizerConfig(BaseOptimizerConfig):
        """Layer-wise Adaptive Moments optimizer for Batch training.

        Paper: https://arxiv.org/pdf/1904.00962.pdf
        """

        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.LAMB

        type: str = schema_utils.ProtectedString("lamb")

        bias_correction: bool = schema_utils.Boolean(
            default=True,
        )

        betas: tuple[float, float] = schema_utils.FloatRangeTupleDataclassField(
            default=(0.9, 0.999),
            description="Coefficients used for computing running averages of gradient and its square.",
            parameter_metadata=OPTIMIZER_METADATA["betas"],
        )

        eps: float = schema_utils.NonNegativeFloat(
            default=1e-08,
            description="Term added to the denominator to improve numerical stability.",
            parameter_metadata=OPTIMIZER_METADATA["eps"],
        )

        weight_decay: float = schema_utils.NonNegativeFloat(
            default=0.0,
            description="Weight decay (L2 penalty).",
            parameter_metadata=OPTIMIZER_METADATA["weight_decay"],
        )

        amsgrad: bool = schema_utils.Boolean(
            default=False,
            description=(
                "Whether to use the AMSGrad variant of this algorithm from the paper "
                "'On the Convergence of Adam and Beyond'."
            ),
            parameter_metadata=OPTIMIZER_METADATA["amsgrad"],
        )

        adam_w_mode: bool = schema_utils.Boolean(
            default=True,
            description="Whether to use the AdamW mode of this algorithm from the paper "
            "'Decoupled Weight Decay Regularization'.",
        )

        percentile_clipping: int = schema_utils.IntegerRange(
            default=100,
            min=0,
            max=100,
            description="Percentile clipping.",
        )

        block_wise: bool = schema_utils.Boolean(
            default=False,
            description="Whether to use block wise update.",
        )

        max_unorm: float = schema_utils.FloatRange(
            default=1.0,
            min=0.0,
            max=1.0,
        )

    @DeveloperAPI
    @register_optimizer(name="lamb_8bit")
    @ludwig_dataclass
    class LAMB8BitOptimizerConfig(LAMBOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.LAMB8bit

        type: str = schema_utils.ProtectedString("lamb_8bit")

        @property
        def is_8bit(self) -> bool:
            return True


if bnb is not None:

    @DeveloperAPI
    @register_optimizer(name="lars")
    @ludwig_dataclass
    class LARSOptimizerConfig(BaseOptimizerConfig):
        """Layerwise Adaptive Rate Scaling.

        Paper: https://arxiv.org/pdf/1708.03888.pdf
        """

        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.LARS

        type: str = schema_utils.ProtectedString("lars")

        # 0.9 taken from the original paper - momentum requires a non zero value
        # https://arxiv.org/pdf/1708.03888v3.pdf
        momentum: float = schema_utils.FloatRange(
            default=0.9,
            min=0.0,
            max=1.0,
            min_inclusive=False,
            description="Momentum factor.",
            parameter_metadata=OPTIMIZER_METADATA["momentum"],
        )

        dampening: float = schema_utils.FloatRange(
            default=0.0,
            min=0.0,
            max=1.0,
            description="Dampening for momentum.",
            parameter_metadata=OPTIMIZER_METADATA["dampening"],
        )

        weight_decay: float = schema_utils.NonNegativeFloat(
            default=0.0,
            description="Weight decay (L2 penalty).",
            parameter_metadata=OPTIMIZER_METADATA["weight_decay"],
        )

        nesterov: bool = schema_utils.Boolean(
            default=False,
            description="Enables Nesterov momentum.",
            parameter_metadata=OPTIMIZER_METADATA["nesterov"],
        )

        percentile_clipping: int = schema_utils.IntegerRange(
            default=100,
            min=0,
            max=100,
            description="Percentile clipping.",
        )

        max_unorm: float = schema_utils.FloatRange(
            default=1.0,
            min=0.0,
            max=1.0,
        )

    @DeveloperAPI
    @register_optimizer(name="lars_8bit")
    @ludwig_dataclass
    class LARS8BitOptimizerConfig(LARSOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.LARS8bit

        type: str = schema_utils.ProtectedString("lars_8bit")

        @property
        def is_8bit(self) -> bool:
            return True


if bnb is not None:

    @DeveloperAPI
    @register_optimizer(name="lion")
    @ludwig_dataclass
    class LIONOptimizerConfig(BaseOptimizerConfig):
        """Evolved Sign Momentum.

        Paper: https://arxiv.org/pdf/2302.06675.pdf
        """

        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.Lion

        type: str = schema_utils.ProtectedString("lion")

        betas: tuple[float, float] = schema_utils.FloatRangeTupleDataclassField(
            default=(0.9, 0.999),
            description="Coefficients used for computing running averages of gradient and its square.",
            parameter_metadata=OPTIMIZER_METADATA["betas"],
        )

        weight_decay: float = schema_utils.NonNegativeFloat(
            default=0.0,
            description="Weight decay (L2 penalty).",
            parameter_metadata=OPTIMIZER_METADATA["weight_decay"],
        )

        percentile_clipping: int = schema_utils.IntegerRange(
            default=100,
            min=0,
            max=100,
            description="Percentile clipping.",
        )

        block_wise: bool = schema_utils.Boolean(
            default=True,
            description="Whether to use block wise update.",
        )

    @DeveloperAPI
    @register_optimizer(name="lion_8bit")
    @ludwig_dataclass
    class LION8BitOptimizerConfig(LIONOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.Lion8bit

        type: str = schema_utils.ProtectedString("lion_8bit")

        @property
        def is_8bit(self) -> bool:
            return True

    @DeveloperAPI
    @register_optimizer(name="paged_lion")
    @ludwig_dataclass
    class PagedLionOptimizerConfig(LIONOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.PagedLion

        type: str = schema_utils.ProtectedString("paged_lion")

        @property
        def is_paged(self) -> bool:
            return True

    @DeveloperAPI
    @register_optimizer(name="paged_lion_8bit")
    @ludwig_dataclass
    class PagedLion8BitOptimizerConfig(PagedLionOptimizerConfig):
        optimizer_class: ClassVar[torch.optim.Optimizer] = bnb.optim.PagedLion8bit

        type: str = schema_utils.ProtectedString("paged_lion_8bit")

        @property
        def is_8bit(self) -> bool:
            return True


@DeveloperAPI
def get_optimizer_conds():
    """Returns a JSON schema of conditionals to validate against optimizer types defined in
    `ludwig.modules.optimization_modules.optimizer_registry`."""
    conds = []
    for optimizer in optimizer_registry:
        optimizer_cls = optimizer_registry[optimizer][1]
        other_props = schema_utils.unload_jsonschema_from_marshmallow_class(optimizer_cls)["properties"]
        schema_utils.remove_duplicate_fields(other_props)
        preproc_cond = schema_utils.create_cond(
            {"type": optimizer},
            other_props,
        )
        conds.append(preproc_cond)
    return conds


@DeveloperAPI
def OptimizerDataclassField(default="adam", description="", parameter_metadata: ParameterMetadata = None):
    """Custom dataclass field that when used inside of a dataclass will allow any optimizer in
    `ludwig.modules.optimization_modules.optimizer_registry`.

    Sets default optimizer to 'adam'.

    :param default: Dict specifying an optimizer with a `type` field and its associated parameters. Will attempt to use
           `type` to load optimizer from registry with given params. (default: {"type": "adam"}).
    :return: Initialized dataclass field that converts untyped dicts with params to optimizer dataclass instances.
    """

    class OptimizerSelection(schema_utils.TypeSelection):
        """Custom marshmallow field that deserializes a dict to a valid optimizer from
        `ludwig.modules.optimization_modules.optimizer_registry` and creates a corresponding `oneOf` JSON schema
        for external usage."""

        def __init__(self):
            super().__init__(
                registry=optimizer_registry,
                default_value=default,
                description=description,
                parameter_metadata=parameter_metadata,
            )

        def get_schema_from_registry(self, key: str) -> type[schema_utils.BaseMarshmallowConfig]:
            return get_optimizer_cls(key)

        def _jsonschema_type_mapping(self):
            # Note that this uses the same conditional pattern as combiners:
            return {
                "type": "object",
                "properties": {
                    "type": {
                        "type": "string",
                        "enum": list(optimizer_registry.keys()),
                        "default": default,
                        "description": "The type of optimizer to use during the learning process",
                    },
                },
                "title": "optimizer_options",
                "allOf": get_optimizer_conds(),
                "required": ["type"],
                "description": description,
            }

    return OptimizerSelection().get_default_field()


@DeveloperAPI
@ludwig_dataclass
class GradientClippingConfig(schema_utils.BaseMarshmallowConfig):
    """Dataclass that holds gradient clipping parameters."""

    clipglobalnorm: float | None = schema_utils.FloatRange(
        default=0.5,
        allow_none=True,
        description="Maximum allowed norm of the gradients",
        parameter_metadata=OPTIMIZER_METADATA["gradient_clipping"],
    )

    # TODO(travis): is this redundant with `clipglobalnorm`?
    clipnorm: float | None = schema_utils.FloatRange(
        default=None,
        allow_none=True,
        description="Maximum allowed norm of the gradients",
        parameter_metadata=OPTIMIZER_METADATA["gradient_clipping"],
    )

    clipvalue: float | None = schema_utils.FloatRange(
        default=None,
        allow_none=True,
        description="Maximum allowed value of the gradients",
        parameter_metadata=OPTIMIZER_METADATA["gradient_clipping"],
    )


@DeveloperAPI
def GradientClippingDataclassField(description: str, default: dict = {}):
    """Returns custom dataclass field for `ludwig.modules.optimization_modules.GradientClippingConfig`. Allows
    `None` by default.

    :param description: Description of the gradient dataclass field
    :param default: dict that specifies clipping param values that will be loaded by its schema class (default: {}).
    """
    allow_none = True

    class GradientClippingMarshmallowField(schema_utils.LudwigSchemaField):
        """Custom field class for gradient clipping.

        Deserializes a dict to a valid instance of `ludwig.modules.optimization_modules.GradientClippingConfig` and
        creates a corresponding JSON schema for external usage.
        """

        def _deserialize(self, value, attr, data, **kwargs):
            if value is None:
                return value
            if isinstance(value, dict):
                try:
                    return GradientClippingConfig.Schema().load(value)
                except (TypeError, ConfigValidationError):
                    raise ConfigValidationError(
                        f"Invalid params for gradient clipping: {value}, see GradientClippingConfig class."
                    )
            raise ConfigValidationError("Field should be None or dict")

        def _jsonschema_type_mapping(self):
            return {
                "oneOf": [
                    {"type": "null", "title": "disabled", "description": "Disable gradient clipping."},
                    {
                        **schema_utils.unload_jsonschema_from_marshmallow_class(GradientClippingConfig),
                        "title": "enabled_options",
                    },
                ],
                "title": "gradient_clipping_options",
                "description": description,
            }

    if not isinstance(default, dict):
        raise ConfigValidationError(f"Invalid default: `{default}`")

    def load_default():
        return GradientClippingConfig.Schema().load(default)

    dump_default = GradientClippingConfig.Schema().dump(default)

    return field(
        metadata={
            "marshmallow_field": GradientClippingMarshmallowField(
                allow_none=allow_none,
                load_default=load_default,
                dump_default=dump_default,
                metadata={
                    "description": description,
                    "parameter_metadata": convert_metadata_to_json(OPTIMIZER_METADATA["gradient_clipping"]),
                },
            )
        },
        default_factory=load_default,
    )
