"""
Implementation of normalized SGD algorithm. We additionally apply 0.2 * \sqrt{mn} to align the RMS norm update with Adam.
"""

import torch
import math

class NormalizedSGD(torch.optim.Optimizer):
    """
    Implements a normalized version of Stochastic Gradient Descent (with momentum).

    The update for each parameter is normalized by its Frobenius norm before
    being applied. This is a form of layer-wise adaptivity.
    """
    def __init__(self, params, lr=1e-3, momentum=0.9, weight_decay=0.1, scale=0.2, eps=1e-12):
        # scale is used for aligning the update RMS with Adam
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if momentum < 0.0:
            raise ValueError(f"Invalid momentum value: {momentum}")
        if weight_decay < 0.0:
            raise ValueError(f"Invalid weight_decay value: {weight_decay}")

        defaults = dict(lr=lr, momentum=momentum, weight_decay=weight_decay, scale=scale, eps=eps)
        super(NormalizedSGD, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(NormalizedSGD, self).__setstate__(state)

    @torch.no_grad()
    def step(self, closure=None):
        """Performs a single optimization step."""
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            lr = group['lr']
            momentum = group['momentum']
            weight_decay = group['weight_decay']
            eps = group['eps']
            scale = group['scale']

            for p in group['params']:
                if p.grad is None:
                    continue
                
                # Get the gradient
                d_p = p.grad

                # Apply momentum
                if momentum != 0:
                    param_state = self.state[p]
                    if 'momentum_buffer' not in param_state:
                        buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()
                    else:
                        buf = param_state['momentum_buffer']
                        buf.mul_(momentum).add_(d_p, alpha=1-momentum) # m_t = beta * m_{t-1} + (1-beta) * g_t
                    
                    # This is the update tensor (m_t)
                    update_vec = buf
                else:
                    # If no momentum, the gradient is the update tensor
                    update_vec = d_p

                # +++ CORE MODIFICATION: Normalize the update vector +++
                # Calculate the Frobenius norm of the update vector
                norm = torch.linalg.norm(update_vec)

                # Apply weight decay
                if weight_decay > 0.0:
                    p.mul_(1.0 - lr * weight_decay)
                
                # Normalize and apply the update
                # p_new = p_old - lr * (m_t / ||m_t||_F)
                scale_align = scale * math.sqrt(p.numel())
                p.add_(update_vec, alpha=-lr * scale_align / (norm + eps))

        return loss