"""Muon (Newton-Schulz orthogonalised momentum) for 2-D hidden matrices (and 3-D stacks of MoE expert matrices); scale matched to AdamW RMS (Moonlight)."""
import torch


@torch.no_grad()
def newton_schulz(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G.bfloat16()
    transposed = X.shape[0] > X.shape[1]
    if transposed:
        X = X.T
    X = X / (X.norm() + 1e-7)
    for _ in range(steps):
        A = X @ X.T
        X = a * X + (b * A + c * (A @ A)) @ X
    return X.T if transposed else X


class Muon(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-3, momentum=0.95, weight_decay=0.1, split=None):
        # split: dict param-id -> number of equal row-chunks to orthogonalise independently (fused gate/up)
        super().__init__(params, dict(lr=lr, momentum=momentum, weight_decay=weight_decay))
        self.split = split or {}
        self.bad = None   # device counter of updates dropped for being non-finite (read at log time, no sync here)

    @torch.no_grad()
    def step(self):
        for g in self.param_groups:
            for p in g["params"]:
                if p.grad is None:
                    continue
                st = self.state[p]
                if "buf" not in st:
                    st["buf"] = torch.zeros_like(p.grad)
                buf = st["buf"]
                buf.lerp_(p.grad, 1 - g["momentum"])
                upd = p.grad.lerp(buf, g["momentum"])  # nesterov
                # 3-D = stacked MoE experts: each expert's matrix is orthogonalised on its own
                mats = list(upd.unbind(0)) if upd.ndim == 3 else [upd]
                done = []
                for m in mats:
                    outs = []
                    for ch in m.chunk(self.split.get(id(p), 1), dim=0):
                        o = newton_schulz(ch)
                        outs.append(o * (0.2 * max(ch.shape) ** 0.5))
                    done.append(torch.cat(outs, 0) if len(outs) > 1 else outs[0])
                o = torch.stack(done) if upd.ndim == 3 else done[0]
                # a rare non-finite Newton-Schulz result (seen on the RX 7600 with finite grads) would poison the
                # weight for good: drop this parameter's update for the step instead
                ok = torch.isfinite(o).all()
                if self.bad is None: self.bad = torch.zeros((), device=o.device, dtype=torch.int64)
                self.bad += (~ok).long()
                o = torch.where(ok, o, torch.zeros((), device=o.device, dtype=o.dtype))
                p.mul_(1 - g["lr"] * g["weight_decay"])
                p.add_(o.to(p.dtype), alpha=-g["lr"])