r"""Implementation of the QHAdamW (Quasi-hyperbolic AdamW) optimizer.

QHAdam Paper: https://arxiv.org/abs/1810.06801
AdamW Paper (Decoupled Weight Decay): https://arxiv.org/abs/1711.05101

Quasi-Hyperbolic Adam (QHAdam) is an optimizer that modifies Adam's update
rule by introducing a quasi-hyperbolic (QH) momentum term. It decouples the
momentum term from the current gradient when performing the update.

The update numerator is a convex combination of the current gradient and the
bias-corrected momentum estimate, controlled by the hyperparameter :math:`\nu_1` (v1).

This implementation assumes the denominator is the standard Adam denominator
(i.e., the second QH parameter v2=1).

This implementation supports both standard L2 regularization (like Adam)
and decoupled weight decay (like AdamW), controlled by the `decouple` parameter.
When decoupled weight decay is used, the decay magnitude follows the learning rate
schedule (scaling by lr/initial_lr) but is not scaled by the learning rate
magnitude itself.

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

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

import torch
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
)

# Imports from the environment context provided in the prompt
try:
    from composer.utils import dist
except ImportError:
    # Fallback if composer is not available
    dist = None

try:
    from llmfoundry.registry import optimizers

    # Define a decorator wrapper if optimizers registry exists
    def register_optimizer(name):
        return optimizers.register_class(name)

except ImportError:
    # Define a dummy decorator if registry doesn't exist
    def register_optimizer(name):
        def decorator(cls):
            return cls

        return decorator


try:
    # Assuming these imports exist based on the prompt context
    from mt_dao.constants import HIGH_WEIGHT_DECAY_WARNING
except ImportError:
    HIGH_WEIGHT_DECAY_WARNING = 0.1  # Default fallback if not available

try:
    from mt_dao.optimizers.utils import get_report_curvature
except ImportError:
    # Fallback if the utility is not available
    def get_report_curvature():
        return None


log = logging.getLogger(__name__)

__all__ = ["QHAdamW", "qhadamw"]


@register_optimizer("qhadamw")
class QHAdamW(Optimizer):
    r"""Implements the QHAdamW algorithm.

    QHAdamW is a variant of Adam that incorporates a quasi-hyperbolic update
    rule, combining the current gradient and the bias-corrected momentum.
    It also supports decoupled weight decay (AdamW style).

    The update rule is given by:

    .. math::
        \text{update}_t = \frac{\nu_1 \hat{m}_t + (1 - \nu_1) g_t}{\sqrt{\hat{v}_t} + \epsilon}

    where :math:`\nu_1` (v1) is the quasi-hyperbolic parameter, and :math:`\hat{m}_t`,
    :math:`\hat{v}_t` are the bias-corrected first and second moments from Adam.

    Argument defaults are aligned with the provided DecoupledAdamW implementation:
    * ``betas`` defaults to ``(0.9, 0.95)``.
    * ``weight_decay`` defaults to ``1e-5``.
    * ``decouple`` defaults to ``True``.

    Notes
    -----
        - This optimizer does not support sparse gradients.
        - Supports the AMSGrad variant.

    """

    metric_functions: ClassVar = {
        "l2_norm/moment": lambda _param,
        optim_state,
        step_tensor: torch.linalg.vector_norm(  # noqa: ARG005
            optim_state["exp_avg"],
        ),
        "l2_norm/moment2": lambda _param,
        optim_state,
        step_tensor: torch.linalg.vector_norm(  # noqa: ARG005
            optim_state["exp_avg_sq"],
        ),
        "min/moment2": (
            lambda param, optim_state, step_tensor: torch.min(  # noqa: ARG005
                optim_state["exp_avg_sq"],
            )
        ),
        "max/moment2": (
            lambda param, optim_state, step_tensor: torch.max(  # noqa: ARG005
                optim_state["exp_avg_sq"],
            )
        ),
        "l2_norm/param": lambda param,
        _optim_state,
        step_tensor: torch.linalg.vector_norm(  # noqa: ARG005
            param.data,
        ),
        "l2_norm/update": lambda _param,
        _optim_state,
        step_tensor: torch.linalg.vector_norm(
            step_tensor,
        ),
        "l2_norm/grad": lambda param,
        _optim_state,
        step_tensor: torch.linalg.vector_norm(  # noqa: ARG005
            param.grad,
        ),
    }

    def __init__(  # noqa: PLR0917, PLR0913, C901
        self,
        params: ParamsT,
        lr: float | Tensor = 1e-3,
        betas: tuple[float, float] = (0.9, 0.95),
        v1: float = 0.7,
        eps: float = 1e-8,
        weight_decay: float = 1e-5,
        *,
        amsgrad: bool = False,
        decouple: bool = True,
        foreach: bool | None = None,
        maximize: bool = False,
        capturable: bool = False,
        differentiable: bool = False,
        fused: bool | None = None,
        report_curvature: bool = False,
    ) -> None:
        """Initialise a QHAdamW optimizer.

        Parameters
        ----------
        params : ParamsT
            Iterable of parameters to optimize or
            dicts defining parameter groups.
        lr : Union[float, Tensor], optional
            Learning rate, by defaut 1e-3.
        betas : tuple[float, float], optional
            Coefficients used for
            computing running averages of gradient and its square,
            by default (0.9, 0.95)
        v1 : float, optional
            Quasi-hyperbolic parameter for the first moment (nu1), by default 0.7
        eps : float, optional
            Term added to the denominator to improve
            numerical stability, by default 1e-8
        weight_decay : float, optional
            Weight decay coefficient, by default 1e-5.
        amsgrad : bool, optional
            Whether to use the AMSGrad variant of this algorithm, by default False
        decouple : bool, optional
            Whether to decouple weight decay (AdamW style)
            or use L2 regularization (Adam style), by default True
        foreach : Optional[bool], optional
            Whether to use the multi-tensor APIs.
        maximize : bool, optional
            Wheter to maximize rather than minimize the loss, by default False
        capturable : bool, optional
            Whether the optimizer should allow gradient capture, by default False
        differentiable : bool, optional
            _description_, by default False
        fused : Optional[bool], optional
            Whether to use a fused version of the kernel, by default None
        report_curvature : bool, optional
            Whether to report curvature metrics, by default False

        """
        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 not 0.0 <= betas[0] < 1.0:
            msg = f"Invalid beta parameter at index 0: {betas[0]}"
            raise ValueError(msg)
        if not 0.0 <= betas[1] < 1.0:
            msg = f"Invalid beta parameter at index 1: {betas[1]}"
            raise ValueError(msg)
        if not weight_decay >= 0.0:
            msg = f"Invalid weight_decay value: {weight_decay}"
            raise ValueError(msg)
        if not 0.0 <= v1 <= 1.0:
            msg = f"Invalid v1 parameter: {v1}. Must be between 0 and 1."
            raise ValueError(msg)

        if decouple and weight_decay >= HIGH_WEIGHT_DECAY_WARNING:
            msg = (
                f"You are using a high value of `weight_decay={weight_decay}` with"
                " decoupled weight decay (AdamW style). Are you sure you want to do this?"
                " Your model's weights will be decayed significantly on every step (scaled by LR schedule)!"
            )
            log.warning(msg)

        defaults = {
            "lr": lr,
            "betas": betas,
            "vs": (v1,),
            "eps": eps,
            "weight_decay": weight_decay,
            "amsgrad": amsgrad,
            "decouple": decouple,
            "maximize": maximize,
            "foreach": foreach,
            "capturable": capturable,
            "differentiable": differentiable,
            "fused": fused,
        }
        super().__init__(params, defaults)
        for group in self.param_groups:
            # Store initial_lr for scheduled decoupled weight decay
            group["initial_lr"] = group["lr"]

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

        # NOTE: Added to avoid expensive metrics calculations
        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("amsgrad", False)
            group.setdefault("maximize", False)
            group.setdefault("foreach", None)
            group.setdefault("capturable", False)
            group.setdefault("differentiable", False)
            group.setdefault("decouple", True)  # Default to True for QHAdamW
            fused = group.setdefault("fused", None)
            for p in group["params"]:
                p_state = self.state.get(p, [])
                if len(p_state) != 0 and not torch.is_tensor(p_state["step"]):  # type: ignore[reportGeneralTypeIssues]
                    step_val = float(p_state["step"])  # type: ignore[reportGeneralTypeIssues]
                    p_state["step"] = (  # type: ignore[reportGeneralTypeIssues]
                        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_sqs: list[Tensor],
        max_exp_avg_sqs: list[Tensor],  # Added for AMSGrad support
        state_steps: list[Tensor],
    ) -> bool:
        """Initialize a parameter group for QHAdamW updates."""
        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 = "QHAdamW does not support sparse gradients"
                raise RuntimeError
            grads.append(p.grad)

            state = self.state[p]

            # Lazy state initialization
            if len(state) == 0:
                if group["fused"]:
                    _device_dtype_check_for_fused(p)
                # note(crcrpar): Deliberately host `step` on CPU if both
                # capturable and fused are off.
                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
                state["exp_avg"] = 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,
                )
                # Maintains max of all exp. moving avg. of sq. grad. values (for AMSGrad)
                if group["amsgrad"]:
                    state["max_exp_avg_sq"] = torch.zeros_like(
                        p,
                        memory_format=torch.preserve_format,
                    )

            exp_avgs.append(state["exp_avg"])
            exp_avg_sqs.append(state["exp_avg_sq"])

            if group["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 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(  # type: ignore # noqa: PGH003
        self,
        closure: Callable[[], torch.Tensor] | None = None,
    ) -> Tensor | None:
        """Perform a single QHAdamW optimization step."""
        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_sqs: list[Tensor] = []
            max_exp_avg_sqs: list[Tensor] = []
            state_steps: list[Tensor] = []
            beta1, beta2 = cast("tuple[float, float]", group["betas"])
            # Quasi-hyperbolic params
            v1, *_ = cast("tuple[float,...]", group["vs"])
            amsgrad = group["amsgrad"]

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

            qhadamw(
                params_with_grad,
                grads,
                exp_avgs,
                exp_avg_sqs,
                max_exp_avg_sqs,
                state_steps,
                initial_lr=group["initial_lr"],
                decouple=group["decouple"],
                amsgrad=amsgrad,
                beta1=beta1,
                beta2=beta2,
                v1=v1,
                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."""
        if dist is None or 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 = sorted(all_keys)  # type: ignore[assignment]

        # Check if running on GPU for tensor creation device
        device = torch.cuda.current_device() if torch.cuda.is_available() else "cpu"

        for metric in all_keys:
            if metric.startswith("l2_norm"):
                reduced = optimizer_metrics.get(
                    metric,
                    torch.tensor(0.0, device=device),
                )
                # Ensure reduced is a tensor before reduction
                if not isinstance(reduced, Tensor):
                    reduced = torch.tensor(reduced, device=device)

                dist.all_reduce(reduced, reduce_operation="SUM")
                # We keep the result as a Tensor if possible, matching QHADOPT style
                optimizer_metrics[metric] = torch.sqrt(reduced)
            else:
                reduced = optimizer_metrics.get(
                    metric,
                    torch.tensor(0.0, device=device),
                )
                if not isinstance(reduced, Tensor):
                    reduced = torch.tensor(reduced, device=device)

                dist.all_reduce(reduced, reduce_operation="SUM")
                optimizer_metrics[metric] = reduced / 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."""
        # Only L2 norm metric keys are present
        for metric in optimizer_metrics:
            if metric.startswith("l2_norm"):
                # L2 norms need to be squared, before they are reduced via summation
                optimizer_metrics[metric] = optimizer_metrics[metric] ** 2

        return optimizer_metrics

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

        This reconstructs the update step applied to the parameter,
        including both the optimization step and the decoupled weight decay step.
        This follows the inversion logic used in the reference DecoupledAdamW.
        """
        # Retrieve hyperparameters (assuming homogeneity across groups)
        group = self.param_groups[0]
        lr = group["lr"]
        weight_decay = group["weight_decay"]
        initial_lr = group["initial_lr"]
        decouple = group["decouple"]
        amsgrad = group["amsgrad"]
        maximize = group["maximize"]
        beta1, beta2 = cast("tuple[float, float]", group["betas"])
        v1, *_ = cast("tuple[float,...]", group["vs"])
        eps = cast("float", group["eps"])

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

            # Retrieve needed buffers from the optimizer's state
            grad = param.grad

            if grad is None:
                # Nothing to report
                return optimizer_metrics

            exp_avg = cast("Tensor", param_optim_state["exp_avg"])
            exp_avg_sq = cast("Tensor", param_optim_state["exp_avg_sq"])
            step_t = cast("Tensor", param_optim_state["step"])
            # We assume step() has already been called, so step_t is the current step count.
            step = _get_value(step_t)

            if step == 0:
                return optimizer_metrics

            # 1) Bias correction (scalar math)
            step_val = float(step)
            bias_correction1 = 1.0 - beta1**step_val
            bias_correction2 = 1.0 - beta2**step_val
            bc2_sqrt = math.sqrt(bias_correction2)

            # 2) Recompute the denominator
            if amsgrad:
                max_exp_avg_sq = cast("Tensor", param_optim_state["max_exp_avg_sq"])
                denom = (max_exp_avg_sq.sqrt() / bc2_sqrt).add_(eps)
            else:
                denom = (exp_avg_sq.sqrt() / bc2_sqrt).add_(eps)

            # 3) Calculate m_hat
            m_hat = exp_avg / bias_correction1

            # 4) Reconstruct the gradient used in the update (g_t)
            # Account for maximization
            current_grad = grad if not maximize else -grad

            # Account for L2 regularization (if not decoupled)
            # If L2 regularization was used, the gradient used for momentum update was (g_t + wd * theta_{t-1}).
            # We approximate theta_{t-1} with theta_t (param) for this reconstruction.
            if weight_decay != 0 and not decouple:
                # Standard Adam L2 regularization is typically unscheduled.
                current_grad = current_grad.add(param, alpha=weight_decay)

            # 5) Calculate QH numerator: (1-v1) * g_t + v1 * m_hat_t
            # Use lerp for numerical stability: g_t.lerp(m_hat_t, v1)
            qh_numerator = current_grad.lerp(m_hat, v1)

            # 6) Calculate the optimization step tensor
            # lr might be a tensor.
            step_tensor = lr * qh_numerator / denom

            # 7) Account for decoupled weight decay contribution to the total update
            # This logic follows the inversion method used in the reference DecoupledAdamW
            if weight_decay != 0 and decouple:
                # initial_lr might be a tensor or 0.
                if (
                    initial_lr is None
                    or (isinstance(initial_lr, Tensor) and (initial_lr == 0).any())
                    or initial_lr == 0.0
                ):
                    decay_factor = 1.0
                else:
                    decay_factor = lr / initial_lr

                W = decay_factor * weight_decay

                # Check stability (W < 1). W might be a tensor.
                is_stable = (W < 1.0).all() if torch.is_tensor(W) else W < 1.0

                if is_stable:
                    # scaling_factor = W / (1 - W)
                    scaling_factor = W / (1.0 - W)
                    # Total update magnitude approx = update_t * (1+S) + param_t * S
                    step_tensor.mul_(1.0 + scaling_factor).add_(
                        param, alpha=scaling_factor,
                    )
                # If W >= 1, reporting might be unstable.

            for metric in self.metric_functions:
                optimizer_metrics[f"{metric}/{name}"] = self.metric_functions[metric](
                    param,
                    param_optim_state,
                    step_tensor,
                )
            # NOTE: these are heavy and require extra memory
            if self.curvature_metric_function is not None:
                optimizer_metrics.update(self.curvature_metric_function(param, name))

        return optimizer_metrics


def _single_tensor_qhadamw(  # noqa: PLR0917, PLR0913, C901
    params: list[Tensor],
    grads: list[Tensor],
    exp_avgs: 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 | None,
    decouple: bool,
    amsgrad: bool,
    beta1: float,
    beta2: float,
    v1: float,
    lr: float | Tensor,
    weight_decay: float,
    eps: float,
    maximize: bool,
    capturable: bool,
    differentiable: bool,
    has_complex: bool,  # noqa: ARG001
) -> None:
    """Single-tensor implementation of QHAdamW step."""
    assert grad_scale is None
    assert found_inf is None

    if torch.jit.is_scripting():  # type: ignore # noqa: PGH003
        # JIT workaround
        assert isinstance(lr, float)

    for i, param in enumerate(params):
        grad = grads[i] if not maximize else -grads[i]
        exp_avg = exp_avgs[i]
        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:  # type: ignore[reportGeneralTypeIssues] # 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}."
            )

        # update step (Adam requires incrementing step before bias correction)
        step_t += 1
        step = step_t if capturable or differentiable else _get_value(step_t)

        # Weight decay (L2 style - Standard Adam)
        if weight_decay != 0 and not decouple:
            # L2 regularization (applied to gradient, typically unscheduled)
            grad = grad.add(param, alpha=weight_decay)

        # Handle Complex Numbers
        if torch.is_complex(param):
            grad = torch.view_as_real(grad)
            exp_avg = torch.view_as_real(exp_avg)
            exp_avg_sq = torch.view_as_real(exp_avg_sq)
            param_data = torch.view_as_real(param)
            if amsgrad:
                # Ensure max_exp_avg_sqs view is also updated if complex
                max_exp_avg_sqs[i] = torch.view_as_real(max_exp_avg_sqs[i])
        else:
            param_data = param

        # Weight decay (Decoupled style - AdamW)
        if weight_decay != 0 and decouple:
            # Decoupled WD scaled by LR schedule ratio (DecoupledAdamW style)
            # lr and initial_lr might be Tensors.

            # Handle potential division by zero or None initial_lr
            if (
                initial_lr is None
                or (isinstance(initial_lr, Tensor) and (initial_lr == 0).any())
                or initial_lr == 0.0
            ):
                decay_factor = 1.0
            else:
                decay_factor = lr / initial_lr

            param_data.mul_(1.0 - decay_factor * weight_decay)

        # Momentum updates (Adam style)
        # m_t
        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
        # v_t
        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

        # Bias correction
        # We use the ** operator which handles both scalar (float conversion needed) and tensor steps.
        if capturable or differentiable:
            bias_correction1 = 1.0 - beta1**step
            bias_correction2 = 1.0 - beta2**step
        else:
            step_val = float(step)
            bias_correction1 = 1.0 - beta1**step_val
            bias_correction2 = 1.0 - beta2**step_val

        # Denominator calculation (sqrt(v_hat_t) + eps)

        # Determine sqrt implementation for bias_correction2 based on mode
        if capturable or differentiable:
            bc2_sqrt = torch.sqrt(bias_correction2)
        else:
            bc2_sqrt = math.sqrt(bias_correction2)

        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 normalizing running avg. of gradient
            denom = (max_exp_avg_sq.sqrt() / bc2_sqrt).add_(eps)
        else:
            denom = (exp_avg_sq.sqrt() / bc2_sqrt).add_(eps)

        # QHAdamW Update Calculation
        # Numerator = (1-v1) * grad + v1 * m_hat_t
        # m_hat_t = exp_avg / bias_correction1

        # Calculate m_hat
        m_hat = exp_avg / bias_correction1

        # Calculate QH numerator using lerp for stability: grad.lerp(m_hat, v1)
        # If differentiable=True, we must ensure operations support autograd. lerp supports it.
        qh_numerator = grad.lerp(m_hat, v1)

        # Final update: param -= lr * qh_numerator / denom
        # lr might be a tensor.
        param_data.addcdiv_(qh_numerator, denom, value=-lr)


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

    # Foreach requires scalar LR if not capturable.
    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:  # type: ignore[reportGeneralTypeIssues] # 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

    # Prepare tensor lists for grouping, handling AMSGrad buffers
    tensor_lists = [params, grads, exp_avgs, exp_avg_sqs, state_steps]
    if amsgrad:
        tensor_lists.insert(4, max_exp_avg_sqs)

    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(  # noqa: SLF001
        tensor_lists,  # type: ignore[reportArgumentType]
    )

    for device_tuple, _ in grouped_tensors.values():
        # Unpack tensors
        if amsgrad:
            (
                device_params_,
                device_grads_,
                device_exp_avgs_,
                device_exp_avg_sqs_,
                device_max_exp_avg_sqs_,
                device_state_steps_,
            ) = 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_sqs_,
                device_state_steps_,
            ) = device_tuple
            device_max_exp_avg_sqs = []  # Dummy list

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

        # Update steps
        # If steps are on CPU, optimize the addition.
        if not torch._utils.is_compiling() and device_state_steps[0].is_cpu:  # type: ignore[reportGeneralTypeIssues] # 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)

        # Handle complex parameters
        if has_complex:
            args_to_view = [
                device_params,
                device_grads,
                device_exp_avgs,
                device_exp_avg_sqs,
            ]
            if amsgrad:
                args_to_view.append(device_max_exp_avg_sqs)
            _view_as_real(*args_to_view)

        if maximize:
            device_grads = torch._foreach_neg(device_grads)  # type: ignore[assignment] # noqa: SLF001

        # Weight decay (L2 style - Standard Adam)
        if weight_decay != 0 and not decouple:
            # L2 regularization (applied to gradient, typically unscheduled)
            # grad += weight_decay * param
            # Re-use memory if maximize=True, otherwise allocate new list.
            if maximize:
                torch._foreach_add_(
                    device_grads,
                    device_params,
                    alpha=weight_decay,
                )
            else:
                device_grads = torch._foreach_add(  # type: ignore[assignment] # noqa: SLF001
                    device_grads,
                    device_params,
                    alpha=weight_decay,
                )

        # Weight decay (Decoupled style - AdamW)
        if weight_decay != 0 and decouple:
            # Decoupled WD scaled by LR schedule ratio

            # Handle potential division by zero or None initial_lr
            if (
                initial_lr is None
                or (isinstance(initial_lr, Tensor) and (initial_lr == 0).any())
                or initial_lr == 0.0
            ):
                decay_factor = 1.0
            else:
                decay_factor = lr / initial_lr

            # If capturable=True, decay_factor might be a Tensor.
            if capturable:
                weight_decay_term = decay_factor * weight_decay
                torch._foreach_mul_(device_params, 1.0 - weight_decay_term)
            else:
                # If not capturable, we need scalar values.
                weight_decay_term = _get_value(decay_factor) * weight_decay
                torch._foreach_mul_(device_params, 1.0 - weight_decay_term)

        # Momentum updates (Adam style)
        # m_t
        torch._foreach_mul_(device_exp_avgs, beta1)
        torch._foreach_add_(device_exp_avgs, device_grads, alpha=1 - beta1)
        # v_t
        torch._foreach_mul_(device_exp_avg_sqs, beta2)
        torch._foreach_addcmul_(
            device_exp_avg_sqs,
            device_grads,
            device_grads,
            value=1 - beta2,
        )

        # Bias correction
        # We need the step value. Assumes uniform steps across the group.
        step = device_state_steps[0]

        if capturable:
            bias_correction1 = 1.0 - torch.pow(beta1, step)
            bias_correction2 = 1.0 - torch.pow(beta2, step)
            bc2_sqrt = torch.sqrt(bias_correction2)
        else:
            step_val = float(_get_value(step))
            bias_correction1 = 1.0 - beta1**step_val
            bias_correction2 = 1.0 - beta2**step_val
            bc2_sqrt = math.sqrt(bias_correction2)

        # Denominator calculation
        if amsgrad:
            # Maintains the maximum of all 2nd moment running avg. till now
            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) / bc2_sqrt) + eps
        torch._foreach_div_(exp_avg_sq_sqrt, bc2_sqrt)
        torch._foreach_add_(exp_avg_sq_sqrt, eps)
        # exp_avg_sq_sqrt now holds the denominators (denom)

        # QHAdamW Update Calculation

        # Calculate m_hat = exp_avg / bias_correction1
        m_hats = torch._foreach_div(device_exp_avgs, bias_correction1)

        # Calculate QH numerator using lerp: grad.lerp(m_hat, v1)
        qh_numerators = torch._foreach_lerp(device_grads, m_hats, v1)

        # Final update: param -= lr * qh_numerator / denom
        # denom is stored in exp_avg_sq_sqrt
        # If capturable=True, lr can be a tensor for _foreach_addcdiv_.
        if capturable:
            torch._foreach_addcdiv_(
                device_params, qh_numerators, exp_avg_sq_sqrt, value=-lr,
            )
        else:
            torch._foreach_addcdiv_(
                device_params, qh_numerators, exp_avg_sq_sqrt, value=-_get_value(lr),
            )


def _fused_qhadamw(
    *args, **kwargs,
) -> None:
    """Fused implementation of the QHAdamW step."""
    raise NotImplementedError("Fused QHAdamW is not implemented.")


@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_qhadamw)
def qhadamw(  # noqa: PLR0913
    params: list[Tensor],
    grads: list[Tensor],
    exp_avgs: list[Tensor],
    exp_avg_sqs: list[Tensor],
    max_exp_avg_sqs: list[Tensor],
    state_steps: list[Tensor],
    # kwonly args with defaults are not supported by functions compiled with
    # torchscript issue #70627
    *,
    initial_lr: float | 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,
    amsgrad: bool,
    beta1: float,
    beta2: float,
    v1: float,
    lr: float | Tensor,
    weight_decay: float,
    decouple: bool,
    eps: float,
    maximize: bool,
) -> None:
    r"""Functional API that performs the QHAdamW algorithm step.

    It directly applies the parameter updates in-place.
    """
    # Determine implementation (fused, foreach, or single-tensor).
    if fused is None and foreach is None:
        _, foreach = _default_to_fused_or_foreach(
            params,
            differentiable,
            use_fused=False,
        )
        # 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():  # type: ignore[PGH003]
        msg = "torch.jit.script not supported with foreach optimizers"
        raise RuntimeError(msg)
    if fused and torch.jit.is_scripting():  # type: ignore[PGH003]
        msg = "torch.jit.script not supported with fused optimizers"
        raise RuntimeError(msg)

    if fused and not torch.jit.is_scripting():  # type: ignore[PGH003]
        func = _fused_qhadamw
    elif foreach and not torch.jit.is_scripting():  # type: ignore[PGH003]
        func = _multi_tensor_qhadamw
    else:
        func = _single_tensor_qhadamw

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