"""GrokAdamW: AdamW c уклоном в обобщение.

  - decoupled weight decay (только на матрицах, не на нормах/эмбеддингах);
  - cautious-маскинг (C-AdamW): применяем компоненты шага, согласованные по
    знаку с градиентом -> стабильнее и чуть быстрее сходимость;
  - Grokfast (опц., gf_lambda>0): усиление медленной компоненты градиента.
    Для большого претрейна обычно держим gf_lambda=0 (нужен для grokking на
    маленьких алгоритмических задачах, не для масштабного обучения).
"""

import math
import torch


class GrokAdamW(torch.optim.Optimizer):
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.95), eps=1e-8,
                 weight_decay=0.1, grokfast_lambda=0.0, grokfast_alpha=0.98,
                 cautious=True):
        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay,
                        grokfast_lambda=grokfast_lambda, grokfast_alpha=grokfast_alpha,
                        cautious=cautious)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        loss = closure() if closure is not None else None
        for grp in self.param_groups:
            lr, (b1, b2), eps = grp["lr"], grp["betas"], grp["eps"]
            wd, gfl, gfa, caut = (grp["weight_decay"], grp["grokfast_lambda"],
                                  grp["grokfast_alpha"], grp["cautious"])
            for p in grp["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 gfl > 0:
                        st["ema"] = g.clone()
                if gfl > 0:
                    ema = st["ema"]; ema.mul_(gfa).add_(g, alpha=1 - gfa)
                    g = g.add(ema, alpha=gfl)
                m, v = st["m"], st["v"]
                st["step"] += 1; t = st["step"]
                m.mul_(b1).add_(g, alpha=1 - b1)
                v.mul_(b2).addcmul_(g, g, value=1 - b2)
                denom = (v.sqrt() / math.sqrt(1 - b2 ** t)).add_(eps)
                step_size = lr / (1 - b1 ** t)
                if wd != 0:
                    p.mul_(1 - lr * wd)
                upd = m / denom
                if caut:
                    mask = (upd * g > 0).to(upd.dtype)
                    mask.mul_(mask.numel() / (mask.sum() + 1))
                    upd.mul_(mask)
                p.add_(upd, alpha=-step_size)
        return loss


def param_groups(model, weight_decay):
    decay, no_decay = [], []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        (no_decay if p.ndim < 2 or name.endswith(".bias") else decay).append(p)
    return [{"params": decay, "weight_decay": weight_decay},
            {"params": no_decay, "weight_decay": 0.0}]