# =============================================================================
# optimizers.py - Part 2: optimizers implemented from scratch.
# Only torch.optim.Optimizer is used, strictly as a base class.
#
#   Category (Wen et al. 2025, Table 1)       Implemented here
#   ---------------------------------------   -------------------------------
#   AdamW (baseline)                          AdamW       (Alg. 1)
#   Variance-reduced AdamW variants           MARS        (Alg. 5, MARS-AdamW)
#   Memory-efficient optimizers               Lion        (Alg. 3)
#   Matrix-based optimizers                   Muon        (Alg. 8) (+AdamW for non-matrix params)
#   Hessian-approximation (bonus)             Sophia      (Alg. 4; Hutchinson "H" or Gauss-Newton-Bartlett "G")
# =============================================================================
import math
import torch
from torch.optim import Optimizer   # base class only


def clip_grad_norm(params, max_norm):
    """Global L2 gradient clipping (re-implemented; returns the pre-clip norm)."""
    grads = [p.grad for p in params if p.grad is not None]
    if not grads:
        return torch.zeros(())
    total = torch.linalg.vector_norm(torch.stack([torch.linalg.vector_norm(g.float()) for g in grads]))
    if max_norm and max_norm > 0:
        coef = (max_norm / (total + 1e-6)).clamp(max=1.0)
        for g in grads:
            g.mul_(coef.to(g.dtype))
    return total


def split_decay_groups(model, weight_decay):
    decay, no_decay = [], []
    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue
        (decay if p.dim() >= 2 else no_decay).append(p)
    return [{"params": decay, "weight_decay": weight_decay},
            {"params": no_decay, "weight_decay": 0.0}]


# -----------------------------------------------------------------------------
# 1. AdamW
# -----------------------------------------------------------------------------
class AdamW(Optimizer):
    """m_t = b1 m + (1-b1) g ; v_t = b2 v + (1-b2) g^2
    w <- w - lr * (m_hat / (sqrt(v_hat) + eps)) - lr * wd * w"""
    def __init__(self, params, lr=1e-3, 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 = closure() if closure is not None else None
        for group in self.param_groups:
            lr, (b1, b2), eps, wd = group["lr"], group["betas"], group["eps"], group["weight_decay"]
            for p in group["params"]:
                if p.grad is None:
                    continue
                g = p.grad
                st = self.state[p]
                if not st:
                    st["step"] = 0
                    st["m"] = torch.zeros_like(p, memory_format=torch.preserve_format)
                    st["v"] = torch.zeros_like(p, memory_format=torch.preserve_format)
                st["step"] += 1
                t = st["step"]
                m, v = st["m"], st["v"]
                m.lerp_(g, 1 - b1)
                v.mul_(b2).addcmul_(g, g, value=1 - b2)
                bc1, bc2 = 1 - b1 ** t, 1 - b2 ** t
                if wd:
                    p.mul_(1 - lr * wd)
                denom = (v.sqrt() / math.sqrt(bc2)).add_(eps)
                p.addcdiv_(m, denom, value=-lr / bc1)
        return loss


# -----------------------------------------------------------------------------
# 2. MARS (variance reduction, MARS-AdamW, approximate version with g_{t-1})
# -----------------------------------------------------------------------------
class MARS(Optimizer):
    """c_t = g_t + gamma * b1/(1-b1) * (g_t - g_{t-1})
    c_t <- c_t / max(1, ||c_t||)        (per-tensor clipping of the corrected gradient)
    then AdamW moments / update on c_t.
    Groups with gamma=0 (1-D params, as in the official MARS code) reduce to AdamW."""
    def __init__(self, params, lr=3e-3, betas=(0.95, 0.99), eps=1e-8, weight_decay=0.1,
                 gamma=0.025, max_c_norm=1.0):
        super().__init__(params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay,
                                      gamma=gamma, max_c_norm=max_c_norm))

    @torch.no_grad()
    def step(self, closure=None):
        loss = closure() if closure is not None else None
        for group in self.param_groups:
            lr, (b1, b2), eps, wd = group["lr"], group["betas"], group["eps"], group["weight_decay"]
            gamma, max_c = group["gamma"], group["max_c_norm"]
            for p in group["params"]:
                if p.grad is None:
                    continue
                g = p.grad
                st = self.state[p]
                if not st:
                    st["step"] = 0
                    st["m"] = torch.zeros_like(p)
                    st["v"] = torch.zeros_like(p)
                    if gamma:
                        st["prev_grad"] = torch.zeros_like(p)
                st["step"] += 1
                t = st["step"]
                if gamma:
                    if t == 1:
                        c = g.clone()
                    else:
                        c = g + (gamma * b1 / (1 - b1)) * (g - st["prev_grad"])
                    st["prev_grad"].copy_(g)
                    if max_c:
                        n = torch.linalg.vector_norm(c)
                        c.mul_((max_c / (n + 1e-6)).clamp(max=1.0))
                else:
                    c = g
                m, v = st["m"], st["v"]
                m.lerp_(c, 1 - b1)
                v.mul_(b2).addcmul_(c, c, value=1 - b2)
                bc1, bc2 = 1 - b1 ** t, 1 - b2 ** t
                if wd:
                    p.mul_(1 - lr * wd)
                denom = (v.sqrt() / math.sqrt(bc2)).add_(eps)
                p.addcdiv_(m, denom, value=-lr / bc1)
        return loss


# -----------------------------------------------------------------------------
# 3. Lion (memory-efficient: a single momentum buffer, sign update)
# -----------------------------------------------------------------------------
class Lion(Optimizer):
    """u = sign(b1 m + (1-b1) g) ; w <- w - lr*(u + wd*w) ; m <- b2 m + (1-b2) g"""
    def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=1.0):
        super().__init__(params, dict(lr=lr, betas=betas, weight_decay=weight_decay))

    @torch.no_grad()
    def step(self, closure=None):
        loss = closure() if closure is not None else None
        for group in self.param_groups:
            lr, (b1, b2), wd = group["lr"], group["betas"], group["weight_decay"]
            for p in group["params"]:
                if p.grad is None:
                    continue
                g = p.grad
                st = self.state[p]
                if not st:
                    st["m"] = torch.zeros_like(p)
                m = st["m"]
                u = torch.lerp(g, m, b1).sign_()          # b1*m + (1-b1)*g
                if wd:
                    p.mul_(1 - lr * wd)
                p.add_(u, alpha=-lr)
                m.lerp_(g, 1 - b2)
        return loss


# -----------------------------------------------------------------------------
# 4. Muon (matrix-based: orthogonalised momentum via Newton-Schulz)
# -----------------------------------------------------------------------------
def zeropower_via_newtonschulz5(G, steps=5, eps=1e-7):
    """Quintic Newton-Schulz iteration approximating the orthogonal polar factor U V^T
    of G = U S V^T (coefficients from Jordan et al. 2024). Runs in fp32 (T4/P100 lack bf16)."""
    assert G.ndim == 2
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G.float()
    transposed = X.size(0) > X.size(1)
    if transposed:
        X = X.T
    X = X / (X.norm() + eps)
    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):
    """Hidden 2-D weight matrices: Nesterov momentum -> Newton-Schulz orthogonalisation ->
    scale by sqrt(max(1, rows/cols)).  Everything else (embeddings / tied LM head,
    LayerNorm gains, biases) is updated with AdamW, as in Alg. 8 of the paper."""
    def __init__(self, muon_params, adamw_params, lr=0.02, momentum=0.95, nesterov=True,
                 ns_steps=5, weight_decay=0.01, adamw_lr=3e-3, adamw_betas=(0.9, 0.95),
                 adamw_eps=1e-8, adamw_wd=0.1):
        groups = [dict(params=list(muon_params), use_muon=True, lr=lr, momentum=momentum,
                       nesterov=nesterov, ns_steps=ns_steps, weight_decay=weight_decay)]
        adamw_params = list(adamw_params)
        dec = [p for p in adamw_params if p.dim() >= 2]
        nodec = [p for p in adamw_params if p.dim() < 2]
        for ps, wd in ((dec, adamw_wd), (nodec, 0.0)):
            if ps:
                groups.append(dict(params=ps, use_muon=False, lr=adamw_lr, betas=adamw_betas,
                                   eps=adamw_eps, weight_decay=wd))
        super().__init__(groups, dict(lr=lr))

    @torch.no_grad()
    def step(self, closure=None):
        loss = closure() if closure is not None else None
        for group in self.param_groups:
            lr, wd = group["lr"], group["weight_decay"]
            if group["use_muon"]:
                mom = group["momentum"]
                for p in group["params"]:
                    if p.grad is None:
                        continue
                    g = p.grad
                    st = self.state[p]
                    if not st:
                        st["momentum_buffer"] = torch.zeros_like(p)
                    buf = st["momentum_buffer"]
                    buf.mul_(mom).add_(g)
                    u = g.add(buf, alpha=mom) if group["nesterov"] else buf
                    u = zeropower_via_newtonschulz5(u, group["ns_steps"])
                    scale = max(1.0, p.size(0) / p.size(1)) ** 0.5
                    if wd:
                        p.mul_(1 - lr * wd)
                    p.add_(u.to(p.dtype), alpha=-lr * scale)
            else:
                (b1, b2), eps = group["betas"], group["eps"]
                for p in group["params"]:
                    if p.grad is None:
                        continue
                    g = p.grad
                    st = self.state[p]
                    if not st:
                        st["step"] = 0
                        st["m"] = torch.zeros_like(p)
                        st["v"] = torch.zeros_like(p)
                    st["step"] += 1
                    t = st["step"]
                    st["m"].lerp_(g, 1 - b1)
                    st["v"].mul_(b2).addcmul_(g, g, value=1 - b2)
                    if wd:
                        p.mul_(1 - lr * wd)
                    denom = (st["v"].sqrt() / math.sqrt(1 - b2 ** t)).add_(eps)
                    p.addcdiv_(st["m"], denom, value=-lr / (1 - b1 ** t))
        return loss


def muon_param_split(model):
    """2-D weights inside transformer blocks -> Muon; embeddings / lm_head / 1-D -> AdamW."""
    muon, other, seen = [], [], set()
    for n, p in model.named_parameters():
        if id(p) in seen or not p.requires_grad:
            continue
        seen.add(id(p))
        if p.dim() == 2 and ".h." in f".{n}" and "wte" not in n and "wpe" not in n and "lm_head" not in n:
            muon.append(p)
        else:
            other.append(p)
    return muon, other


# -----------------------------------------------------------------------------
# 5. Sophia (bonus) - clipped diagonal-Hessian preconditioned momentum
# -----------------------------------------------------------------------------
class Sophia(Optimizer):
    """m_t = b1 m + (1-b1) g
    every k steps: h = b2 h + (1-b2) h_hat   (h_hat from Hutchinson u*(Hu) or GNB  B*g_hat^2)
    w <- w - lr*wd*w ;  w <- w - lr * sign(m) * min(|m| / max(rho*h, eps), 1)"""
    def __init__(self, params, lr=6e-4, betas=(0.965, 0.99), rho=0.05, weight_decay=0.2, eps=1e-15):
        super().__init__(params, dict(lr=lr, betas=betas, rho=rho, weight_decay=weight_decay, eps=eps))
        self.last_clip_frac = float("nan")

    @torch.no_grad()
    def update_hessian(self, h_hats):
        """h_hats: dict param -> diagonal Hessian estimate (same shape)."""
        for group in self.param_groups:
            b2 = group["betas"][1]
            for p in group["params"]:
                if p not in h_hats:
                    continue
                st = self.state[p]
                if "h" not in st:
                    st["h"] = torch.zeros_like(p)
                st["h"].mul_(b2).add_(h_hats[p], alpha=1 - b2)

    @torch.no_grad()
    def step(self, closure=None):
        loss = closure() if closure is not None else None
        clipped, total = 0, 0
        for group in self.param_groups:
            lr, (b1, _), rho, wd, eps = (group["lr"], group["betas"], group["rho"],
                                         group["weight_decay"], group["eps"])
            for p in group["params"]:
                if p.grad is None:
                    continue
                st = self.state[p]
                if "m" not in st:
                    st["m"] = torch.zeros_like(p)
                if "h" not in st:
                    st["h"] = torch.zeros_like(p)
                m, h = st["m"], st["h"]
                m.lerp_(p.grad, 1 - b1)
                if wd:
                    p.mul_(1 - lr * wd)
                ratio = (m.abs() / (rho * h).clamp(min=eps)).clamp(max=1.0)
                p.addcmul_(m.sign(), ratio, value=-lr)
                clipped += int((ratio >= 1.0).sum()); total += ratio.numel()
        self.last_clip_frac = clipped / max(total, 1)
        return loss


def sophia_hutchinson_estimate(model, x, y, params):
    """h_hat = u * (H u), u ~ Rademacher. Needs double backward -> fp32 + math attention."""
    ctx = None
    try:
        from torch.nn.attention import sdpa_kernel, SDPBackend
        ctx = sdpa_kernel(SDPBackend.MATH)
    except Exception:
        import contextlib
        ctx = torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False,
                                             enable_math=True) if torch.cuda.is_available() else contextlib.nullcontext()
    with ctx:
        _, loss, _ = model(x, y)
        grads = torch.autograd.grad(loss, params, create_graph=True)
        us = [torch.randint_like(p, 0, 2) * 2 - 1 for p in params]
        hvps = torch.autograd.grad(grads, params, grad_outputs=us)
    return {p: (u * hv).detach() for p, u, hv in zip(params, us, hvps)}


def sophia_gnb_estimate(model, x, params, amp_ctx, scaler=None):
    """Gauss-Newton-Bartlett: sample y_hat ~ softmax(logits), h_hat = N_tokens * grad(CE(y_hat))^2."""
    import torch.nn.functional as F
    model.zero_grad(set_to_none=True)
    with amp_ctx:
        logits, _, _ = model(x)
    with torch.no_grad():
        probs = F.softmax(logits.float().view(-1, logits.size(-1)), -1)
        y_hat = torch.multinomial(probs, 1).view(x.shape)
        del probs
    loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), y_hat.view(-1))
    scale = scaler.get_scale() if (scaler is not None and scaler.is_enabled()) else 1.0
    (loss * scale).backward()
    n_tok = x.numel()
    out = {}
    ok = True
    for p in params:
        if p.grad is None:
            continue
        g = p.grad.float() / scale
        if not torch.isfinite(g).all():
            ok = False; break
        out[p] = n_tok * g * g
    model.zero_grad(set_to_none=True)
    return out if ok else {}


# -----------------------------------------------------------------------------
# factory
# -----------------------------------------------------------------------------
OPTIMIZER_DEFAULTS = {
    # tuned for a ~27M-param GPT (d=512, 6 layers) with 64x256-token batches;
    # override from the notebook or let the sweep pick the LR.
    "adamw":  dict(lr=2e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1, grad_clip=1.0),
    "mars":   dict(lr=2e-3, betas=(0.95, 0.99), eps=1e-8, weight_decay=0.1, gamma=0.025,
                   max_c_norm=1.0, grad_clip=0.0),   # MARS clips its corrected gradient itself
    "lion":   dict(lr=2e-4, betas=(0.9, 0.99), weight_decay=1.0, grad_clip=1.0),
    "muon":   dict(lr=0.02, momentum=0.95, nesterov=True, ns_steps=5, weight_decay=0.01,
                   adamw_lr=2e-3, adamw_betas=(0.9, 0.95), adamw_wd=0.1, grad_clip=1.0),
    "sophia": dict(lr=6e-4, betas=(0.965, 0.99), rho=0.05, weight_decay=0.2, grad_clip=1.0,
                   hess_interval=10, estimator="gnb", hess_batch=16),
}

OPTIMIZER_CATEGORY = {
    "adamw": "AdamW (baseline)",
    "mars": "Variance-reduced AdamW variant",
    "lion": "Memory-efficient",
    "muon": "Matrix-based",
    "sophia": "Hessian-approximation (bonus)",
}

SWEEP_GRIDS = {
    "adamw": [5e-4, 1e-3, 2e-3, 4e-3],
    "mars": [5e-4, 1e-3, 2e-3, 4e-3],
    "lion": [5e-5, 1e-4, 2e-4, 4e-4],
    "muon": [0.005, 0.01, 0.02, 0.04],
    "sophia": [2e-4, 4e-4, 8e-4, 1.6e-3],
}


def build_optimizer(name, model, hp):
    """Returns (optimizer, hp_used). hp overrides OPTIMIZER_DEFAULTS[name]."""
    h = dict(OPTIMIZER_DEFAULTS[name]); h.update(hp or {})
    if name == "adamw":
        opt = AdamW(split_decay_groups(model, h["weight_decay"]), lr=h["lr"], betas=tuple(h["betas"]), eps=h["eps"])
    elif name == "mars":
        dec, nodec = [], []
        for n, p in model.named_parameters():
            (dec if p.dim() >= 2 else nodec).append(p)
        groups = [dict(params=dec, weight_decay=h["weight_decay"], gamma=h["gamma"]),
                  dict(params=nodec, weight_decay=0.0, gamma=0.0)]
        opt = MARS(groups, lr=h["lr"], betas=tuple(h["betas"]), eps=h["eps"], max_c_norm=h["max_c_norm"])
    elif name == "lion":
        opt = Lion(split_decay_groups(model, h["weight_decay"]), lr=h["lr"], betas=tuple(h["betas"]))
    elif name == "muon":
        mp, op = muon_param_split(model)
        opt = Muon(mp, op, lr=h["lr"], momentum=h["momentum"], nesterov=h["nesterov"], ns_steps=h["ns_steps"],
                   weight_decay=h["weight_decay"], adamw_lr=h["adamw_lr"], adamw_betas=tuple(h["adamw_betas"]),
                   adamw_wd=h["adamw_wd"])
    elif name == "sophia":
        opt = Sophia(split_decay_groups(model, h["weight_decay"]), lr=h["lr"], betas=tuple(h["betas"]), rho=h["rho"])
    else:
        raise ValueError(name)
    return opt, h


def optimizer_state_bytes(opt):
    n = 0
    for st in opt.state.values():
        for v in st.values():
            if torch.is_tensor(v):
                n += v.numel() * v.element_size()
    return n