r"""Implementation of the AdemamixAdopt optimizer.

AdemamixAdopt is based on QHADOPT (Quasi-hyperbolic ADOPT).
ADOPT Paper: https://arxiv.org/abs/2411.02853
Quasi-hyperbolic ADAM: https://arxiv.org/abs/1810.06801
Ademamix Paper: https://arxiv.org/abs/2409.03137

AdemamixAdopt (Adam/Adopt with Momentum Mixture) extends the Quasi-Hyperbolic
approach by utilizing two independent first momentums (exp_avg and exp_avg_2)
with separate beta values (beta1_a and beta1_b).

The final update direction is a convex combination of the two momentums and the
current (normalized and clipped) gradient, controlled by hyperparameters v1 and v2.

Update Direction D_t:
D_t = v1 * exp_avg_t + v2 * exp_avg_2_t + (1 - v1 - v2) * normed_grad_t

This allows for a more flexible adaptation by mixing momentums with different
time scales, while retaining the adaptive clipping stabilization from ADOPT.

This module provides both the high-level :class:`AdemamixAdopt` class, an `Optimizer`
subclass suitable for direct usage, and the low-level functional API
:func:`ademamixadopt` that performs the same computations.
"""

import math
from collections.abc import Callable
from typing import Any, ClassVar, cast

import torch
from composer.utils import dist
from llmfoundry.registry import optimizers
from torch import Tensor
from torch.optim.optimizer import (
    Optimizer,
    ParamsT,
    _default_to_fused_or_foreach,  # noqa: PLC2701
    _device_dtype_check_for_fused,  # noqa: PLC2701
    _disable_dynamo_if_unsupported,  # noqa: PLC2701
    _get_capturable_supported_devices,  # noqa: PLC2701
    _get_scalar_dtype,  # noqa: PLC2701
    _get_value,  # noqa: PLC2701
    _use_grad_for_differentiable,  # noqa: PLC2701
    _view_as_real,  # noqa: PLC2701
)
from torch.types import Number

from mt_dao.optimizers.utils import get_report_curvature

__all__ = ["AdemamixAdopt", "ademamixadopt"]


def _default_clip_lambda(step: Number | Tensor) -> float:
    r"""Return the default clipping value for the given step.

    By default, it uses :math:`\text{clip} = \text{step}^{0.25}`.
    Clipping is only active for step >= 1.

    Parameters
    ----------
    step : Number | Tensor | Any
         The current iteration/step value.

    Returns
    -------
    float
        The clipping threshold for the given step.

    Raises
    ------
    TypeError
        Step needs to be a Number or Tensor.

    """
    internal_step: int
    if isinstance(step, (Tensor)):
        try:
            internal_step = int(step.item())
        except (ValueError, TypeError):
            return 0.0
    elif isinstance(step, (Number)):
        internal_step = int(step)
    else:
        msg = f"step must be a Number or Tensor, but got {type(step)}"
        raise TypeError(msg)

    if internal_step < 1:
        return 0.0  # No clipping at step 0 (initialization)

    return internal_step**0.25


@optimizers.register_class("ademamixadopt")
class AdemamixAdopt(Optimizer):
    r"""Implements the AdemamixAdopt algorithm.

    This is a variant of Adam/Adopt featuring an adaptive clipping
    mechanism.
    Its update rule incorporates a convex combination of two independent first momentums
    and the current normalized gradient.

    Update = v1 * M1_t + v2 * M2_t + (1 - v1 - v2) * NormedGrad_t

    Args:
        params (ParamsT): Iterable of parameters to optimize
                        or dicts defining parameter groups.
        lr (float | Tensor, optional): Learning rate. Defaults to 1e-3.
        betas (tuple[float, float, float], optional): Coefficients for running averages.
            (beta1_a for exp_avg, beta1_b for exp_avg_2, beta2 for exp_avg_sq).
            Defaults to (0.9, 0.999, 0.9999).
        v1 (float, optional): Convex combination coefficient for exp_avg.
                            Defaults to 0.45.
        v2 (float, optional): Convex combination coefficient for exp_avg_2.
                            Defaults to 0.45.
        eps (float, optional): Term added to the denominator for stability.
                            Defaults to 1e-6.
        clip_lambda (Callable | None, optional): Function returning clipping threshold.
        weight_decay (float, optional): Weight decay coefficient. Defaults to 0.0.
        decouple (bool, optional): Whether to use decoupled weight decay (AdamW style).
                            Defaults to False.
        foreach (bool | None, optional): Whether to use the multi-tensor APIs.
        maximize (bool, optional): Maximize the objective. Defaults to False.
        capturable (bool, optional): Whether to enable CUDA graph capture.
        differentiable (bool, optional): Whether to enable differentiable optimization.
        fused (bool | None, optional): Whether to use fused kernels (Not Implemented).
        report_curvature (bool, optional): Whether to report curvature metrics.

    """

    metric_functions: ClassVar = {
        # Metric for the first momentum (exp_avg)
        "l2_norm/moment": lambda _param,
        optim_state,
        _step_tensor: torch.linalg.vector_norm(
            optim_state["exp_avg"],
        ),
        # Metric for the second momentum (exp_avg_2)
        "l2_norm/moment_b": lambda _param,
        optim_state,
        _step_tensor: torch.linalg.vector_norm(
            optim_state["exp_avg_2"],
        ),
        "l2_norm/moment2": lambda _param,
        optim_state,
        _step_tensor: torch.linalg.vector_norm(
            optim_state["exp_avg_sq"],
        ),
        "min/moment2": (
            lambda _param, optim_state, _step_tensor: torch.min(
                optim_state["exp_avg_sq"],
            )
        ),
        "max/moment2": (
            lambda _param, optim_state, _step_tensor: torch.max(
                optim_state["exp_avg_sq"],
            )
        ),
        "l2_norm/param": lambda param,
        _optim_state,
        _step_tensor: torch.linalg.vector_norm(
            param.data,
        ),
        "l2_norm/update": lambda _param,
        _optim_state,
        step_tensor: torch.linalg.vector_norm(
            step_tensor,
        )
        if step_tensor is not None
        else torch.tensor(0.0),
        "l2_norm/grad": lambda param,
        _optim_state,
        _step_tensor: torch.linalg.vector_norm(
            param.grad,
        ),
    }

    MAX_NUM_BETAS = 3

    def __init__(  # noqa: PLR0917, PLR0913, C901, PLR0912, PLR0915
        self,
        params: ParamsT,
        lr: float | Tensor = 1e-3,
        # Updated betas to include three values: (beta1_a, beta1_b, beta2)
        betas: tuple[float, float, float] = (0.95, 0.999, 0.9999),
        v1: float = 0.1,
        v2: float = 0.9,  # Added v2
        eps: float = 1e-6,
        clip_lambda: (
            Callable[[Number | Tensor | Any], float] | None
        ) = _default_clip_lambda,
        weight_decay: float = 0.0,
        *,
        decouple: bool = False,
        foreach: bool | None = None,
        maximize: bool = False,
        capturable: bool = False,
        differentiable: bool = False,
        fused: bool | None = None,
        report_curvature: bool = False,
    ) -> None:
        """AdemamixAdopt optimizer constructor.

        Parameters
        ----------
        params : ParamsT
            Iterable of parameters to optimize or dicts defining parameter groups.
        lr : float | Tensor, optional
            Learning rate. Defaults to 1e-3.
        betas : tuple[float, float, float], optional
            Coefficients used for computing running averages of gradient and its square.
            The first two values are for the two first momentums (exp_avg, exp_avg_2),
            and the third is for the second moment (exp_avg_sq).
            Defaults to (0.9, 0.999, 0.9999).
        v1 : float, optional
            Convex combination coefficient for exp_avg. Defaults to 0.1.
        v2 : float, optional
            Convex combination coefficient for exp_avg_2. Defaults to 0.9.
        eps : float, optional
            Term added to the denominator to improve numerical stability.
            Defaults to 1e-6.
        clip_lambda : Callable | None, optional
            A function that takes the current step (int or Tensor) and returns
            the clipping threshold (float). If None, no clipping is applied.
            Defaults to a function that returns step**0.25.
        weight_decay : float, optional
            Weight decay (L2 penalty) coefficient. Defaults to 0.0.
        decouple : bool, optional
            If True, uses decoupled weight decay as in AdamW.
            Defaults to False.
        foreach : bool | None, optional
            Whether to use the foreach implementation.
            If None, defaults to True if there are multiple parameter groups.
            Defaults to None.
        maximize : bool, optional
            If True, maximize the objective instead of minimizing.
            Defaults to False.
        capturable : bool, optional
            If True, enables CUDA graph capture.
            Defaults to False.
        differentiable : bool, optional
            If True, enables differentiable optimization.
            Defaults to False.
        fused : bool | None, optional
            If True, uses fused kernels for optimization.
            If None, defaults to False.
            Defaults to None.
        report_curvature : bool, optional
            If True, reports curvature metrics during optimization.
            Defaults to False.

        Raises
        ------
        ValueError
            If any of the hyperparameters are out of valid range.
        RuntimeError
            If `fused` is True (not implemented) or if incompatible options are set.

        """
        if not lr >= 0.0:
            msg = f"Invalid learning rate: {lr}"
            raise ValueError(msg)
        if isinstance(lr, Tensor) and foreach and not capturable:
            msg = (
                "lr as a Tensor is not supported for capturable=False and foreach=True"
            )
            raise ValueError(
                msg,
            )
        if not eps >= 0.0:
            msg = f"Invalid epsilon value: {eps}"
            raise ValueError(msg)
        # Validate all three betas
        if len(betas) != self.MAX_NUM_BETAS:
            msg = f"betas must be a tuple of 3 floats, but got {betas}"
            raise ValueError(msg)
        if not 0.0 <= betas[0] < 1.0:
            msg = f"Invalid beta parameter at index 0 (beta1_a): {betas[0]}"
            raise ValueError(msg)
        if not 0.0 <= betas[1] < 1.0:
            msg = f"Invalid beta parameter at index 1 (beta1_b): {betas[1]}"
            raise ValueError(msg)
        if not 0.0 < betas[2] < 1.0:
            # NOTE: beta2=0 is allowed by the optimizer logic,
            # but metrics reconstruction requires beta2 > 0.
            msg = f"Invalid beta parameter at index 2 (beta2): {betas[2]}"
            raise ValueError(msg)

        # Validate v1 and v2 for convex combination
        if not 0.0 <= v1 <= 1.0:
            msg = f"Invalid v1 parameter: {v1}"
            raise ValueError(msg)
        if not 0.0 <= v2 <= 1.0:
            msg = f"Invalid v2 parameter: {v2}"
            raise ValueError(msg)
        # Use a small tolerance for floating point comparison
        if not (v1 + v2) <= 1.0 + 1e-9:
            msg = f"""Invalid combination of v1 and v2
            (must sum to <= 1.0): {v1} + {v2} = {v1 + v2}"""
            raise ValueError(msg)

        if not weight_decay >= 0.0:
            msg = f"Invalid weight_decay value: {weight_decay}"
            raise ValueError(msg)

        self.clip_lambda = clip_lambda

        defaults = {
            "lr": lr,
            "betas": betas,
            "vs": (v1, v2),  # Updated vs tuple
            "eps": eps,
            "weight_decay": weight_decay,
            "decouple": decouple,
            "maximize": maximize,
            "foreach": foreach,
            "capturable": capturable,
            "differentiable": differentiable,
            "fused": fused,
        }
        super().__init__(params, defaults)
        for group in self.param_groups:
            group["initial_lr"] = group["lr"]

        if fused:
            if differentiable:
                msg = "`fused` does not support `differentiable`"
                raise RuntimeError(msg)
            self._step_supports_amp_scaling = True
            if foreach:
                msg = "`fused` and `foreach` cannot be `True` together."
                raise RuntimeError(msg)
            # NOTE: support fused
            msg = "`fused` is not currently supported for AdemamixAdopt"
            raise RuntimeError(msg)

        # Initialize curvature reporting
        self.curvature_metric_function: (
            Callable[[Tensor, str], dict[str, Tensor]] | None
        ) = None
        if report_curvature:
            self.curvature_metric_function = get_report_curvature()

    def __setstate__(self, state: dict) -> None:
        """Set the state of the optimizer for backward compatibility.

        Parameters
        ----------
        state : dict
            State dictionary containing optimizer state.

        """
        super().__setstate__(state)
        for group in self.param_groups:
            group.setdefault("maximize", False)
            group.setdefault("foreach", None)
            group.setdefault("capturable", False)
            group.setdefault("differentiable", False)
            group.setdefault("decouple", False)
            fused = group.setdefault("fused", None)
            for p in group["params"]:
                p_state = self.state.get(p, [])
                if len(p_state) != 0:
                    # Ensure exp_avg_2 exists if loading from a checkpoint
                    # that might not have it
                    if "exp_avg_2" not in p_state:
                        p_state["exp_avg_2"] = torch.zeros_like(
                            p,
                            memory_format=torch.preserve_format,
                        )

                    if not torch.is_tensor(p_state["step"]):
                        step_val = float(p_state["step"])
                        p_state["step"] = (
                            torch.tensor(
                                step_val,
                                dtype=_get_scalar_dtype(is_fused=fused),
                                device=p.device,
                            )
                            if group["capturable"] or group["fused"]
                            else torch.tensor(step_val, dtype=_get_scalar_dtype())
                        )

    def _init_group(  # noqa: PLR0917, PLR0913
        self,
        group: dict,
        params_with_grad: list[Tensor],
        grads: list[Tensor],
        exp_avgs: list[Tensor],
        exp_avg_2s: list[Tensor],  # Added exp_avg_2s
        exp_avg_sqs: list[Tensor],
        state_steps: list[Tensor],
    ) -> bool:
        """Initialize a parameter group for AdemamixAdopt updates.

        Parameters
        ----------
        group : dict
            Parameter group from the optimizer.
        params_with_grad : list[Tensor]
            List to populate with parameters that have gradients.
        grads : list[Tensor]
            List to populate with gradients of the parameters.
        exp_avgs : list[Tensor]
            List to populate with the first moment estimates (exp_avg).
        exp_avg_2s : list[Tensor]
            List to populate with the second first moment estimates (exp_avg_2).
        exp_avg_sqs : list[Tensor]
            List to populate with the second moment estimates (exp_avg_sq).
        state_steps : list[Tensor]
            List to populate with the step counts for each parameter.

        Returns
        -------
        bool
            True if any parameter is complex, False otherwise.

        Raises
        ------
        RuntimeError
            If sparse gradients are encountered or if incompatible options are set.

        """
        has_complex = False
        for p in group["params"]:
            if p.grad is None:
                continue
            has_complex |= torch.is_complex(p)
            params_with_grad.append(p)
            if p.grad.is_sparse:
                msg = "AdemamixAdopt does not support sparse gradients"
                raise RuntimeError(msg)
            grads.append(p.grad)

            state = self.state[p]

            # Lazy state initialization
            if len(state) == 0:
                if group["fused"]:
                    _device_dtype_check_for_fused(p)

                state["step"] = (
                    torch.zeros(
                        (),
                        dtype=_get_scalar_dtype(is_fused=group["fused"]),
                        device=p.device,
                    )
                    if group["capturable"] or group["fused"]
                    else torch.tensor(0.0, dtype=_get_scalar_dtype())
                )
                # Exponential moving average of gradient values (Momentum A)
                state["exp_avg"] = torch.zeros_like(
                    p,
                    memory_format=torch.preserve_format,
                )
                # Exponential moving average of gradient values (Momentum B)
                state["exp_avg_2"] = torch.zeros_like(
                    p,
                    memory_format=torch.preserve_format,
                )
                # Exponential moving average of squared gradient values
                state["exp_avg_sq"] = torch.zeros_like(
                    p,
                    memory_format=torch.preserve_format,
                )

            exp_avgs.append(state["exp_avg"])
            exp_avg_2s.append(state["exp_avg_2"])  # Added this line
            exp_avg_sqs.append(state["exp_avg_sq"])

            if group["differentiable"] and state["step"].requires_grad:
                msg = "`requires_grad` not supported for `step` in differentiable mode"
                raise RuntimeError(
                    msg,
                )

            # Foreach without capturable does not support a tensor lr
            if (
                group["foreach"]
                and isinstance(group["lr"], Tensor)
                and not group["capturable"]
            ):
                msg = "Tensor lr not supported for capturable=False and foreach=True"
                raise RuntimeError(
                    msg,
                )

            state_steps.append(state["step"])
        return has_complex

    @_use_grad_for_differentiable
    def step(  # pyright: ignore[reportIncompatibleMethodOverride]
        self,
        closure: Callable[[], torch.Tensor] | None = None,
    ) -> Tensor | None:
        """Perform a single AdemamixAdopt optimization step.

        Parameters
        ----------
        closure : Callable, optional
            A closure that reevaluates the model and returns the loss.

        Returns
        -------
        Tensor | None
            The loss value if a closure is provided, otherwise None.

        """
        self._cuda_graph_capture_health_check()

        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            params_with_grad: list[Tensor] = []
            grads: list[Tensor] = []
            exp_avgs: list[Tensor] = []
            exp_avg_2s: list[Tensor] = []  # Added list
            exp_avg_sqs: list[Tensor] = []
            state_steps: list[Tensor] = []
            # Updated betas extraction: (beta1_a, beta1_b, beta2)
            beta1_a, beta1_b, beta2 = cast("tuple[float, float, float]", group["betas"])
            # Updated vs extraction: (v1, v2)
            v1, v2 = cast("tuple[float, float]", group["vs"])

            # Handle initial_lr type checking for functional API compatibility
            initial_lr = group.get("initial_lr")

            has_complex = self._init_group(
                group,
                params_with_grad,
                grads,
                exp_avgs,
                exp_avg_2s,  # Added argument
                exp_avg_sqs,
                state_steps,
            )

            ademamixadopt(
                params_with_grad,
                grads,
                exp_avgs,
                exp_avg_2s,  # Added argument
                exp_avg_sqs,
                state_steps,
                initial_lr=initial_lr,
                decouple=group["decouple"],
                clip_lambda=self.clip_lambda,
                beta1_a=beta1_a,
                beta1_b=beta1_b,
                beta2=beta2,
                v1=v1,
                v2=v2,
                lr=group["lr"],
                weight_decay=group["weight_decay"],
                eps=group["eps"],
                maximize=group["maximize"],
                foreach=group["foreach"],
                capturable=group["capturable"],
                differentiable=group["differentiable"],
                fused=group["fused"],
                grad_scale=getattr(self, "grad_scale", None),
                found_inf=getattr(self, "found_inf", None),
                has_complex=has_complex,
            )

        return loss

    def dist_reduce_metrics(self, optimizer_metrics: dict) -> dict:  # noqa: PLR6301
        """Reduce metrics across workers.

        Parameters
        ----------
        optimizer_metrics : dict
            Dictionary of metrics to reduce.

        Returns
        -------
        dict
            Dictionary with reduced metrics.

        """
        if dist.get_world_size() <= 1:
            return optimizer_metrics

        local_keys = list(optimizer_metrics.keys())
        all_gathered_keys = dist.all_gather_object(local_keys)
        all_keys = set()
        for keys in all_gathered_keys:
            all_keys.update(keys)

        # Sort keys to ensure every rank has the same keys order
        all_keys_list = sorted(all_keys)

        # Determine device for reduction tensors
        device = (
            torch.cuda.current_device()
            if torch.cuda.is_available()
            else torch.device("cpu")
        )

        for metric in all_keys_list:
            value = optimizer_metrics.get(metric, torch.tensor(0.0, device=device))

            # Ensure value is a tensor before reduction
            if not isinstance(value, Tensor):
                value = torch.tensor(float(value), device=device)

            if metric.startswith("l2_norm"):
                # L2 norms were squared in pre_reduce_metrics, so we sum and then sqrt
                dist.all_reduce(value, reduce_operation="SUM")
                # Use math.sqrt for scalar result consistency
                optimizer_metrics[metric] = math.sqrt(value.item())
            elif metric.startswith("min"):
                dist.all_reduce(value, reduce_operation="MIN")

                optimizer_metrics[metric] = value
            elif metric.startswith("max"):
                dist.all_reduce(value, reduce_operation="MAX")

                optimizer_metrics[metric] = value
            else:
                if dist.get_world_size() > 1:
                    dist.all_reduce(value, reduce_operation="SUM")
                optimizer_metrics[metric] = value / dist.get_world_size()

        return optimizer_metrics

    def pre_reduce_metrics(self, optimizer_metrics: dict) -> dict:  # noqa: PLR6301
        """Preprocess metrics to reduce across ranks correctly.

        Parameters
        ----------
        optimizer_metrics : dict
            Dictionary of metrics to preprocess.

        Returns
        -------
        dict
            Dictionary with preprocessed metrics.

        """
        # L2 norms need to be squared, before they are reduced via summation
        for metric, value in optimizer_metrics.items():
            if metric.startswith("l2_norm"):
                if isinstance(value, Tensor):
                    optimizer_metrics[metric] = value**2
                else:
                    # Handle potential non-tensor metrics
                    optimizer_metrics[metric] = value**2

        return optimizer_metrics

    def report_per_parameter_metrics(  # noqa: PLR0914
        self,
        param: torch.Tensor,
        name: str,
        optimizer_metrics: dict,
    ) -> dict:
        """Report the metrics per parameter.

        This function reconstructs the exact update applied to the parameter
        during the optimization step, including the contributions from the
        momentum mixture and decoupled weight decay, if applicable.

        Parameters
        ----------
        param : torch.Tensor
            The parameter tensor for which to report metrics.
        name : str
            The name of the parameter.
        optimizer_metrics : dict
            Dictionary to populate with computed metrics.

        Returns
        -------
        dict
            Updated dictionary with computed metrics.

        """
        # Retrieve hyperparameters from the first param group (assuming homogeneity)
        group = self.param_groups[0]
        lr = group["lr"]
        weight_decay = group["weight_decay"]
        initial_lr = group.get("initial_lr")
        decouple = group["decouple"]
        # Extract betas (only beta2 needed for reconstruction) and vs
        _, _, beta2 = cast("tuple[float, float, float]", group["betas"])
        v1, v2 = cast("tuple[float, float]", group["vs"])
        eps = cast("float", group["eps"])

        if param in self.state:
            param_optim_state = self.state[param]
            # NOTE: This section reconstructs the AdemamixAdopt update (step_tensor).

            grad = param.grad

            if grad is None:
                return optimizer_metrics

            # Retrieve buffers
            # (these are the updated ones, at T+1, if called after step())
            exp_avg = cast("Tensor", param_optim_state["exp_avg"])  # 1st moment A
            exp_avg_2 = cast("Tensor", param_optim_state["exp_avg_2"])  # 1st moment B
            exp_avg_sq = cast("Tensor", param_optim_state["exp_avg_sq"])  # 2nd moment
            step_t = cast("Tensor", param_optim_state["step"])

            # The functional API increments the step at the very end.
            # If this function is called after optimizer.step(), step_t holds T+1.
            step_t_plus_1 = _get_value(step_t)

            # ADOPT initializes at step 0. The first update occurs at T=1.
            if step_t_plus_1 < 1:
                return optimizer_metrics

            # The update we are reconstructing occurred at step T.
            step_t = step_t_plus_1 - 1

            if step_t < 1:
                # If T=0, only initialization happened.
                return optimizer_metrics

            # 1) Recover the previous second moment (V_{T-1}) used for normalization.
            # V_T = beta2 * V_{T-1} + (1-beta2) * G_T^2
            # V_{T-1} = (V_T - (1-beta2) * G_T^2) / beta2

            # Use torch.square for element-wise square
            old_exp_avg_sq = (exp_avg_sq - torch.square(grad) * (1.0 - beta2)) / beta2

            step_tensor = None

            # 2) Recompute the denominator
            denom = torch.clamp(old_exp_avg_sq.sqrt(), eps)

            # 3) Normalize the gradient
            normed_grad = grad.div(denom)

            # 4) Apply clipping using the correct step count (T)
            if self.clip_lambda is not None:
                # We use step_T as this was the value
                # used during the functional update.
                clip = self.clip_lambda(step_t)

                normed_grad.clamp_(-clip, clip)

            # 5) Calculate the update direction (D_T).
            # D_T = (v1 * M1_T + v2 * M2_T +
            #          (1 - v1 - v2) * NormedGrad_T)
            # We use the updated momentums (M1_T, M2_T) stored in the state.

            # In-place modification of normed_grad to become D_T for efficiency.
            v_grad = 1.0 - v1 - v2
            normed_grad.mul_(v_grad)
            normed_grad.add_(exp_avg, alpha=v1)
            normed_grad.add_(exp_avg_2, alpha=v2)

            # 6) Calculate the gradient-based update (lr * D_T)
            # Handle tensor LR if capturable=True

            step_tensor = _get_value(lr) * normed_grad

            if weight_decay != 0 and decouple:
                decay_factor = (lr / initial_lr) if initial_lr else 1.0
                scaling_factor = (decay_factor * weight_decay) / (
                    1 - decay_factor * weight_decay
                )
                step_tensor.mul_(1 + scaling_factor).add_(
                    param,
                    alpha=scaling_factor,
                )

            for metric in self.metric_functions:
                optimizer_metrics[f"{metric}/{name}"] = self.metric_functions[metric](
                    param,
                    param_optim_state,
                    step_tensor,
                )

            # Report curvature metrics if enabled
            if self.curvature_metric_function is not None:
                optimizer_metrics.update(self.curvature_metric_function(param, name))

        return optimizer_metrics


# --- Functional API Implementation ---


def _single_tensor_ademamixadopt(  # noqa: PLR0917, PLR0913
    params: list[Tensor],
    grads: list[Tensor],
    exp_avgs: list[Tensor],
    exp_avg_2s: list[Tensor],  # Added
    exp_avg_sqs: list[Tensor],
    state_steps: list[Tensor],
    grad_scale: Tensor | None,
    found_inf: Tensor | None,
    *,
    initial_lr: float | Tensor | None,
    decouple: bool,
    clip_lambda: Callable[[Number | Tensor | Any], float] | None,
    beta1_a: float,  # Renamed
    beta1_b: float,  # Added
    beta2: float,
    v1: float,
    v2: float,  # Added
    lr: float | Tensor,
    weight_decay: float,
    eps: float,
    maximize: bool,
    capturable: bool,
    differentiable: bool,
    has_complex: bool,  # noqa: ARG001
) -> None:
    """Single-tensor implementation of AdemamixAdopt step."""
    assert grad_scale is None
    assert found_inf is None

    if torch.jit.is_scripting():  # type: ignore # noqa: PGH003
        assert isinstance(lr, float)
        if initial_lr is not None:
            assert isinstance(initial_lr, float)

    for i, param in enumerate(params):
        # Prepare gradient
        grad = grads[i] if not maximize else -grads[i]

        exp_avg = exp_avgs[i]
        exp_avg_2 = exp_avg_2s[i]  # Added
        exp_avg_sq = exp_avg_sqs[i]
        step_t = state_steps[i]

        # If compiling, the compiler will handle cudagraph checks
        if not torch._utils.is_compiling() and capturable:  # pyright: ignore[reportAttributeAccessIssue]  # noqa: SLF001
            device = _get_capturable_supported_devices()
            assert param.device.type == step_t.device.type
            assert param.device.type in device, (
                f"If capturable=True, params, state_steps must be on: {device}."
            )

        step = step_t if capturable or differentiable else _get_value(step_t)

        # Perform weight decay if not decoupled (L2 regularization)
        if weight_decay != 0 and not decouple:
            # Standard L2 regularization: grad += WD * param
            # This applies the penalty correctly regardless of the maximization flag.
            grad = grad.add(param, alpha=weight_decay)

        if torch.is_complex(param):
            grad = torch.view_as_real(grad)
            exp_avg = torch.view_as_real(exp_avg)
            exp_avg_2 = torch.view_as_real(exp_avg_2)  # Added
            exp_avg_sq = torch.view_as_real(exp_avg_sq)
            param = torch.view_as_real(param)  # noqa: PLW2901

        # NOTE: During the zero-th step (when `step == 0``),
        # the algorithm initializes the second momenta vectors
        if step == 0:
            exp_avg_sq.addcmul_(grad, grad)
            # update step
            step_t += 1
            continue

        # Perform decoupled weight decay (AdamW style)
        # Perform stepweight decay if decoupled
        if weight_decay != 0 and decouple:
            # NOTE: The decay factor follows the same schedule of the learning rate
            # w/o being scaled by it
            decay_factor = (lr / initial_lr) if initial_lr else 1.0
            param.mul_(1 - decay_factor * weight_decay)

        # Calculate normalized gradient
        # Uses the second moment from the previous step (V_{t-1}),
        # which is currently in exp_avg_sq.
        denom = torch.clamp(exp_avg_sq.sqrt(), eps)
        normed_grad = grad.div(denom)

        # Apply clipping
        if clip_lambda is not None:
            # step is >= 1 here.
            clip = clip_lambda(step)
            normed_grad.clamp_(-clip, clip)

        # Update first momentums
        exp_avg.lerp_(normed_grad, 1 - beta1_a)
        exp_avg_2.lerp_(normed_grad, 1 - beta1_b)

        # Calculate the update direction (AdemamixAdopt update)
        # D_t := v1 * M1_t + v2 * M2_t + (1 - v1 - v2) * NormedGrad_t

        # In-place modification of normed_grad to become D_t for efficiency.
        v_grad = 1.0 - v1 - v2
        normed_grad.mul_(v_grad)
        normed_grad.add_(exp_avg, alpha=v1)
        normed_grad.add_(exp_avg_2, alpha=v2)

        # Final parameter update (normed_grad now holds D_t)

        param.add_(normed_grad, alpha=-_get_value(lr))

        # Update second momentum (V_t) using the raw gradient
        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

        # update step
        step_t += 1


def _multi_tensor_ademamixadopt(  # noqa: PLR0915, PLR0912, PLR0917, PLR0913, C901, PLR0914
    params: list[Tensor],
    grads: list[Tensor],
    exp_avgs: list[Tensor],
    exp_avg_2s: list[Tensor],  # Added
    exp_avg_sqs: list[Tensor],
    state_steps: list[Tensor],
    grad_scale: Tensor | None,
    found_inf: Tensor | None,
    *,
    initial_lr: float | Tensor | None,
    has_complex: bool,
    beta1_a: float,  # Renamed
    beta1_b: float,  # Added
    beta2: float,
    v1: float,
    v2: float,  # Added
    lr: float | Tensor,
    clip_lambda: Callable[[Number | Tensor | Any], float] | None,
    weight_decay: float,
    decouple: bool,
    eps: float,
    maximize: bool,
    capturable: bool,
    differentiable: bool,
) -> None:
    """Multi-tensor implementation of the AdemamixAdopt step.

    Parameters
    ----------
    params : list[Tensor]
        List of parameters to be updated.
    grads : list[Tensor]
        List of gradients corresponding to the parameters.
    exp_avgs : list[Tensor]
        List of first moment estimates (exp_avg).
    exp_avg_2s : list[Tensor]
        List of second first moment estimates (exp_avg_2).
    exp_avg_sqs : list[Tensor]
        List of second moment estimates (exp_avg_sq).
    state_steps : list[Tensor]
        List of step counts for each parameter.
    grad_scale : Tensor | None
        Gradient scaling factor for mixed precision training.
    found_inf : Tensor | None
        Indicator for found infinity in gradients.
    initial_lr : float | Tensor | None
        Initial learning rate for decoupled weight decay.
    has_complex : bool
        Indicates if any parameter is complex.
    beta1_a : float
        Exponential decay rate for the first moment estimate (exp_avg).
    beta1_b : float
        Exponential decay rate for the second first moment estimate (exp_avg_2).
    beta2 : float
        Exponential decay rate for the second moment estimate (exp_avg_sq).
    v1 : float
        Convex combination coefficient for exp_avg.
    v2 : float
        Convex combination coefficient for exp_avg_2.
    lr : float | Tensor
        Learning rate.
    clip_lambda : Callable[[Number | Tensor | Any], float] | None
        Function to compute clipping threshold based on the current step.
    weight_decay : float
        Weight decay (L2 penalty) coefficient.
    decouple : bool
        If True, uses decoupled weight decay as in AdamW.
    eps : float
        Term added to the denominator to improve numerical stability.
    maximize : bool
        If True, maximize the objective instead of minimizing.
    capturable : bool
        If True, enables CUDA graph capture.
    differentiable : bool
        If True, enables differentiable optimization.

    Raises
    ------
    RuntimeError
        If incompatible options are set or if capturable conditions are not met.

    """
    if len(params) == 0:
        return

    if isinstance(lr, Tensor) and not capturable:
        msg = "lr as a Tensor is not supported for capturable=False and foreach=True"
        raise RuntimeError(
            msg,
        )

    # If compiling, the compiler will handle cudagraph checks
    if not torch._utils.is_compiling() and capturable:  # pyright: ignore[reportAttributeAccessIssue]  # noqa: SLF001
        device = _get_capturable_supported_devices(
            supports_xla=False,
        )
        assert all(
            p.device.type == step.device.type and p.device.type in device
            for p, step in zip(params, state_steps, strict=False)
        ), f"If capturable=True, params, state_steps must be on: {device}."

    assert not differentiable, "_foreach ops don't support autograd"

    assert grad_scale is None
    assert found_inf is None

    # Updated grouping to include exp_avg_2s
    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(  # noqa: SLF001
        [params, grads, exp_avgs, exp_avg_2s, exp_avg_sqs, state_steps],  # pyright: ignore[reportArgumentType]
    )
    for (
        device_params_,
        device_grads_,
        device_exp_avgs_,
        device_exp_avg_2s_,  # Added
        device_exp_avg_sqs_,
        device_state_steps_,
    ), _ in grouped_tensors.values():
        device_params = cast("list[Tensor]", device_params_)
        device_grads = cast("list[Tensor]", device_grads_)
        device_exp_avgs = cast("list[Tensor]", device_exp_avgs_)
        device_exp_avg_2s = cast("list[Tensor]", device_exp_avg_2s_)  # Added
        device_exp_avg_sqs = cast("list[Tensor]", device_exp_avg_sqs_)
        device_state_steps = cast("list[Tensor]", device_state_steps_)

        # Handle complex parameters
        if has_complex:
            _view_as_real(
                device_params,
                device_grads,
                device_exp_avgs,
                device_exp_avg_2s,  # Added
                device_exp_avg_sqs,
            )

        if maximize:
            torch._foreach_neg_(device_grads)

        # L2 Regularization (not decoupled)
        if weight_decay != 0 and not decouple:
            # Standard L2 regularization: grad += WD * param.
            # This applies correctly regardless of maximize flag.
            torch._foreach_add_(
                device_grads,
                device_params,
                alpha=weight_decay,
            )

        # NOTE: Check for the zero-th step (initialization)
        # Assuming synchronized steps, check the first one.
        current_step = _get_value(device_state_steps[0])

        if current_step == 0:
            torch._foreach_addcmul_(device_exp_avg_sqs, device_grads, device_grads)

            # Update steps
            if not torch._utils.is_compiling() and device_state_steps[0].is_cpu:  # pyright: ignore[reportAttributeAccessIssue]  # noqa: SLF001
                torch._foreach_add_(
                    device_state_steps,
                    torch.tensor(1.0, device="cpu"),
                    alpha=1.0,
                )
            else:
                torch._foreach_add_(device_state_steps, 1)
            # Skip the rest of the loop
            continue

        # Decoupled Weight Decay (AdamW style)
        if weight_decay != 0 and decouple:
            # Decay factor R = lr / initial_lr
            decay_factor = 1.0
            if initial_lr:
                lr_scalar = _get_value(lr)
                initial_lr_scalar = _get_value(initial_lr)
                if initial_lr_scalar != 0:
                    decay_factor = lr_scalar / initial_lr_scalar

            # P_t = P_{t-1} * (1 - WD_scaled)

            weight_decay_scaled = _get_value(decay_factor) * weight_decay
            torch._foreach_mul_(device_params, 1.0 - weight_decay_scaled)

        # Calculate denominator (sqrt(V_{t-1}))
        exp_avg_sq_sqrt = torch._foreach_sqrt(device_exp_avg_sqs)  # noqa: SLF001
        torch._foreach_maximum_(exp_avg_sq_sqrt, eps)

        # Calculate normalized gradient
        normed_grad = torch._foreach_div(device_grads, exp_avg_sq_sqrt)  # noqa: SLF001

        # Apply clipping
        if clip_lambda is not None:
            # step is >= 1 here.
            clip = clip_lambda(current_step)

            torch._foreach_clamp_max_(normed_grad, clip)
            torch._foreach_clamp_min_(normed_grad, -clip)

        # Update first momentums
        torch._foreach_lerp_(device_exp_avgs, normed_grad, 1 - beta1_a)
        torch._foreach_lerp_(device_exp_avg_2s, normed_grad, 1 - beta1_b)

        # Calculate the update direction (AdemamixAdopt update)
        # D_t := v1 * M1_t + v2 * M2_t + (1 - v1 - v2) * NormedGrad_t

        # In-place modification of normed_grad to become D_t.
        v_grad = 1.0 - v1 - v2
        torch._foreach_mul_(normed_grad, v_grad)
        torch._foreach_add_(normed_grad, device_exp_avgs, alpha=v1)
        torch._foreach_add_(normed_grad, device_exp_avg_2s, alpha=v2)

        # Apply update (normed_grad now holds D_t)

        torch._foreach_add_(device_params, normed_grad, alpha=-_get_value(lr))

        # Update second momentum (V_t)
        torch._foreach_mul_(device_exp_avg_sqs, beta2)
        torch._foreach_addcmul_(
            device_exp_avg_sqs,
            device_grads,
            device_grads,
            value=1 - beta2,
        )

        # Update steps
        if not torch._utils.is_compiling() and device_state_steps[0].is_cpu:  # pyright: ignore[reportAttributeAccessIssue] # noqa: SLF001
            torch._foreach_add_(
                device_state_steps,
                torch.tensor(1.0, device="cpu"),
                alpha=1.0,
            )
        else:
            torch._foreach_add_(device_state_steps, 1)


def _fused_ademamixadopt(*args: Any, **kwargs: Any) -> None:  # noqa: ANN401
    """Fused implementation of the AdemamixAdopt step.

    Parameters
    ----------
    *args: Any
        Positional arguments for the fused AdemamixAdopt function.
    **kwargs: Any
        Keyword arguments for the fused AdemamixAdopt function.

    """
    msg = "Fused implementation of AdemamixAdopt is not available."
    raise NotImplementedError(msg)


# We need to define the single tensor function reference before the decorator
@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_ademamixadopt)
def ademamixadopt(  # noqa: PLR0913, PLR0917
    params: list[Tensor],
    grads: list[Tensor],
    exp_avgs: list[Tensor],
    exp_avg_2s: list[Tensor],  # Added
    exp_avg_sqs: list[Tensor],
    state_steps: list[Tensor],
    # kwonly args
    *,
    initial_lr: float | Tensor | None = None,
    foreach: bool | None = None,
    capturable: bool = False,
    differentiable: bool = False,
    fused: bool | None = None,
    grad_scale: Tensor | None = None,
    found_inf: Tensor | None = None,
    has_complex: bool = False,
    beta1_a: float,
    beta1_b: float,
    beta2: float,
    v1: float,
    v2: float,
    lr: float | Tensor,
    clip_lambda: Callable[[Number | Tensor | Any], float] | None,
    weight_decay: float,
    decouple: bool,
    eps: float,
    maximize: bool,
) -> None:
    r"""Functional API that performs the AdemamixAdopt algorithm step.

    It directly applies the parameter updates in-place.

    Parameters
    ----------
    params : list[Tensor]
        List of parameters to be updated.
    grads : list[Tensor]
        List of gradients corresponding to the parameters.
    exp_avgs : list[Tensor]
        List of first moment estimates (exp_avg).
    exp_avg_2s : list[Tensor]
        List of second first moment estimates (exp_avg_2).
    exp_avg_sqs : list[Tensor]
        List of second moment estimates (exp_avg_sq).
    state_steps : list[Tensor]
        List of step counts for each parameter.
    initial_lr : float | Tensor | None, optional
        Initial learning rate for decoupled weight decay. Default is None.
    foreach : bool | None, optional
        Whether to use the foreach implementation. If None, it will be
        determined automatically. Default is None.
    capturable : bool, optional
        If True, enables CUDA graph capture. Default is False.
    differentiable : bool, optional
        If True, enables differentiable optimization. Default is False.
    fused : bool | None, optional
        Whether to use the fused implementation. If None, it will be
        determined automatically. Default is None.
    grad_scale : Tensor | None, optional
        Gradient scaling factor for mixed precision training. Default is None.
    found_inf : Tensor | None, optional
        Indicator for found infinity in gradients. Default is None.
    has_complex : bool, optional
        Indicates if any parameter is complex. Default is False.
    beta1_a : float
        Exponential decay rate for the first moment estimate (exp_avg).
    beta1_b : float
        Exponential decay rate for the second first moment estimate (exp_avg_2).
    beta2 : float
        Exponential decay rate for the second moment estimate (exp_avg_sq).
    v1 : float
        Convex combination coefficient for exp_avg.
    v2 : float
        Convex combination coefficient for exp_avg_2.
    lr : float | Tensor
        Learning rate.
    clip_lambda : Callable[[Number | Tensor | Any], float] | None
        Function to compute clipping threshold based on the current step.
    weight_decay : float
        Weight decay (L2 penalty) coefficient.
    decouple : bool
        If True, uses decoupled weight decay as in AdamW.
    eps : float
        Term added to the denominator to improve numerical stability.
    maximize : bool
        If True, maximize the objective instead of minimizing.

    Raises
    ------
    RuntimeError
        If incompatible options are set or if capturable conditions are not met.

    """
    # Logic for selecting implementation (foreach vs fused vs single-tensor)
    if fused is None and foreach is None:
        _, foreach = _default_to_fused_or_foreach(
            params,
            differentiable,
            use_fused=False,  # Fused not supported
        )
        # NOTE: Do not flip on foreach for the unsupported case where lr
        # is a Tensor and capturable=False.
        if foreach and isinstance(lr, Tensor) and not capturable:
            foreach = False
    if fused is None:
        fused = False
    if foreach is None:
        foreach = False

    if foreach and torch.jit.is_scripting():  # pyright: ignore[reportPrivateImportUsage]
        msg = "torch.jit.script not supported with foreach optimizers"
        raise RuntimeError(msg)
    if fused and torch.jit.is_scripting():  # pyright: ignore[reportPrivateImportUsage]
        msg = "torch.jit.script not supported with fused optimizers"
        raise RuntimeError(msg)

    if fused:
        # Fused implementation raises NotImplementedError
        func = _fused_ademamixadopt
    elif foreach:
        func = _multi_tensor_ademamixadopt
    else:
        func = _single_tensor_ademamixadopt

    func(
        params,
        grads,
        exp_avgs,
        exp_avg_2s,  # Added
        exp_avg_sqs,
        state_steps,
        initial_lr=initial_lr,
        decouple=decouple,
        clip_lambda=clip_lambda,
        beta1_a=beta1_a,
        beta1_b=beta1_b,
        beta2=beta2,
        v1=v1,
        v2=v2,
        lr=lr,
        weight_decay=weight_decay,
        eps=eps,
        maximize=maximize,
        capturable=capturable,
        differentiable=differentiable,
        grad_scale=grad_scale,
        found_inf=found_inf,
        has_complex=has_complex,
    )
