"""Optimizers for Task 2. Rules: subclass torch.optim.Optimizer directly; no other torch.optim.* anywhere.

Paper: "Fantastic Pretraining Optimizers and Where to Find Them" (Wen et al., 2025), Appendix A has one algorithm per
optimizer. Every step() below follows the same pattern:
    for each param group -> for each parameter p with a gradient:
        g = p.grad (same shape as p), state = self.state[p] (dict that persists across steps, empty on step 1)
        -> initialise state on the first step -> update the running statistics -> apply the update to p in place.
Hyperparameters come from `group` (the dict built from `defaults`, overridable per group). The training loop sets
group["lr"] every step (learning-rate schedule), so always read the lr from the group, never cache it.
"""
import torch
from torch.optim import Optimizer


def _adamw_update(p, g, state, group):
    lr, (beta1, beta2), eps, wd = group["lr"], group["betas"], group["eps"], group["weight_decay"]
    if len(state) == 0:
        state["step"] = 0
        state["exp_avg"] = torch.zeros_like(p)
        state["exp_avg_sq"] = torch.zeros_like(p)
    state["step"] += 1
    t = state["step"]
    exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
    exp_avg.mul_(beta1).add_(g, alpha=1 - beta1)  # m_t = beta_1 * m_t-1 + (1 - beta_1) * g_t
    exp_avg_sq.mul_(beta2).addcmul_(g, g, value=1 - beta2)  # 
    bc1 = 1 - beta1 ** t
    bc2 = 1 - beta2 ** t
    p.mul_(1 - lr * wd)
    denom = (exp_avg_sq / bc2).sqrt_().add_(eps)
    p.addcdiv_(exp_avg, denom, value=-lr / bc1)  


class AdamW(Optimizer):
    """Category 1 (baseline). Appendix A, Algorithm 1."""

    def __init__(self, params, lr=6e-4, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1):
        super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay))

    @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:
            for p in group["params"]:
                if p.grad is None:
                    continue
                g, state = p.grad, self.state[p]
                _adamw_update(p, g, state, group)
        return loss


class NAdamW(Optimizer):
    """Category 2 (variance-reduced AdamW variants). Appendix A, Algorithm 2.
    Same state as AdamW; the difference is how the momentum enters the update (Nesterov look-ahead)."""

    def __init__(self, params, lr=6e-4, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1):
        super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay))

    @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, (beta1, beta2), eps, wd = group["lr"], group["betas"], group["eps"], group["weight_decay"]
            for p in group["params"]:
                if p.grad is None:
                    continue
                g, state = p.grad, self.state[p]
                if len(state) == 0:
                    state["step"] = 0
                    state["exp_avg"] = torch.zeros_like(p)
                    state["exp_avg_sq"] = torch.zeros_like(p)
                state["step"] += 1
                t = state["step"]
                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
                exp_avg.mul_(beta1).add_(g, alpha=1 - beta1)
                exp_avg_sq.mul_(beta2).addcmul_(g, g, value=1 - beta2)
                bc1 = 1 - beta1 ** t
                bc2 = 1 - beta2 ** t
                nesterov_m = exp_avg.mul(beta1).add_(g, alpha=1 - beta1)
                p.mul_(1 - lr * wd)
                denom = (exp_avg_sq / bc2).sqrt_().add_(eps)
                p.addcdiv_(nesterov_m, denom, value=-lr / bc1)
        return loss


class Lion(Optimizer):
    """Category 3 (memory-efficient). Appendix A, Algorithm 3.
    One state buffer per parameter; the update is a sign, so every coordinate moves by exactly lr
    -> needs a much smaller lr (and usually larger weight decay) than AdamW."""

    def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.1):
        super().__init__(params, dict(lr=lr, betas=betas, weight_decay=weight_decay))

    @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, (beta1, beta2), wd = group["lr"], group["betas"], group["weight_decay"]
            for p in group["params"]:
                if p.grad is None:
                    continue
                g, state = p.grad, self.state[p]
                if len(state) == 0:
                    state["exp_avg"] = torch.zeros_like(p)
                exp_avg = state["exp_avg"]
                update = exp_avg.mul(beta1).add_(g, alpha=1 - beta1).sign_()
                p.mul_(1 - lr * wd)
                p.add_(update, alpha=-lr)
                exp_avg.mul_(beta2).add_(g, alpha=1 - beta2)
        return loss


def zeropower_via_newtonschulz(G, steps):
    """G: (m, n) matrix -> matrix of the same shape with the same singular vectors and singular values pushed to ~1
    (an approximate orthogonalisation), computed with `steps` Newton-Schulz iterations. Used by Muon."""
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G.float()
    X = X / (X.norm() + 1e-7)
    transposed = G.size(0) > G.size(1)
    if transposed:
        X = X.T
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * (A @ A)
        X = a * X + B @ X
    if transposed:
        X = X.T
    return X


class Muon(Optimizer):
    """Category 4 (matrix-based). Appendix A, Algorithm 8.
    Two kinds of param groups (built in train.py):
      use_muon=True  -> hidden 2-D weight matrices (attention q/k/v/out, FFN linears): Muon update
      use_muon=False -> embedding (= tied output layer), norm gains, biases: AdamW update (reuse your AdamW logic)
    The two kinds use different learning rates (Muon's is much larger), set per group in train.py."""

    def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5, weight_decay=0.0,
                 betas=(0.9, 0.95), eps=1e-8):
        super().__init__(params, dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps,
                                      weight_decay=weight_decay, betas=betas, eps=eps, use_muon=True))

    @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:
            for p in group["params"]:
                if p.grad is None:
                    continue
                g, state = p.grad, self.state[p]
                if group["use_muon"]:
                    lr, momentum, wd = group["lr"], group["momentum"], group["weight_decay"]
                    if len(state) == 0:
                        state["momentum_buffer"] = torch.zeros_like(p)
                    buf = state["momentum_buffer"]
                    buf.lerp_(g, 1 - momentum)
                    update = g.lerp(buf, momentum) if group["nesterov"] else buf
                    update = zeropower_via_newtonschulz(update, group["ns_steps"])
                    scale = max(1.0, p.size(0) / p.size(1)) ** 0.5
                    p.mul_(1 - lr * wd)
                    p.add_(update.to(p.dtype), alpha=-lr * scale)
                else:
                    _adamw_update(p, g, state, group)
        return loss


OPTIMIZERS = {"adamw": AdamW, "nadamw": NAdamW, "lion": Lion, "muon": Muon}