r"""Implementation of the AdemamixAdam optimizer.

AdemamixAdam (Adam with Momentum Mixture) is a variant of the Adam optimizer
that incorporates ideas from Quasi-Hyperbolic Adam (QHAdam). It utilizes two
independent first momentums (M1 and M2) with separate beta values
(beta1_a and beta1_b).

QHAdam Paper: https://arxiv.org/abs/1810.06801
Adam Paper: https://arxiv.org/abs/1412.6980

The final update direction is a convex combination of the two bias-corrected
momentums and the current gradient (G_t), controlled by hyperparameters v1 and v2.
This combination is then normalized by the standard Adam denominator.

Update Numerator N_t (QHAdam style):
N_t = v1 * M1_hat_t + v2 * M2_hat_t + (1 - v1 - v2) * G_t

Update Rule:
P_t = P_{t-1} - lr * (N_t / (sqrt(V_hat_t) + eps))

This allows for flexible adaptation by mixing momentums with different time scales.
It supports both standard L2 regularization and decoupled weight decay (AdamW style)
via the 'decouple' flag.
"""

import logging
import math
from collections.abc import Callable
from typing import 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 mt_dao.optimizers.utils import get_report_curvature

log = logging.getLogger(__name__)

__all__ = ["AdemamixAdam", "ademamixadam"]


@optimizers.register_class("ademamixadam")
class AdemamixAdam(Optimizer):
    r"""Implements the AdemamixAdam algorithm.

    A variant of Adam incorporating a Quasi-Hyperbolic update rule using a convex
    combination of two independent first momentums and the current gradient.

    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, beta1_b, beta2). Defaults to (0.9, 0.99, 0.999).
        v1 (float, optional): Convex combination coefficient for exp_avg (M1).
                            Defaults to 0.1.
        v2 (float, optional): Convex combination coefficient for exp_avg_2 (M2).
                            Defaults to 0.9.
        eps (float, optional): Term added to the denominator for stability.
                                Defaults to 1e-8.
        weight_decay (float, optional): Weight decay coefficient. Defaults to 1e-2.
        decouple (bool, optional): Whether to use decoupled weight decay (AdamW style).
                                    Defaults to True.
        amsgrad (bool, optional): Whether to use the AMSGrad variant. 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 = {
        "l2_norm/moment1_a": lambda _p, s, _t: torch.linalg.vector_norm(s["exp_avg"]),
        "l2_norm/moment1_b": lambda _p, s, _t: torch.linalg.vector_norm(s["exp_avg_2"]),
        "l2_norm/moment2": lambda _p, s, _t: torch.linalg.vector_norm(s["exp_avg_sq"]),
        "min/moment2": lambda _p, s, _t: torch.min(s["exp_avg_sq"]),
        "max/moment2": lambda _p, s, _t: torch.max(s["exp_avg_sq"]),
        "l2_norm/param": lambda p, _s, _t: torch.linalg.vector_norm(p.data),
        "l2_norm/update": lambda _p, _s, t: torch.linalg.vector_norm(t)
        if t is not None
        else torch.tensor(0.0),
        "l2_norm/grad": lambda p, _s, _t: torch.linalg.vector_norm(p.grad)
        if p.grad is not None
        else torch.tensor(0.0),
    }

    MAX_NUM_BETAS = 3

    def __init__(  # noqa: PLR0917, PLR0913, C901, PLR0912
        self,
        params: ParamsT,
        lr: float | Tensor = 1e-3,
        betas: tuple[float, float, float] = (0.9, 0.99, 0.999),
        v1: float = 0.1,
        v2: float = 0.9,
        eps: float = 1e-8,
        weight_decay: float = 1e-2,
        *,
        decouple: bool = True,  # Defaulting to True for AdamW style
        amsgrad: bool = False,
        foreach: bool | None = None,
        maximize: bool = False,
        capturable: bool = False,
        differentiable: bool = False,
        fused: bool | None = None,
        report_curvature: bool = False,
    ) -> None:
        """Initialize the AdemamixAdam optimizer.

        Parameters
        ----------
        params: ParamsT
            Iterable of parameters to optimize or dicts defining parameter groups.
        lr: float | Tensor, optional
            Learning rate. Default: 1e-3.
        betas: tuple[float, float, float], optional
            Coefficients used for computing running averages of gradient and its square.
            Should be a tuple of three floats: (beta1_a, beta1_b, beta2).
            Default: (0.9, 0.99, 0.999).
        v1: float, optional
            Convex combination coefficient for the first momentum (exp_avg).
            Default: 0.1.
        v2: float, optional
            Convex combination coefficient for the second momentum (exp_avg_2).
            Default: 0.9.
        eps: float, optional
            Term added to the denominator to improve numerical stability.
            Default: 1e-8.
        weight_decay: float, optional
            Weight decay (L2 penalty) coefficient.
            Default: 1e-2.
        decouple: bool, optional
            If True, uses decoupled weight decay (AdamW style).
            If False, uses standard L2 regularization.
            Default: True.
        amsgrad: bool, optional
            If True, uses the AMSGrad variant of this algorithm.
            Default: False.
        foreach: bool | None, optional
            If True, uses the foreach implementation to potentially speed up
            the optimizer step. If None, defaults to fused or foreach based on
            other parameters.
            Default: None.
        maximize: bool, optional
            If True, maximizes the objective instead of minimizing.
            Default: False.
        capturable: bool, optional
            If True, enables the optimizer to be used in CUDA graph capture.
            Default: False.
        differentiable: bool, optional
            If True, enables higher order gradients through the optimizer step.
            Default: False.
        fused: bool | None, optional
            If True, uses fused kernels to potentially speed up the optimizer step.
            If None, defaults to fused or foreach based on other parameters.
            Default: None.
        report_curvature: bool, optional
            If True, reports curvature metrics during optimization.
            Default: False.

        Raises
        ------
        ValueError:
            If any of the input parameters are invalid.
        RuntimeError:
            If fused=True is specified (not implemented).

        """
        # Input validation checks
        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)
        if len(betas) != self.MAX_NUM_BETAS:
            msg = (
                f"betas must be a tuple of {self.MAX_NUM_BETAS} floats, but got {betas}"
            )
            raise ValueError(msg)
        if not all(0.0 <= b < 1.0 for b in betas):
            msg = f"Invalid beta parameters: {betas}"
            raise ValueError(msg)

        # Validate v1 and v2 for convex combination
        if not (0.0 <= v1 <= 1.0 and 0.0 <= v2 <= 1.0):
            msg = f"Invalid v parameters: v1={v1}, v2={v2}"
            raise ValueError(msg)

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

        defaults = {
            "lr": lr,
            "betas": betas,
            "vs": (v1, v2),
            "eps": eps,
            "weight_decay": weight_decay,
            "decouple": decouple,
            "amsgrad": amsgrad,
            "maximize": maximize,
            "foreach": foreach,
            "capturable": capturable,
            "differentiable": differentiable,
            "fused": fused,
        }
        super().__init__(params, defaults)
        for group in self.param_groups:
            # Store initial_lr for decoupled weight decay calculation if decouple=True
            if group["decouple"]:
                if isinstance(group["lr"], Tensor):
                    # Detach tensor LR to store as initial_lr if needed for tracking
                    group["initial_lr"] = group["lr"].clone().detach()
                else:
                    group["initial_lr"] = float(group["lr"])
            else:
                group["initial_lr"] = None

        if fused:
            msg = "`fused` is not currently supported for AdemamixAdam"
            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."""
        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", True)
            group.setdefault("amsgrad", False)
            fused = group.setdefault("fused", None)

            # Ensure initial_lr is correctly set
            # if missing (e.g., loading older checkpoints)
            if "initial_lr" not in group:
                if group["decouple"]:
                    if isinstance(group["lr"], Tensor):
                        group["initial_lr"] = group["lr"].clone().detach()
                    else:
                        group["initial_lr"] = float(group["lr"])
                else:
                    group["initial_lr"] = None

            for p in group["params"]:
                p_state = self.state.get(p, {})
                if len(p_state) != 0:
                    # Ensure exp_avg_2 exists
                    if "exp_avg_2" not in p_state:
                        p_state["exp_avg_2"] = torch.zeros_like(
                            p,
                            memory_format=torch.preserve_format,
                        )

                    # Ensure max_exp_avg_sq exists if amsgrad is enabled
                    if group["amsgrad"] and "max_exp_avg_sq" not in p_state:
                        p_state["max_exp_avg_sq"] = torch.zeros_like(
                            p,
                            memory_format=torch.preserve_format,
                        )

                    # Handle step tensor conversion if loaded as float
                    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],
        exp_avg_sqs: list[Tensor],
        max_exp_avg_sqs: list[Tensor],
        state_steps: list[Tensor],
    ) -> bool:
        """Initialize a parameter group for AdemamixAdam updates.

        Parameters
        ----------
        group: dict
            The parameter group to initialize.
        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 first moment estimates (M1).
        exp_avg_2s: list[Tensor]
            List to populate with second first moment estimates (M2).
        exp_avg_sqs: list[Tensor]
            List to populate with second moment estimates (V).
        max_exp_avg_sqs: list[Tensor]
            List to populate with max of second moment estimates (for AMSGrad).
        state_steps: list[Tensor]
            List to populate with step counts for each parameter.

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

        Raises
        ------
        RuntimeError
            If any parameter has sparse gradient/if invalid configurations are detected.

        """
        has_complex = False
        amsgrad = group["amsgrad"]

        for p in group["params"]:
            if p.grad is None:
                continue
            has_complex |= torch.is_complex(p)

            if p.grad.is_sparse:
                msg = "AdemamixAdam does not support sparse gradients"
                raise RuntimeError(msg)

            params_with_grad.append(p)
            grads.append(p.grad)

            state = self.state[p]

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

                # Initialize step to 0. Adam increments step at the beginning of the update.
                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())
                )
                # Initialize momentums
                state["exp_avg"] = torch.zeros_like(
                    p,
                    memory_format=torch.preserve_format,
                )
                state["exp_avg_2"] = torch.zeros_like(
                    p,
                    memory_format=torch.preserve_format,
                )
                state["exp_avg_sq"] = torch.zeros_like(
                    p,
                    memory_format=torch.preserve_format,
                )

                if amsgrad:
                    state["max_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"])
            exp_avg_sqs.append(state["exp_avg_sq"])

            if amsgrad:
                max_exp_avg_sqs.append(state["max_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/Tensor LR check
            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 AdemamixAdam optimization step.

        Parameters
        ----------
        closure: Callable[[], torch.Tensor] | None, optional
            A closure that reevaluates the model and returns the loss.
            Optional for most optimizers. Default: None.

        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] = []
            exp_avg_sqs: list[Tensor] = []
            max_exp_avg_sqs: list[Tensor] = []
            state_steps: list[Tensor] = []

            # Extract hyperparameters
            beta1_a, beta1_b, beta2 = cast("tuple[float, float, float]", group["betas"])
            v1, v2 = cast("tuple[float, float]", group["vs"])

            # Handle initial_lr (can be float, Tensor, or None)
            initial_lr = group.get("initial_lr")

            has_complex = self._init_group(
                group,
                params_with_grad,
                grads,
                exp_avgs,
                exp_avg_2s,
                exp_avg_sqs,
                max_exp_avg_sqs,
                state_steps,
            )

            ademamixadam(
                params_with_grad,
                grads,
                exp_avgs,
                exp_avg_2s,
                exp_avg_sqs,
                max_exp_avg_sqs,
                state_steps,
                initial_lr=initial_lr,
                decouple=group["decouple"],
                amsgrad=group["amsgrad"],
                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

    # Metric reduction utilities (Generic L2 reduction)

    def dist_reduce_metrics(self, optimizer_metrics: dict) -> dict:
        """Reduce metrics across workers."""
        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
        all_keys_list = sorted(list(all_keys))

        # Determine device
        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
            if not isinstance(value, Tensor):
                try:
                    value = torch.tensor(float(value), device=device)
                except (ValueError, TypeError):
                    continue

            if metric.startswith("l2_norm"):
                # L2 norms were squared, sum and then sqrt
                dist.all_reduce(value, reduce_operation="SUM")
                optimizer_metrics[metric] = math.sqrt(value.item())
            else:
                # Other metrics are averaged
                dist.all_reduce(value, reduce_operation="SUM")
                optimizer_metrics[metric] = value.item() / dist.get_world_size()

        return optimizer_metrics

    def pre_reduce_metrics(self, optimizer_metrics: dict) -> dict:
        """Preprocess metrics to reduce across ranks correctly."""
        # L2 norms need to be squared
        for metric, value in optimizer_metrics.items():
            if metric.startswith("l2_norm"):
                if isinstance(value, Tensor):
                    optimizer_metrics[metric] = value**2
                else:
                    try:
                        optimizer_metrics[metric] = float(value) ** 2
                    except (ValueError, TypeError):
                        pass
        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 calculates the exact update applied to the parameter
        during the optimization step based on the AdemamixAdam (QHAdam style) formula,
        correctly handling both L2 and AdamW weight decay styles.
        """
        # Retrieve hyperparameters
        group = self.param_groups[0]
        lr = group["lr"]
        weight_decay = group["weight_decay"]
        initial_lr = group.get("initial_lr")
        decouple = group["decouple"]
        amsgrad = group["amsgrad"]
        # Extract betas and vs
        beta1_a, beta1_b, beta2 = cast("tuple[float, float, float]", group["betas"])
        v1, v2 = cast("tuple[float, float]", group["vs"])
        v_grad = 1.0 - v1 - v2
        eps = cast("float", group["eps"])

        if param in self.state:
            param_optim_state = self.state[param]

            # Retrieve raw gradient
            grad_raw = param.grad
            if grad_raw is None:
                return optimizer_metrics

            # Retrieve buffers
            exp_avg = cast("Tensor", param_optim_state["exp_avg"])
            exp_avg_2 = cast("Tensor", param_optim_state["exp_avg_2"])
            exp_avg_sq = cast("Tensor", param_optim_state["exp_avg_sq"])
            step_t = cast("Tensor", param_optim_state["step"])

            # The functional API increments the step at the start.
            step = _get_value(step_t)

            if step < 1:
                return optimizer_metrics

            # 1. Determine the gradient used in the update (G_t)
            if weight_decay != 0 and not decouple:
                # L2 regularization: G_t = G_raw + WD * P
                # We must use the L2-modified gradient for accurate reconstruction.
                grad_used = grad_raw.add(param, alpha=weight_decay)
            else:
                # AdamW or No WD: G_t = G_raw
                grad_used = grad_raw

            # 2. Calculate Bias Corrections
            bias_correction1_a = 1.0 - beta1_a**step
            bias_correction1_b = 1.0 - beta1_b**step
            bias_correction2 = 1.0 - beta2**step

            # 3. Calculate Denominator
            if amsgrad:
                max_exp_avg_sq = cast("Tensor", param_optim_state["max_exp_avg_sq"])
                denom_sq = max_exp_avg_sq
            else:
                denom_sq = exp_avg_sq

            # denom = (sqrt(V_t) / sqrt(bc2)) + eps = sqrt(V_hat_t) + eps
            denom = (denom_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)

            # 4. Calculate Numerator (Ademamix QH Update)
            # N_t = v1 * (M1_t/bc1_a) + v2 * (M2_t/bc1_b) + v_grad * G_t

            # Use the determined gradient (grad_used)
            numerator = grad_used.mul(v_grad)
            numerator.add_(exp_avg, alpha=v1 / bias_correction1_a)
            numerator.add_(exp_avg_2, alpha=v2 / bias_correction1_b)

            # 5. Calculate the gradient-based update (lr * N_t / denom)

            # Handle tensor LR if capturable
            if isinstance(lr, Tensor) and self.defaults["capturable"]:
                step_tensor = torch.div(numerator, denom).mul_(lr)
            else:
                step_tensor = torch.div(numerator, denom).mul_(_get_value(lr))

            # 6. Include decoupled weight decay contribution (if applicable).

            if weight_decay != 0 and decouple:
                # Reconstruction logic for AdamW style update.
                # R = lr/initial_lr (decay factor) and lambda = weight_decay.
                # TotalUpdate = (GradUpdate + P_T*R*lambda) / (1 - R*lambda)

                # Calculate decay factor R
                decay_factor = 1.0
                if initial_lr:
                    if isinstance(lr, Tensor) or isinstance(initial_lr, Tensor):
                        # Assumes safe division (initial_lr != 0)
                        decay_factor = lr / initial_lr
                    elif initial_lr != 0:
                        decay_factor = lr / initial_lr

                R_lambda = decay_factor * weight_decay

                # Ensure R_lambda < 1 for stable reconstruction
                # Efficient calculation based on DecoupledAdamW example:
                # scaling_factor = R_lambda / (1 - R_lambda)
                # step_tensor.mul_(1 + scaling_factor).add_(param, alpha=scaling_factor)

                if isinstance(R_lambda, Tensor):
                    if _get_value(R_lambda) < 1.0:
                        scaling_factor = R_lambda / (1.0 - R_lambda)
                        # Use in-place operations for efficiency
                        step_tensor.mul_(1.0 + scaling_factor).add_(
                            param,
                            alpha=scaling_factor,
                        )
                elif R_lambda < 1.0:
                    scaling_factor = R_lambda / (1.0 - R_lambda)
                    step_tensor.mul_(1.0 + scaling_factor).add_(
                        param,
                        alpha=scaling_factor,
                    )
                # If R_lambda >= 1.0, optimization is unstable.

            # Report metrics
            for metric, func in self.metric_functions.items():
                optimizer_metrics[f"{metric}/{name}"] = func(
                    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_ademamixadam(  # noqa: PLR0917, PLR0913
    params: list[Tensor],
    grads: list[Tensor],
    exp_avgs: list[Tensor],
    exp_avg_2s: list[Tensor],
    exp_avg_sqs: list[Tensor],
    max_exp_avg_sqs: list[Tensor],
    state_steps: list[Tensor],
    grad_scale: Tensor | None,
    found_inf: Tensor | None,
    *,
    initial_lr: float | Tensor | None,
    decouple: bool,
    amsgrad: bool,
    beta1_a: float,
    beta1_b: float,
    beta2: float,
    v1: float,
    v2: float,
    lr: float | Tensor,
    weight_decay: float,
    eps: float,
    maximize: bool,
    capturable: bool,
    differentiable: bool,
    has_complex: bool,
) -> None:
    """Single-tensor implementation of AdemamixAdam step (QHAdam style)."""
    assert grad_scale is None
    assert found_inf is None

    if torch.jit.is_scripting():
        # JIT compatibility checks
        assert isinstance(lr, float)
        if initial_lr is not None:
            assert isinstance(initial_lr, float)

    v_grad = 1.0 - v1 - v2

    for i, param in enumerate(params):
        grad = grads[i]
        exp_avg = exp_avgs[i]
        exp_avg_2 = exp_avg_2s[i]
        exp_avg_sq = exp_avg_sqs[i]
        step_t = state_steps[i]

        # Capturable checks
        if not torch._utils.is_compiling() and capturable:
            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}."
            )

        # Update step
        step_t += 1

        if maximize:
            grad = -grad

        # Weight Decay
        if weight_decay != 0:
            if decouple:
                # Decoupled Weight Decay (AdamW style)
                # Decay factor R = lr / initial_lr
                decay_factor = 1.0
                if initial_lr:
                    # Handle potential Tensor types
                    if capturable and (
                        isinstance(lr, Tensor) or isinstance(initial_lr, Tensor)
                    ):
                        # Assumes initial_lr != 0
                        decay_factor = lr / initial_lr
                    else:
                        # Extract scalar values
                        lr_scalar = _get_value(lr)
                        initial_lr_scalar = _get_value(initial_lr)
                        if initial_lr_scalar != 0:
                            decay_factor = lr_scalar / initial_lr_scalar

                # param.mul_(1 - decay_factor * weight_decay)
                if capturable and isinstance(decay_factor, Tensor):
                    param.mul_(1.0 - decay_factor * weight_decay)
                else:
                    param.mul_(1.0 - _get_value(decay_factor) * weight_decay)

            else:
                # L2 Regularization (Standard Adam style)
                # grad += WD * param
                grad = grad.add(param, alpha=weight_decay)

        # Handle complex parameters
        if has_complex:
            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)
                exp_avg_sq = torch.view_as_real(exp_avg_sq)
                param_view = torch.view_as_real(param)
                if amsgrad:
                    max_exp_avg_sqs[i] = torch.view_as_real(max_exp_avg_sqs[i])
            else:
                param_view = param
        else:
            param_view = param

        # Update Momentums
        # Note: 'grad' might be L2-modified if decouple=False.
        exp_avg.mul_(beta1_a).add_(grad, alpha=1.0 - beta1_a)
        exp_avg_2.mul_(beta1_b).add_(grad, alpha=1.0 - beta1_b)
        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)

        # Calculate Bias Corrections
        step = step_t if capturable or differentiable else _get_value(step_t)

        # NOTE: When capturable=True, calculations must be Tensors.
        if capturable:
            bias_correction1_a = 1.0 - torch.pow(beta1_a, step)
            bias_correction1_b = 1.0 - torch.pow(beta1_b, step)
            bias_correction2 = 1.0 - torch.pow(beta2, step)
            bias_correction2_sqrt = torch.sqrt(bias_correction2)
        else:
            bias_correction1_a = 1.0 - beta1_a**step
            bias_correction1_b = 1.0 - beta1_b**step
            bias_correction2 = 1.0 - beta2**step
            bias_correction2_sqrt = math.sqrt(bias_correction2)

        # Calculate Denominator
        if amsgrad:
            # Maintains the maximum of all 2nd moment running avg. till now
            max_exp_avg_sq = max_exp_avg_sqs[i]
            torch.maximum(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
            # Use the max. for normalization
            denom = (max_exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)
        else:
            denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)

        # Calculate Numerator (Ademamix QHAdam Update)
        # N_t = v1 * (M1_t/bc1_a) + v2 * (M2_t/bc1_b) + v_grad * G_t

        # Start with the gradient component
        # 'grad' here is G_t (potentially L2 modified).
        numerator = grad.mul(v_grad)

        # Add momentum components
        # NOTE: Handle potential Tensor bias corrections for capturable mode
        if capturable:
            alpha1_a = v1 / bias_correction1_a
            alpha1_b = v2 / bias_correction1_b
        else:
            alpha1_a = v1 / bias_correction1_a
            alpha1_b = v2 / bias_correction1_b

        numerator.add_(exp_avg, alpha=alpha1_a)
        numerator.add_(exp_avg_2, alpha=alpha1_b)

        # Final parameter update
        # P_t = P_{t-1} - lr * (N_t / denom)

        # Handle Tensor lr for capturable mode
        if capturable and isinstance(lr, Tensor):
            param_view.addcdiv_(numerator, denom, value=-lr)
        else:
            param_view.addcdiv_(numerator, denom, value=-_get_value(lr))


def _multi_tensor_ademamixadam(  # noqa: PLR0915, PLR0912, PLR0917, PLR0913, C901
    params: list[Tensor],
    grads: list[Tensor],
    exp_avgs: list[Tensor],
    exp_avg_2s: list[Tensor],
    exp_avg_sqs: list[Tensor],
    max_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,
    beta1_b: float,
    beta2: float,
    v1: float,
    v2: float,
    lr: float | Tensor,
    weight_decay: float,
    decouple: bool,
    amsgrad: bool,
    eps: float,
    maximize: bool,
    capturable: bool,
    differentiable: bool,
) -> None:
    """Multi-tensor implementation of the AdemamixAdam step (QHAdam style)."""
    if len(params) == 0:
        return

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

    # Checks for capturable mode
    if not torch._utils.is_compiling() and capturable:
        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 and found_inf is None

    v_grad = 1.0 - v1 - v2

    # Grouping tensors
    tensors_to_group = [params, grads, exp_avgs, exp_avg_2s, exp_avg_sqs, state_steps]
    if amsgrad:
        tensors_to_group.append(max_exp_avg_sqs)

    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(tensors_to_group)

    for device_tuple, _ in grouped_tensors.values():
        # Unpack the tuple
        if amsgrad:
            (
                device_params_,
                device_grads_,
                device_exp_avgs_,
                device_exp_avg_2s_,
                device_exp_avg_sqs_,
                device_state_steps_,
                device_max_exp_avg_sqs_,
            ) = device_tuple
            device_max_exp_avg_sqs = cast("list[Tensor]", device_max_exp_avg_sqs_)
        else:
            (
                device_params_,
                device_grads_,
                device_exp_avgs_,
                device_exp_avg_2s_,
                device_exp_avg_sqs_,
                device_state_steps_,
            ) = device_tuple
            device_max_exp_avg_sqs = []

        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_)
        device_exp_avg_sqs = cast("list[Tensor]", device_exp_avg_sqs_)
        device_state_steps = cast("list[Tensor]", device_state_steps_)

        # Update steps
        # Optimized handling for CPU steps.
        if (
            not torch._utils.is_compiling()
            and device_state_steps[0].is_cpu
            and not capturable
        ):
            torch._foreach_add_(
                device_state_steps,
                torch.tensor(1.0, device="cpu"),
                alpha=1.0,
            )
        else:
            torch._foreach_add_(device_state_steps, 1)

        # Handle complex parameters (view as real)
        if has_complex:
            tensors_to_view = [
                device_params,
                device_grads,
                device_exp_avgs,
                device_exp_avg_2s,
                device_exp_avg_sqs,
            ]
            if amsgrad:
                tensors_to_view.append(device_max_exp_avg_sqs)
            # _view_as_real modifies the lists in-place
            _view_as_real(*tensors_to_view)

        if maximize:
            torch._foreach_neg_(device_grads)

        # Weight Decay
        if weight_decay != 0:
            if decouple:
                # Decoupled Weight Decay (AdamW style)
                # Decay factor R = lr / initial_lr
                decay_factor = 1.0
                if initial_lr:
                    # Handle potential Tensor types
                    if capturable and (
                        isinstance(lr, Tensor) or isinstance(initial_lr, Tensor)
                    ):
                        # Assumes initial_lr != 0
                        decay_factor = lr / initial_lr
                    else:
                        # Extract values
                        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)
                if capturable and isinstance(decay_factor, Tensor):
                    weight_decay_scaled = decay_factor * weight_decay
                    torch._foreach_mul_(device_params, 1.0 - weight_decay_scaled)
                else:
                    weight_decay_scaled = _get_value(decay_factor) * weight_decay
                    torch._foreach_mul_(device_params, 1.0 - weight_decay_scaled)

            else:
                # L2 Regularization (Standard Adam style)
                # grad += WD * param.
                torch._foreach_add_(
                    device_grads,
                    device_params,
                    alpha=weight_decay,
                )

        # Update Momentums
        # Note: device_grads might be L2-modified if decouple=False.
        # We operate directly on device_grads (G_t).
        torch._foreach_mul_(device_exp_avgs, beta1_a)
        torch._foreach_add_(device_exp_avgs, device_grads, alpha=1.0 - beta1_a)

        torch._foreach_mul_(device_exp_avg_2s, beta1_b)
        torch._foreach_add_(device_exp_avg_2s, device_grads, alpha=1.0 - beta1_b)

        torch._foreach_mul_(device_exp_avg_sqs, beta2)
        torch._foreach_addcmul_(
            device_exp_avg_sqs,
            device_grads,
            device_grads,
            value=1.0 - beta2,
        )

        # Calculate Bias Corrections
        # Assuming synchronized steps, use the first one (already incremented).
        step = device_state_steps[0]

        if capturable:
            # Calculate as Tensors
            bias_correction1_a = 1.0 - torch.pow(beta1_a, step)
            bias_correction1_b = 1.0 - torch.pow(beta1_b, step)
            bias_correction2 = 1.0 - torch.pow(beta2, step)
            bias_correction2_sqrt = torch.sqrt(bias_correction2)
        else:
            # Calculate as floats
            step_scalar = _get_value(step)
            bias_correction1_a = 1.0 - beta1_a**step_scalar
            bias_correction1_b = 1.0 - beta1_b**step_scalar
            bias_correction2 = 1.0 - beta2**step_scalar
            bias_correction2_sqrt = math.sqrt(bias_correction2)

        # Calculate Denominator
        if amsgrad:
            # Maintain max
            torch._foreach_maximum_(device_max_exp_avg_sqs, device_exp_avg_sqs)
            exp_avg_sq_sqrt = torch._foreach_sqrt(device_max_exp_avg_sqs)
        else:
            exp_avg_sq_sqrt = torch._foreach_sqrt(device_exp_avg_sqs)

        # denom = (sqrt(V_t) / sqrt(bc2)) + eps
        if capturable:
            torch._foreach_div_(exp_avg_sq_sqrt, bias_correction2_sqrt)
        else:
            torch._foreach_div_(exp_avg_sq_sqrt, bias_correction2_sqrt)

        torch._foreach_add_(exp_avg_sq_sqrt, eps)
        # exp_avg_sq_sqrt now holds the denominators

        # Calculate Numerator (Ademamix QHAdam Update)
        # N_t = v1 * (M1_t/bc1_a) + v2 * (M2_t/bc1_b) + v_grad * G_t

        # Start with the gradient component. We reuse device_grads memory for the numerator.
        # NOTE: This modifies device_grads (G_t) in-place! This is safe as device_grads is not used later.
        torch._foreach_mul_(device_grads, v_grad)

        # Add momentum components
        # Calculate alphas (Handle Tensor/float for capturable mode)
        if capturable:
            alpha1_a = v1 / bias_correction1_a
            alpha1_b = v2 / bias_correction1_b
        else:
            alpha1_a = v1 / bias_correction1_a
            alpha1_b = v2 / bias_correction1_b

        torch._foreach_add_(device_grads, device_exp_avgs, alpha=alpha1_a)
        torch._foreach_add_(device_grads, device_exp_avg_2s, alpha=alpha1_b)
        # device_grads now holds the numerators

        # Final parameter update
        # P_t = P_{t-1} - lr * (N_t / denom)

        # Handle Tensor lr for capturable mode
        if capturable and isinstance(lr, Tensor):
            # Calculate step_size = -lr.
            # Rely on _foreach_addcdiv_ supporting Tensor value.
            step_size = -lr
            torch._foreach_addcdiv_(
                device_params,
                device_grads,
                exp_avg_sq_sqrt,
                value=step_size,
            )
        else:
            step_size = -_get_value(lr)
            torch._foreach_addcdiv_(
                device_params,
                device_grads,
                exp_avg_sq_sqrt,
                value=step_size,
            )


def _fused_ademamixadam(*args, **kwargs) -> None:
    """Fused implementation of the AdemamixAdam step."""
    raise NotImplementedError("Fused implementation of AdemamixAdam is not available.")


# Define the single tensor function reference before the decorator
@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_ademamixadam)
def ademamixadam(  # noqa: PLR0913
    params: list[Tensor],
    grads: list[Tensor],
    exp_avgs: list[Tensor],
    exp_avg_2s: list[Tensor],
    exp_avg_sqs: list[Tensor],
    max_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,
    weight_decay: float,
    decouple: bool,
    amsgrad: bool,
    eps: float,
    maximize: bool,
) -> None:
    r"""Functional API that performs the AdemamixAdam algorithm step."""
    # Logic for selecting implementation
    if fused is None and foreach is None:
        _, foreach = _default_to_fused_or_foreach(
            params,
            differentiable,
            use_fused=False,  # Fused not supported
        )
        # NOTE: Handle unsupported case: lr is 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 or fused) and torch.jit.is_scripting():
        msg = "torch.jit.script not supported with foreach/fused optimizers"
        raise RuntimeError(msg)

    if fused:
        func = _fused_ademamixadam
    elif foreach:
        func = _multi_tensor_ademamixadam
    else:
        func = _single_tensor_ademamixadam

    func(
        params,
        grads,
        exp_avgs,
        exp_avg_2s,
        exp_avg_sqs,
        max_exp_avg_sqs,
        state_steps,
        initial_lr=initial_lr,
        decouple=decouple,
        amsgrad=amsgrad,
        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,
    )
