"""
Single-device Muon + Aux AdamW (for NoCap kit).

Attribution: Muon algorithm by Keller Jordan et al.
  https://kellerjordan.github.io/posts/muon/
  https://github.com/KellerJordan/Muon

Embedded (not pip) so Modal kit tarball / local tree can ship it without
extra deps. Use only for hidden 2D weights; AdamW for embed/head/1D.
Competition writeup must attribute Modded-NanoGPT adjacency.
"""
from __future__ import annotations

import torch


# zeropower ≈ Ortho(G) = UVᵀ (drop singular values; Keller blog: newtonschulz5)
def zeropower_via_newtonschulz5(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    assert G.ndim >= 2
    a, b, c = (3.4445, -4.7750, 2.0315)  # tuned NS polynomial coeffs (same as blog)
    X = G.bfloat16()
    # Tall matrix: transpose so the cheap matmul is X @ X.mT (min side squared).
    if G.size(-2) > G.size(-1):
        X = X.mT
    # Frobenius normalize → singular values land in [0, 1] (batched over last two dims).
    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
    for _ in range(steps):
        A = X @ X.mT
        B = b * A + c * A @ A  # degree-5 polynomial in A via matmuls
        X = a * X + B @ X      # φ(X) = aX + (bA + cA²)X
    if G.size(-2) > G.size(-1):
        X = X.mT  # undo tall-matrix transpose
    return X


# Momentum first, then Newton–Schulz (Muon order). Default: Nesterov lookahead.
def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True):
    momentum.lerp_(grad, 1 - beta)
    update = grad.lerp_(momentum, beta) if nesterov else momentum  # Nesterov vs plain momentum
    if update.ndim == 4:
        update = update.view(len(update), -1)  # 4D conv weights: flatten spatial dims
    update = zeropower_via_newtonschulz5(update, steps=ns_steps)
    update *= max(1, update.size(-2) / update.size(-1)) ** 0.5  # aspect-ratio scale max(1, n/m)^0.5
    return update


def muon_update_split_qkv(
    grad,
    momentum,
    beta=0.95,
    ns_steps=5,
    nesterov=True,
    match_joint_frobenius=True,
):
    """Apply the frozen Muon polar map independently to fused Q, K and V.

    ``nn.Linear(d, 3*d, bias=False)`` stores its weight and gradient as
    ``(3*d, d)``. Momentum remains one fused tensor and is updated exactly as
    in :func:`muon_update`; only the post-momentum polar map is separated into
    three square blocks.

    Exact polar factors have equal joint/split Frobenius scale. The frozen
    five-step Newton--Schulz map is approximate, so its realized scale can
    differ across the two spectra.

    * ``match_joint_frobenius=True`` (dual-polar isolation path): recompute the
      historical joint NS5 direction and rescale the split update to its
      realized Frobenius norm. Extra joint polar work stays on the clock.
    * ``match_joint_frobenius=False`` (timing-cheap path): apply only the three
      block NS5 maps with **no** second polar and **no** extra closed-form
      scale. Empirical NS5 block norms already land near the joint+sqrt(3)
      scale; multiplying by sqrt(3) again over-scales (~1.8x) and is not used.
    """
    if grad.ndim != 2 or grad.size(0) != 3 * grad.size(1):
        raise ValueError(
            "split-QKV Muon requires a fused (3*d, d) matrix, got %s"
            % (tuple(grad.shape),)
        )
    if momentum.shape != grad.shape:
        raise ValueError("split-QKV momentum and gradient shapes must match")
    momentum.lerp_(grad, 1 - beta)
    update = grad.lerp_(momentum, beta) if nesterov else momentum
    d = update.size(1)
    split = zeropower_via_newtonschulz5(
        update.reshape(3, d, d), steps=ns_steps
    ).reshape(3 * d, d)
    if match_joint_frobenius:
        joint = zeropower_via_newtonschulz5(update, steps=ns_steps)
        joint *= 3.0**0.5
        split_norm = split.float().norm()
        split *= joint.float().norm() / split_norm.clamp_min(1e-12)
    return split

def adam_update(grad, buf1, buf2, step, betas, eps):
    buf1.lerp_(grad, 1 - betas[0])
    buf2.lerp_(grad.square(), 1 - betas[1])
    buf1c = buf1 / (1 - betas[0] ** step)
    buf2c = buf2 / (1 - betas[1] ** step)
    return buf1c / (buf2c.sqrt() + eps)


class SingleDeviceMuonWithAuxAdam(torch.optim.Optimizer):
    """Non-distributed Muon + AdamW hybrid. Each param group needs use_muon bool."""

    def __init__(
        self,
        param_groups,
        *,
        split_qkv_params=(),
        match_joint_frobenius=True,
    ):
        if type(match_joint_frobenius) is not bool:
            raise ValueError("match_joint_frobenius must be a bool")
        for group in param_groups:
            assert "use_muon" in group
            if group["use_muon"]:
                group["lr"] = group.get("lr", 0.02)
                group["momentum"] = group.get("momentum", 0.95)
                group["weight_decay"] = group.get("weight_decay", 0)
                group["ns_steps"] = group.get("ns_steps", 5)
                assert set(group.keys()) == {
                    "params", "lr", "momentum", "weight_decay", "ns_steps",
                    "use_muon"
                }
            else:
                group["lr"] = group.get("lr", 3e-4)
                group["betas"] = group.get("betas", (0.9, 0.95))
                group["eps"] = group.get("eps", 1e-10)
                group["weight_decay"] = group.get("weight_decay", 0)
                assert set(group.keys()) == {
                    "params", "lr", "betas", "eps", "weight_decay", "use_muon"
                }
        super().__init__(param_groups, dict())
        split_qkv_params = tuple(split_qkv_params)
        muon_param_ids = {
            id(parameter)
            for group in self.param_groups
            if group["use_muon"]
            for parameter in group["params"]
        }
        split_ids = {id(parameter) for parameter in split_qkv_params}
        if len(split_ids) != len(split_qkv_params):
            raise ValueError("split-QKV parameter list contains duplicates")
        if not split_ids.issubset(muon_param_ids):
            raise ValueError("every split-QKV parameter must belong to a Muon group")
        for parameter in split_qkv_params:
            if parameter.ndim != 2 or parameter.size(0) != 3 * parameter.size(1):
                raise ValueError(
                    "split-QKV parameters must have shape (3*d, d), got %s"
                    % (tuple(parameter.shape),)
                )
        # Reconstructed from the trainer recipe on resume. Tensor optimizer
        # state remains the historical fused momentum buffer.
        self._split_qkv_param_ids = frozenset(split_ids)
        self._match_joint_frobenius = match_joint_frobenius

    @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:
            if group["use_muon"]:
                for p in group["params"]:
                    if p.grad is None:
                        continue
                    state = self.state[p]
                    if len(state) == 0:
                        state["momentum_buffer"] = torch.zeros_like(p)
                    if id(p) in self._split_qkv_param_ids:
                        update = muon_update_split_qkv(
                            p.grad,
                            state["momentum_buffer"],
                            beta=group["momentum"],
                            ns_steps=group.get("ns_steps", 5),
                            match_joint_frobenius=self._match_joint_frobenius,
                        )
                    else:
                        update = muon_update(
                            p.grad,
                            state["momentum_buffer"],
                            beta=group["momentum"],
                            # Old optimizer checkpoints predate this param-group key.
                            ns_steps=group.get("ns_steps", 5),
                        )
                    p.mul_(1 - group["lr"] * group["weight_decay"])
                    p.add_(update.reshape(p.shape), alpha=-group["lr"])
            else:
                for p in group["params"]:
                    if p.grad is None:
                        continue
                    state = self.state[p]
                    if len(state) == 0:
                        state["exp_avg"] = torch.zeros_like(p)
                        state["exp_avg_sq"] = torch.zeros_like(p)
                        state["step"] = 0
                    state["step"] += 1
                    update = adam_update(
                        p.grad,
                        state["exp_avg"],
                        state["exp_avg_sq"],
                        state["step"],
                        group["betas"],
                        group["eps"],
                    )
                    p.mul_(1 - group["lr"] * group["weight_decay"])
                    p.add_(update, alpha=-group["lr"])
        return loss


def build_muon_adam_optimizer(
    model,
    lr_adam,
    lr_muon,
    weight_decay,
    betas=(0.9, 0.95),
    *,
    muon_weight_decay=None,
    muon_ns_steps=5,
    muon_momentum=0.95,
    prior_lr=None,
    muon_split_qkv=False,
    match_joint_frobenius=True,
):
    """Split GPT params: 2D hidden → Muon; embed / lm_head / 1D → AdamW.

    ``muon_weight_decay=None`` preserves the historical behavior of applying
    ``weight_decay`` to both optimizer groups. ``muon_ns_steps=5`` preserves
    the historical five-step Newton–Schulz update. ``muon_split_qkv=False``
    preserves historical joint orthogonalization of fused attention QKV.
    ``match_joint_frobenius`` is only consumed when ``muon_split_qkv=True``.
    ``muon_momentum`` defaults to the historical 0.95 heavy-ball coefficient.
    """
    if type(muon_split_qkv) is not bool:
        raise ValueError("muon_split_qkv must be a bool")
    if type(match_joint_frobenius) is not bool:
        raise ValueError("match_joint_frobenius must be a bool")
    if not isinstance(muon_momentum, (int, float)) or not 0.0 < float(muon_momentum) < 1.0:
        raise ValueError("muon_momentum must be in (0, 1)")
    muon_momentum = float(muon_momentum)
    if muon_weight_decay is None:
        muon_weight_decay = weight_decay
    hidden_matrix = []
    adam_params = []
    prior_params = []
    split_qkv_params = []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if name == "logit_prior":
            prior_params.append(p)
        elif p.ndim >= 2 and "wte" not in name and "lm_head" not in name:
            hidden_matrix.append(p)
            if muon_split_qkv and name.endswith(".attn.c_attn.weight"):
                if p.ndim != 2 or p.size(0) != 3 * p.size(1):
                    raise ValueError(
                        "fused QKV parameter %s has unexpected shape %s"
                        % (name, tuple(p.shape))
                    )
                split_qkv_params.append(p)
        else:
            adam_params.append(p)
    groups = []
    if adam_params:
        groups.append(
            dict(
                params=adam_params,
                lr=lr_adam,
                betas=betas,
                eps=1e-10,
                weight_decay=weight_decay,
                use_muon=False,
            )
        )
    if prior_params:
        groups.append(
            dict(
                params=prior_params,
                lr=lr_adam if prior_lr is None else prior_lr,
                betas=betas,
                eps=1e-10,
                # The vocabulary intercept is a bias, not a weight matrix.
                weight_decay=0.0,
                use_muon=False,
            )
        )
    if hidden_matrix:
        groups.append(
            dict(
                params=hidden_matrix,
                lr=lr_muon,
                momentum=muon_momentum,
                weight_decay=muon_weight_decay,
                ns_steps=muon_ns_steps,
                use_muon=True,
            )
        )
    if not groups:
        raise RuntimeError("no parameters for Muon/Adam hybrid")
    if muon_split_qkv:
        if not split_qkv_params:
            raise RuntimeError(
                "muon_split_qkv requested but no .attn.c_attn.weight matrices found"
            )
        return SingleDeviceMuonWithAuxAdam(
            groups,
            split_qkv_params=split_qkv_params,
            match_joint_frobenius=match_joint_frobenius,
        )
    return SingleDeviceMuonWithAuxAdam(
        groups, match_joint_frobenius=match_joint_frobenius
    )


def handoff_muon_groups_to_adam(
    optimizer: SingleDeviceMuonWithAuxAdam,
    lr: float,
    betas=(0.9, 0.95),
    eps: float = 1e-10,
) -> int:
    """Switch only Muon matrix groups to fresh AdamW state in place.

    Existing auxiliary Adam groups (embedding/head/bias/prior) retain their
    moments.  Matrix momentum is deliberately discarded at the continuation
    boundary; subsequent calls are idempotent.  Returns the number of matrix
    parameters handed off.
    """
    if not isinstance(optimizer, SingleDeviceMuonWithAuxAdam):
        raise TypeError("Muon-to-Adam handoff requires SingleDeviceMuonWithAuxAdam")
    if not isinstance(lr, (int, float)) or not 0.0 < float(lr):
        raise ValueError("handoff Adam LR must be positive")
    switched = 0
    for group in optimizer.param_groups:
        if not group.get("use_muon", False):
            continue
        for parameter in group["params"]:
            optimizer.state[parameter].clear()
            switched += 1
        group.pop("momentum", None)
        group.pop("ns_steps", None)
        group["use_muon"] = False
        group["betas"] = tuple(betas)
        group["eps"] = float(eps)
        group["lr"] = float(lr)
        # The trainer scales every group by the common outer LR schedule.
        group["base_lr"] = float(lr)
    return switched