import math
from collections.abc import Iterable

import torch
import torch.nn.functional as F


class HybridMuonAdamW(torch.optim.Optimizer):
    """AdamW with optional tiled Muon-style updates for selected matrix params."""

    def __init__(
        self,
        params: Iterable,
        lr: float = 1e-3,
        betas: tuple[float, float] = (0.9, 0.95),
        eps: float = 1e-8,
        weight_decay: float = 0.0,
        tile_size: int = 32,
        ns_steps: int = 5,
    ) -> None:
        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, tile_size=tile_size, ns_steps=ns_steps)
        super().__init__(params, defaults)

    @staticmethod
    def _orthogonalize(matrix: torch.Tensor, steps: int, eps: float) -> torch.Tensor:
        original_dtype = matrix.dtype
        x = matrix.float()
        transposed = False
        if x.shape[-2] > x.shape[-1]:
            x = x.mT
            transposed = True
        x = x / x.norm(dim=(-2, -1), keepdim=True).clamp_min(eps)
        for _ in range(steps):
            gram = x @ x.mT
            x = 1.5 * x - 0.5 * gram @ x
        if transposed:
            x = x.mT
        scale = math.sqrt(max(1, matrix.shape[-2]) / max(1, matrix.shape[-1]))
        return (x * scale).to(dtype=original_dtype)

    def _tiled_muon_update(self, update: torch.Tensor, tile_size: int, steps: int, eps: float) -> torch.Tensor:
        if update.ndim < 2:
            return update
        matrices = update.reshape(-1, update.shape[-2], update.shape[-1])
        out = torch.empty_like(matrices)
        for matrix_idx in range(matrices.shape[0]):
            matrix = matrices[matrix_idx]
            result = torch.empty_like(matrix)
            for row_start in range(0, matrix.shape[0], tile_size):
                for col_start in range(0, matrix.shape[1], tile_size):
                    tile = matrix[row_start : row_start + tile_size, col_start : col_start + tile_size]
                    result[row_start : row_start + tile_size, col_start : col_start + tile_size] = self._orthogonalize(
                        tile, steps, eps
                    )
            out[matrix_idx] = result
        return out.reshape_as(update)

    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        for group in self.param_groups:
            lr = group["lr"]
            beta1, beta2 = group["betas"]
            eps = group["eps"]
            weight_decay = group["weight_decay"]
            tile_size = int(group.get("tile_size", 32))
            ns_steps = int(group.get("ns_steps", 5))
            use_muon = bool(group.get("use_muon", False))
            for param in group["params"]:
                if param.grad is None:
                    continue
                grad = param.grad
                if grad.is_sparse:
                    raise RuntimeError("HybridMuonAdamW does not support sparse gradients")
                state = self.state[param]
                if use_muon and grad.ndim >= 2:
                    if len(state) == 0:
                        state["momentum"] = torch.zeros_like(param)
                    momentum = state["momentum"]
                    momentum.mul_(beta1).add_(grad, alpha=1.0 - beta1)
                    if weight_decay != 0.0:
                        param.mul_(1.0 - lr * weight_decay)
                    update = self._tiled_muon_update(momentum, tile_size, ns_steps, eps)
                    param.add_(update, alpha=-lr)
                    continue
                if len(state) == 0:
                    state["step"] = torch.zeros((), dtype=torch.float32, device=param.device)
                    state["exp_avg"] = torch.zeros_like(param)
                    state["exp_avg_sq"] = torch.zeros_like(param)
                exp_avg = state["exp_avg"]
                exp_avg_sq = state["exp_avg_sq"]
                state["step"].add_(1)
                step = state["step"].item()
                if weight_decay != 0.0:
                    param.mul_(1.0 - lr * weight_decay)
                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
                bias_correction1 = 1.0 - beta1**step
                bias_correction2 = 1.0 - beta2**step
                step_size = lr * math.sqrt(bias_correction2) / bias_correction1
                denom = exp_avg_sq.sqrt().add_(eps)
                param.addcdiv_(exp_avg, denom, value=-step_size)
        return loss


class RMNPGrouped(torch.optim.Optimizer):
    """Authors' grouped RMNP/Adam optimizer for LLaMA pretraining.

    The update equations intentionally follow the official RMNP LLaMA code,
    including its Nesterov-style interpolation, last-dimension row
    normalization, rectangular-matrix scaling, and Adam-branch decay using the
    bias-corrected step size. Parameter groups select the branch with
    ``is_rmnp``.

    Source: https://anonymous.4open.science/r/RMNP-E8E1/
    """

    def __init__(
        self,
        params: Iterable,
        lr_rmnp: float = 0.005,
        lr_adam: float = 0.001,
        momentum: float = 0.95,
        beta: float = 0.95,
        weight_decay: float = 0.0,
        adam_weight_decay: float | None = None,
        adamw_decay: bool = False,
        betas: tuple[float, float] = (0.9, 0.95),
        eps: float = 1e-10,
    ) -> None:
        defaults = dict(
            lr=lr_adam,
            lr_rmnp=lr_rmnp,
            lr_adam=lr_adam,
            momentum=momentum,
            beta=beta,
            weight_decay=weight_decay,
            adam_weight_decay=adam_weight_decay,
            adamw_decay=adamw_decay,
            betas=betas,
            eps=eps,
        )
        super().__init__(params, defaults)
        # LitGPT's scheduler needs one reference LR. Individual groups retain
        # their authors-specified LR as schedule_base_lr.
        self.defaults["lr"] = lr_rmnp
        for group in self.param_groups:
            group_lr = lr_rmnp if group.get("is_rmnp", False) else lr_adam
            group["lr"] = group_lr
            group["schedule_base_lr"] = group_lr

    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            is_rmnp = group.get("is_rmnp", False)
            lr = group["lr"]
            momentum = group["momentum"]
            beta = group.get("beta", 0.95)
            weight_decay = group["weight_decay"]
            if not is_rmnp and group.get("adam_weight_decay") is not None:
                weight_decay = group["adam_weight_decay"]

            for param in group["params"]:
                if param.grad is None:
                    continue
                grad = param.grad
                if grad.is_sparse:
                    raise RuntimeError("RMNPGrouped does not support sparse gradients")
                state = self.state[param]

                if is_rmnp and grad.dim() >= 2:
                    if "momentum_buffer" not in state:
                        state["momentum_buffer"] = torch.zeros_like(grad)
                    buf = state["momentum_buffer"]
                    buf.lerp_(grad, 1 - beta)
                    nesterov_buf = grad.lerp(buf, momentum)
                    normed = F.normalize(nesterov_buf, p=2, dim=-1)
                    scale = max(1, math.sqrt(grad.size(-2) / grad.size(-1)))
                    normed = normed * scale
                    if weight_decay > 0:
                        param.data.mul_(1 - lr * weight_decay)
                    param.data.add_(normed, alpha=-lr)
                else:
                    beta1, beta2 = group["betas"]
                    eps = group["eps"]
                    if "step" not in state:
                        state["step"] = 0
                        state["exp_avg"] = torch.zeros_like(grad)
                        state["exp_avg_sq"] = torch.zeros_like(grad)
                    state["step"] += 1
                    exp_avg = state["exp_avg"]
                    exp_avg_sq = state["exp_avg_sq"]
                    exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                    bias_correction1 = 1 - beta1 ** state["step"]
                    bias_correction2 = 1 - beta2 ** state["step"]
                    step_size = lr * math.sqrt(bias_correction2) / bias_correction1
                    denom = exp_avg_sq.sqrt().add_(eps)
                    if weight_decay > 0:
                        decay_step = lr if group.get("adamw_decay", False) else step_size
                        param.data.mul_(1 - decay_step * weight_decay)
                    adam_update = exp_avg / denom
                    param.data.add_(adam_update, alpha=-step_size)
        return loss


class AdaptiveRowTAdamW(torch.optim.Optimizer):
    """AdamW with a magnitude-preserving adaptive row-direction correction.

    ``row_blend=0`` is ordinary AdamW. For selected hidden matrices, the
    Adam-preconditioned direction is blended with the first-moment row
    direction. The alternate direction is rescaled to the Adam direction's
    per-row norm, and its blend is suppressed when the second moment is highly
    anisotropic. This retains TAdam's update scale instead of substituting an
    RMNP-sized step.
    """

    def __init__(
        self,
        params: Iterable,
        lr: float = 1e-3,
        betas: tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-8,
        weight_decay: float = 0.0,
        row_blend: float = 0.1,
        row_blend_warmup_steps: int = 600,
    ) -> None:
        if not 0.0 <= row_blend <= 1.0:
            raise ValueError("row_blend must be between 0 and 1")
        defaults = dict(
            lr=lr,
            betas=betas,
            eps=eps,
            weight_decay=weight_decay,
            row_blend=row_blend,
            row_blend_warmup_steps=row_blend_warmup_steps,
        )
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            lr = group["lr"]
            beta1, beta2 = group["betas"]
            eps = group["eps"]
            weight_decay = group["weight_decay"]
            max_blend = float(group["row_blend"])
            blend_warmup = int(group["row_blend_warmup_steps"])
            use_row_blend = bool(group.get("use_row_blend", False))

            for param in group["params"]:
                if param.grad is None:
                    continue
                grad = param.grad
                if grad.is_sparse:
                    raise RuntimeError("AdaptiveRowTAdamW does not support sparse gradients")
                state = self.state[param]
                if len(state) == 0:
                    state["step"] = 0
                    state["exp_avg"] = torch.zeros_like(param)
                    state["exp_avg_sq"] = torch.zeros_like(param)
                state["step"] += 1
                step = state["step"]
                exp_avg = state["exp_avg"]
                exp_avg_sq = state["exp_avg_sq"]
                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)

                bias_correction1 = 1.0 - beta1**step
                bias_correction2_sqrt = math.sqrt(1.0 - beta2**step)
                denom = exp_avg_sq.sqrt().div_(bias_correction2_sqrt).add_(eps)
                direction = exp_avg / denom / bias_correction1

                if use_row_blend and direction.ndim >= 2 and max_blend > 0.0:
                    row_norm = direction.norm(p=2, dim=-1, keepdim=True)
                    moment_direction = F.normalize(exp_avg, p=2, dim=-1) * row_norm
                    corrected_variance = exp_avg_sq / (1.0 - beta2**step)
                    log_variance = corrected_variance.clamp_min(eps * eps).log()
                    anisotropy = log_variance.std(dim=-1, keepdim=True, correction=0)
                    reliability = anisotropy.add(1.0).reciprocal()
                    ramp = min(1.0, step / max(1, blend_warmup))
                    blend = max_blend * ramp * reliability
                    direction.lerp_(moment_direction, blend)

                if weight_decay != 0.0:
                    param.mul_(1.0 - lr * weight_decay)
                param.add_(direction, alpha=-lr)
        return loss


def _official_muon_zeropower_newton_schulz5(matrix: torch.Tensor, steps: int) -> torch.Tensor:
    """Pinned Keller Jordan Muon f90a42b quintic Newton-Schulz iteration."""
    if matrix.ndim < 2:
        raise ValueError("Muon requires matrix-shaped parameters")
    a, b, c = (3.4445, -4.7750, 2.0315)
    x = matrix.bfloat16()
    transposed = matrix.size(-2) > matrix.size(-1)
    if transposed:
        x = x.mT
    x = x / (x.norm(dim=(-2, -1), keepdim=True) + 1e-7)
    for _ in range(steps):
        gram = x @ x.mT
        polynomial = b * gram + c * gram @ gram
        x = a * x + polynomial @ x
    if transposed:
        x = x.mT
    return x


class OfficialMuonWithAuxAdam(torch.optim.Optimizer):
    """Single-device MuonWithAuxAdam pinned by the RMNP LLaMA experiments."""

    def __init__(
        self,
        params: Iterable,
        lr_muon: float = 0.01,
        lr_adam: float = 0.001,
        momentum: float = 0.95,
        weight_decay: float = 0.1,
        betas: tuple[float, float] = (0.9, 0.95),
        eps: float = 1e-10,
        ns_steps: int = 5,
    ) -> None:
        defaults = dict(
            lr=lr_muon,
            lr_muon=lr_muon,
            lr_adam=lr_adam,
            momentum=momentum,
            weight_decay=weight_decay,
            betas=betas,
            eps=eps,
            ns_steps=ns_steps,
        )
        super().__init__(params, defaults)
        for group in self.param_groups:
            group_lr = lr_muon if group.get("use_muon", False) else lr_adam
            group["lr"] = group_lr
            group["schedule_base_lr"] = group_lr

    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        for group in self.param_groups:
            for param in group["params"]:
                if param.grad is None:
                    continue
                grad = param.grad
                if grad.is_sparse:
                    raise RuntimeError("OfficialMuonWithAuxAdam does not support sparse gradients")
                state = self.state[param]
                if group.get("use_muon", False):
                    if len(state) == 0:
                        state["momentum_buffer"] = torch.zeros_like(param)
                    momentum_buffer = state["momentum_buffer"]
                    momentum = group["momentum"]
                    momentum_buffer.lerp_(grad, 1.0 - momentum)
                    # The pinned implementation intentionally mutates grad here.
                    update = grad.lerp_(momentum_buffer, momentum)
                    head_splits = int(group.get("muon_head_splits", 1))
                    if head_splits > 1:
                        if update.ndim != 2 or update.size(0) % head_splits:
                            raise ValueError(
                                "Per-head Muon requires a 2D projection whose "
                                "output dimension is divisible by muon_head_splits."
                            )
                        head_update = update.view(
                            head_splits,
                            update.size(0) // head_splits,
                            update.size(1),
                        )
                        head_update = _official_muon_zeropower_newton_schulz5(
                            head_update,
                            int(group["ns_steps"]),
                        )
                        head_update *= max(
                            1,
                            head_update.size(-2) / head_update.size(-1),
                        ) ** 0.5
                        update = head_update.reshape_as(update)
                    else:
                        update = _official_muon_zeropower_newton_schulz5(
                            update,
                            int(group["ns_steps"]),
                        )
                        update *= max(1, grad.size(-2) / grad.size(-1)) ** 0.5
                else:
                    if len(state) == 0:
                        state["exp_avg"] = torch.zeros_like(param)
                        state["exp_avg_sq"] = torch.zeros_like(param)
                        state["step"] = 0
                    state["step"] += 1
                    beta1, beta2 = group["betas"]
                    exp_avg = state["exp_avg"]
                    exp_avg_sq = state["exp_avg_sq"]
                    exp_avg.lerp_(grad, 1.0 - beta1)
                    exp_avg_sq.lerp_(grad.square(), 1.0 - beta2)
                    corrected_avg = exp_avg / (1.0 - beta1 ** state["step"])
                    corrected_sq = exp_avg_sq / (1.0 - beta2 ** state["step"])
                    update = corrected_avg / (corrected_sq.sqrt() + group["eps"])
                lr = group["lr"]
                param.mul_(1.0 - lr * group["weight_decay"])
                param.add_(update.reshape(param.shape), alpha=-lr)
        return loss