"""
Muon Optimizer — Keller Jordan, NanoGPT speedrun.

Newton-Schulz iterasyonu ile orthogonalize edilmiş momentum.
2D ağırlıklara (Linear weight) uygulanır. 1D parametreler (norm weight,
bias, embedding) AdamW'ye verilir.

Referans:
  https://github.com/KellerJordan/modded-nanogpt
  https://kellerjordan.github.io/posts/muon/

Kullanim:
    # Param ayri:
    muon_params = [p for p in model.parameters() if p.dim() >= 2 and p.requires_grad]
    other_params = [p for p in model.parameters() if p.dim() < 2 and p.requires_grad]

    # Embedding ve lm_head'i muon'dan ayir (yaygin best practice)
    embed_params = [model.wte.weight]  # tied ise lm_head dahil
    muon_params = [p for p in muon_params if not any(p is e for e in embed_params)]
    other_params = other_params + embed_params

    optimizer_muon = Muon(muon_params, lr=2e-2, momentum=0.95)
    optimizer_adam = torch.optim.AdamW(other_params, lr=3e-4, ...)
"""

import torch


@torch.no_grad()
def newton_schulz(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """G matrisini orthogonalize et (yaklasik USV^T -> UV^T).

    Newton-Schulz quintic iteration. bf16'da kararli, hizli.
    """
    assert G.ndim == 2
    a, b, c = (3.4445, -4.7750, 2.0315)
    X = G.to(torch.bfloat16)
    # Boyut yonune gore transpose (her iki yonde de calissin)
    if X.size(0) > X.size(1):
        X = X.T
    # Spektral normu yaklasik 1'e cek
    X = X / (X.norm() + 1e-7)
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * (A @ A)
        X = a * X + B @ X
    if G.size(0) > G.size(1):
        X = X.T
    return X.to(G.dtype)


class Muon(torch.optim.Optimizer):
    """Muon: Momentum + orthogonalize edilmiş update.

    Sadece 2D parametreler için. 1D'leri AdamW ile ayrı eğit.

    Args:
        params: 2D parametreler iterable
        lr: 0.02 (AdamW'nin ~50x'i, çünkü update'ler ortonormal)
        momentum: 0.95
        nesterov: True (genelde daha iyi)
        ns_steps: Newton-Schulz iter sayisi (5 default)
    """
    def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5):
        defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_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"]
            momentum = group["momentum"]
            nesterov = group["nesterov"]
            ns_steps = group["ns_steps"]

            for p in group["params"]:
                if p.grad is None:
                    continue
                if p.ndim < 2:
                    raise ValueError(
                        f"Muon sadece >=2D param destekler, {p.ndim}D bulundu. "
                        "1D paramları AdamW'ye ver.")

                g = p.grad
                state = self.state[p]
                if "momentum_buffer" not in state:
                    state["momentum_buffer"] = torch.zeros_like(g)

                buf = state["momentum_buffer"]
                buf.mul_(momentum).add_(g)

                # Nesterov momentum
                if nesterov:
                    g = g.add(buf, alpha=momentum)
                else:
                    g = buf

                # Reshape if needed (e.g., conv weight)
                original_shape = g.shape
                if g.ndim > 2:
                    g = g.view(g.size(0), -1)

                # Newton-Schulz orthogonalization
                g_orth = newton_schulz(g, steps=ns_steps)

                # Scale: sqrt(max(out, in) / min(out, in)) — ~spectral norm
                # Ya da basitce sqrt(d_out / d_in) gibi.
                # Modded-nanogpt: scale = max(1, p.shape[0]/p.shape[1]) ** 0.5
                scale = max(1.0, g_orth.size(0) / g_orth.size(1)) ** 0.5

                # Geri reshape
                g_orth = g_orth.view(original_shape)

                p.add_(g_orth, alpha=-lr * scale)

        return loss