# activation_gradient_preconditioned_adam.py

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

class SOAP(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.999, 0.95), weight_decay=0.1, eps=1e-6):
        # 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)
        super(SOAP, 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']
                beta3_precond = beta2_adam
            else:
                beta1, beta2_adam, beta3_precond = group['betas']
            weight_decay = group['weight_decay']
            eps = group['eps']
            
            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)
                    # L_t: Gradient preconditioner
                    state['left_preconditioner'] = torch.zeros((p.shape[0], p.shape[0]), device=p.device)
                    # R_t: Activation preconditioner
                    state['right_preconditioner'] = torch.zeros((p.shape[1], p.shape[1]), device=p.device)
                    # V'_t: Adam's 2nd moment buffer in the *projected* space
                    state['adam_v_prime'] = torch.zeros_like(p, memory_format=torch.preserve_format)

                state['step'] += 1
                step = state['step']
                
                momentum_buffer = state['momentum_buffer']
                left_preconditioner = state['left_preconditioner']
                right_preconditioner = state['right_preconditioner']
                adam_v_prime = state['adam_v_prime']

                # xxt = p.activation_second_moment
                # gygyt = p.output_gradient_second_moment

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

                # Step 5: Update momentum
                momentum_buffer.mul_(beta1).add_(grad, alpha=1 - beta1)

                # Step 6 & 7: Update preconditioners
                left_preconditioner.mul_(beta3_precond).add_(GGT, alpha=1 - beta3_precond)
                right_preconditioner.mul_(beta3_precond).add_(GTG, alpha=1 - beta3_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)
                # use svd instead of eigh for better numerical stability
                U, S_L, _ = torch.linalg.svd(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)
                # use svd instead of eigh for better numerical stability
                V, S_R, _ = torch.linalg.svd(right_preconditioner.float() + eps_R)

                # Step 10: Project momentum and gradient
                M_prime = torch.mm(U.t(), torch.mm(momentum_buffer, V))
                G_prime = torch.mm(U.t(), torch.mm(grad, V))

                # +++ ADDED: Store G_prime in state for external logging +++
                # will cost more memory, only use it for logging
                state['G_prime'] = G_prime

                # Step 11: Update Adam's second moment in projected space
                adam_v_prime.mul_(beta2_adam).addcmul_(G_prime, G_prime, value=1 - beta2_adam)

                # Bias Correction for Adam
                bias_correction1 = 1 - beta1 ** step
                bias_correction2 = 1 - beta2_adam ** step
                
                M_prime_hat = M_prime / bias_correction1
                adam_v_prime_hat = adam_v_prime / bias_correction2

                # Step 12: Run Adam in the preconditioned space
                denom = adam_v_prime_hat.sqrt().add_(eps)
                Q_t = M_prime_hat / denom

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

                # Step 13: Project back to do the update
                update = torch.mm(U, torch.mm(Q_t, V.t()))
                p.add_(update, alpha=-lr)

                
        return loss