"""
Implementation of normalized steepest descent algorithm. We additionally apply multipliers to align the RMS norm update with Adam.
Note that when p=2 is chosen, the algorithm is the same as normalized SGD.
"""

import torch
import math

class NormalizedSteepestDescent(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, norm_p=2, lr=1e-3, momentum=0.9, weight_decay=0.1, scale=0.2, eps=1e-12, use_nesterov=False, use_grafting=False):
        # 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}")
        # use_grafting not implemented yet
        defaults = dict(norm_p=norm_p, lr=lr, momentum=momentum, weight_decay=weight_decay, scale=scale, eps=eps, use_nesterov=use_nesterov)
        super(NormalizedSteepestDescent, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(NormalizedSteepestDescent, 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:
            norm_p = group['norm_p']
            lr = group['lr']
            momentum = group['momentum']
            weight_decay = group['weight_decay']
            eps = group['eps']
            scale = group['scale']
            use_nesterov = group['use_nesterov']

            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)
                    if use_nesterov:
                        # use nesterov momentum, following Muon and Lion
                        Gt = torch.lerp(d_p, buf, momentum)
                    else:
                        Gt = buf
                else:
                    # If no momentum, the gradient is the update tensor
                    Gt = d_p

                mn = p.numel() # parameter count

                if norm_p <= 1 + eps:
                    # in this case, only the largest coordiantes have update
                    g_max_abs = torch.max(Gt.abs())
                    mask = Gt.abs() >= g_max_abs - eps
                    # Count how many elements have the max value (k in the formula)
                    k = torch.sum(mask).float()
                    # The update direction is sign(G_t) where |G_t| is max, and 0 otherwise.
                    update_direction = torch.zeros_like(Gt)
                    update_direction[mask] = torch.sign(Gt[mask])
                    
                    # Calculate the final update magnitude based on the limit
                    alpha = -lr * scale * torch.sqrt(mn / (k + eps))
                
                elif norm_p == float('inf'):
                    # in this case, we apply SignSGD
                    update_direction = torch.sign(Gt)
                    alpha = -lr * scale 
                
                else:
                    norm_q = 1 / (1 - 1 / norm_p)
                    assert norm_q > 0

                    # Calculate the update direction: |G_t|^(q-1) * sign(G_t)
                    update_direction = torch.pow(Gt.abs(), norm_q - 1) * torch.sign(Gt)
                    
                    update_norm = torch.linalg.norm(update_direction)

                    update_direction = update_direction / (update_norm + eps)
                    
                    # Calculate the pre-factor
                    alpha = -lr * scale * math.sqrt(mn)
                    
                    
                    

                # Apply weight decay
                if weight_decay > 0.0:
                    p.mul_(1.0 - lr * weight_decay)
                
                # Apply the update
                p.add_(update_direction, alpha=alpha)

        return loss