# activation_gradient_preconditioned_adam.py

import torch
import math
from torch.optim.optimizer import Optimizer

class ASGO(Optimizer):
    """
    Implements Activation-Gradient-Preconditioned Adam.
    This optimizer uses second-order statistics from both input activations (X_t)
    and output gradients (G_Yt) to form two preconditioners. This is closely
    related to the Shampoo optimizer.
    """
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.8), weight_decay=0.1, eps=1e-10, scale=0.2, use_grafting=False, use_nesterov=False):
        # betas[0]: momentum (β1)
        # betas[1]: adam second moment (β2 in pseudocode, for V')
        # betas[2]: preconditioners (β3 in pseudocode, for L and R)
        # if not (isinstance(betas, (list, tuple)) and len(betas) == 3):
        #     raise ValueError("Betas must be a tuple/list of 3 elements.")
        
        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay, eps=eps, scale=scale, use_grafting=use_grafting, use_nesterov=use_nesterov)
        super(ASGO, self).__init__(params, defaults)

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

        for group in self.param_groups:
            lr = group['lr']
            if len(group['betas']) == 2:
                beta1, beta2_adam = group['betas']
                beta2_precond = beta2_adam
            else:
                beta1, beta2_adam, beta2_precond = group['betas']
            weight_decay = group['weight_decay']
            eps = group['eps']
            scale = group['scale']
            use_grafting = group['use_grafting']
            use_nesterov = group['use_nesterov']
            
            for p in group['params']:
                if p.grad is None:
                    continue

                grad = p.grad
                state = self.state[p]
                
                if len(state) == 0:
                    state['step'] = 0
                    state['momentum_buffer'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                    # R_t: Right preconditioner
                    state['right_preconditioner'] = torch.zeros((p.shape[-1], p.shape[-1]), device=p.device)
                    # V'_t: Adam's 2nd moment buffer for grafting
                    if use_grafting:
                        # beta2_adam = beta2_precond
                        state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

                state['step'] += 1
                step = state['step']
                
                momentum_buffer = state['momentum_buffer']
                right_preconditioner = state['right_preconditioner']
                
                if use_grafting:
                    exp_avg_sq = state['exp_avg_sq']
                    exp_avg_sq.mul_(beta2_adam).addcmul_(grad, grad, value=1 - beta2_adam)


                
                GTG = torch.mm(grad.t(), grad)

                # Step 5: Update momentum
                momentum_buffer.mul_(beta1).add_(grad, alpha=1 - beta1)
                if use_nesterov:
                    Gt = torch.lerp(grad, momentum_buffer, beta1)
                else:
                    Gt = momentum_buffer

                # Step 6 & 7: Update preconditioners
                # left_preconditioner.mul_(beta2_precond).add_(GGT, alpha=1 - beta2_precond)
                right_preconditioner.mul_(beta2_precond).add_(GTG, alpha=1 - beta2_precond)

                # Step 8 & 9: Eigendecomposition of preconditioners
                # we use float32 for eigendecomposition to ensure numerical stability
                # eps_L = torch.eye(left_preconditioner.size(0), device=p.device) * eps
                # S_L, U = torch.linalg.eigh(left_preconditioner.float() + eps_L)
                
                eps_R = torch.eye(right_preconditioner.size(0), device=p.device) * eps
                
                # S_R, V = torch.linalg.eigh(right_preconditioner.float() + eps_R)
                # update_direction = Gt @ V @ torch.diag(S_R.rsqrt()) @ V.t()
                
                U, S_R, Vh = torch.linalg.svd(right_preconditioner.float() + eps_R)
                update_direction = Gt @ U @ torch.diag(S_R.rsqrt()) @ Vh

                
                # normalize the update
                update_direction = update_direction / (torch.linalg.norm(update_direction) + eps)

                if use_grafting:
                    bias_correction1 = 1 - beta1 ** step
                    bias_correction2 = 1 - beta2_adam ** step
                        
                    denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
                    # The magnitude is calculated based on the effective gradient
                    adam_update_magnitude = torch.linalg.norm((Gt / bias_correction1) / denom)
                    alpha = -lr * adam_update_magnitude
                
                else:
                    # Default scaling if not grafting
                    alpha = -lr * scale * math.sqrt(p.numel())

                
                if weight_decay > 0.0:
                    p.add_(p, alpha=-lr * weight_decay)

                
                p.add_(update_direction, alpha=alpha)

                
        return loss