"""Muon: momentum, then orthogonalise the update of every weight matrix.

For a matrix W, plain SGD adds -lr * G. Muon replaces the momentum-averaged G
by the closest orthogonal matrix U·Vᵀ (from G = U·S·Vᵀ): every direction gets
the same step size, so rare but useful directions are not drowned out by a
few dominant ones. Five Newton–Schulz iterations approximate U·Vᵀ cheaply.

Only for 2-D hidden weights. Embeddings and norm gains stay with AdamW.
After Keller Jordan, https://kellerjordan.github.io/posts/muon/
"""
from collections import defaultdict

import torch


def orthogonalize(G, steps=5):
    """G: (..., rows, cols) -> the nearest semi-orthogonal matrices, batched."""
    a, b, c = 3.4445, -4.7750, 2.0315       # tuned for fast convergence
    X = G.bfloat16()
    tall = X.size(-2) > X.size(-1)
    if tall:
        X = X.mT
    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
    for _ in range(steps):
        A = X @ X.mT
        X = a * X + (b * A + c * A @ A) @ X
    if tall:
        X = X.mT
    return X.to(G.dtype)


class Muon(torch.optim.Optimizer):
    def __init__(self, params, lr=0.02, momentum=0.95, weight_decay=0.0, split=None):
        # split: {param: n} treats a fused weight (like q, k, v) as n separate matrices
        self.split = split or {}
        super().__init__(params, dict(lr=lr, momentum=momentum, weight_decay=weight_decay))

    @torch.no_grad()
    def step(self):
        for group in self.param_groups:
            # matrices of the same shape are orthogonalised together, in one batch
            by_shape = defaultdict(list)
            for p in group['params']:
                if p.grad is None:
                    continue
                state = self.state[p]
                if 'momentum' not in state:
                    state['momentum'] = torch.zeros_like(p)
                buf = state['momentum']
                buf.lerp_(p.grad, 1 - group['momentum'])
                g = p.grad.lerp(buf, group['momentum'])       # Nesterov look-ahead
                n = self.split.get(p, 1)
                by_shape[(p.size(0) // n, p.size(1))].append((p, g.view(n, -1, p.size(1))))
            for (rows, cols), items in by_shape.items():
                updates = orthogonalize(torch.cat([g for _, g in items]))
                scale = max(1.0, rows / cols) ** 0.5
                i = 0
                for p, g in items:
                    u = updates[i:i + g.size(0)].reshape_as(p)
                    i += g.size(0)
                    if group['weight_decay']:
                        p.mul_(1 - group['lr'] * group['weight_decay'])
                    p.add_(u, alpha=-group['lr'] * scale)