"""AdamAtan2 (Everett et al., 2024; used by TRM) and a warmup + cosine LR schedule."""
import math

import torch


class AdamAtan2(torch.optim.Optimizer):
    """Adam with the division m / sqrt(v) replaced by a * atan2(m, b * sqrt(v)); decoupled weight decay."""

    def __init__(self, params, lr=1e-4, betas=(0.9, 0.95), weight_decay=0.0, a=1.27, b=1.0):
        super().__init__(params, dict(lr=lr, betas=betas, weight_decay=weight_decay, a=a, b=b))

    @torch.no_grad()
    def step(self):
        for g in self.param_groups:
            b1, b2 = g["betas"]
            for p in g["params"]:
                if p.grad is None:
                    continue
                s = self.state[p]
                if not s:
                    s["step"], s["m"], s["v"] = 0, torch.zeros_like(p), torch.zeros_like(p)
                s["step"] += 1
                p.mul_(1 - g["lr"] * g["weight_decay"])
                s["m"].lerp_(p.grad, 1 - b1)
                s["v"].lerp_(p.grad * p.grad, 1 - b2)
                m_hat = s["m"] / (1 - b1 ** s["step"])
                v_hat = s["v"] / (1 - b2 ** s["step"])
                p.add_(torch.atan2(m_hat, g["b"] * v_hat.sqrt()), alpha=-g["lr"] * g["a"])


def lr_at(step, base_lr, warmup, total, min_ratio=1.0):
    if step < warmup:
        return base_lr * (step + 1) / warmup
    progress = min(1.0, (step - warmup) / max(1, total - warmup))
    return base_lr * (min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * progress)))


def set_lr(opt, lr):
    for g in opt.param_groups:
        g["lr"] = lr