"""
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. By default the split direction is therefore
    rescaled to the realized norm of the historical joint NS5 direction from
    the same preconditioned tensor. This is an explicit per-update control,
    not an LR tune; its extra joint polar arithmetic remains on the clock.
    """
    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=()):
        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)

    @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),
                        )
                    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,
    prior_lr=None,
    muon_split_qkv=False,
):
    """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.
    """
    if type(muon_split_qkv) is not bool:
        raise ValueError("muon_split_qkv must be a bool")
    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=0.95,
                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
        )
    return SingleDeviceMuonWithAuxAdam(groups)


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