# activation_preconditioned_adam.py

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

class ActivationPreconditionedAdam(Optimizer):
    """
    Implements the Activation-Preconditioned Adam algorithm.
    This optimizer projects gradients into a preconditioned space derived from
    the SVD of the activation's second-moment matrix, performs an Adam update
    in that space, and projects the update back.
    """
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.95, 0.95), weight_decay=0.1, eps=1e-8):
        # beta1: momentum; beta2: adam second moment; beta3: preconditioner;

        # if not 0.0 <= lr:
        #     raise ValueError(f"Invalid learning rate: {lr}")
        # if not (isinstance(betas, (list, tuple)) and len(betas) == 3):
        #     raise ValueError(f"Betas must be a tuple/list of 3 elements.")
        # if not 0.0 <= betas[0] < 1.0:
        #     raise ValueError(f"Invalid beta1 parameter: {betas[0]}")
        # if not 0.0 <= betas[1] < 1.0:
        #     raise ValueError(f"Invalid beta2 parameter: {betas[1]}")
        # if not 0.0 <= betas[2] < 1.0:
        #     raise ValueError(f"Invalid beta3 parameter: {betas[2]}")
        # if not 0.0 <= weight_decay:
        #     raise ValueError(f"Invalid weight_decay value: {weight_decay}")

        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay, eps=eps)
        super(ActivationPreconditionedAdam, 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']
            
            for p in group['params']:
                if p.grad is None:
                    continue
                if not hasattr(p, 'activation_second_moment'):
                    continue

                grad = p.grad
                if grad.is_sparse:
                    raise RuntimeError('This optimizer does not support sparse gradients.')

                state = self.state[p]
                
                # State initialization
                if len(state) == 0:
                    state['step'] = 0
                    # M_t: Momentum buffer for gradients in original space
                    state['momentum_buffer'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                    # Λ_t: Preconditioner matrix
                    state['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']
                preconditioner = state['preconditioner']
                adam_v_prime = state['adam_v_prime']
                xxt = p.activation_second_moment

                # Step 5: Update momentum buffer M_t = β1*M_{t-1} + (1-β1)*G_t
                momentum_buffer.mul_(beta1).add_(grad, alpha=1 - beta1)

                # Step 6: Update preconditioner Λ_t = β2*Λ_{t-1} + (1-β2)*X_t*X_t^T
                preconditioner.mul_(beta2_precond).add_(xxt, alpha=1 - beta2_precond)

                # Step 7: Get unitary matrix U_t via SVD(Λ_t)
                # For a symmetric matrix like Λ_t, SVD (U,S,Vh) has U == Vh.T
                # We add a small epsilon for stability in case Λ_t is rank-deficient
                eps_I = torch.eye(preconditioner.size(0), device=preconditioner.device) * eps
                # we use float32 for eigendecomposition to ensure numerical stability
                # svd is numerically more accurate than eigh
                U, S, Vh = torch.linalg.svd(preconditioner.float() + eps_I)
                # S, U = torch.linalg.eigh(preconditioner.float() + eps_I)  # since Λ_t is symmetric

                # Step 8: Project momentum and gradient
                M_prime = torch.matmul(momentum_buffer, U)
                G_prime = torch.matmul(grad, U)

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

                # Step 9: Update Adam's second moment in projected space
                # V'_t = β3*V'_{t-1} + (1-β3)*G'_t*G'_t
                adam_v_prime.mul_(beta2_adam).addcmul_(G_prime, G_prime, value=1 - beta2_adam)

                # --- Adam Bias Correction ---
                bias_correction1 = 1 - beta1 ** step
                bias_correction3 = 1 - beta2_adam ** step
                
                # Correct the moments
                M_prime_hat = M_prime / bias_correction1
                adam_v_prime_hat = adam_v_prime / bias_correction3

                # Step 10: Run Adam in the preconditioned space
                # Q_t = M'_hat / (sqrt(V'_hat) + eps)
                denom = adam_v_prime_hat.sqrt().add_(eps)
                Q_t = M_prime_hat / denom

                # Apply weight decay (AdamW style) before the main update
                if weight_decay > 0.0:
                    p.add_(p, alpha=-lr * weight_decay)

                # Step 11: Project back to do the update
                # W_{t+1} = W_t - η_t * Q_t * U_t^T
                update = torch.matmul(Q_t, U.t())
                p.add_(update, alpha=-lr)

                # Cleanup
                delattr(p, 'activation_second_moment')
        return loss