import torch
import math

class ColumnWiseAdamW(torch.optim.Optimizer):
    """
    Implements a Column-wise AdamW optimization algorithm.
    
    The first-order momentum (exp_avg) is tracked element-wise, but the 
    second-order momentum (exp_avg_sq) is tracked column-wise for tensors 
    with dim >= 2.
    """
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-2):
        """
        Initializes the optimizer.

        Args:
            params (iterable): Iterable of parameters to optimize or dicts defining parameter groups.
            lr (float, optional): Learning rate (default: 1e-3).
            betas (Tuple[float, float], optional): Coefficients used for computing running averages 
                                                   of gradient and its square (default: (0.9, 0.999)).
            eps (float, optional): Term added to the denominator to improve numerical stability (default: 1e-8).
            weight_decay (float, optional): Weight decay coefficient (default: 1e-2).
        """
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
        if weight_decay < 0.0:
            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
        if eps < 0.0:
            raise ValueError(f"Invalid epsilon value: {eps}")

        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
        super(ColumnWiseAdamW, self).__init__(params, defaults)

    @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']
            beta1, beta2 = group['betas']
            weight_decay = group['weight_decay']
            eps = group['eps']

            for p in group['params']:
                if p.grad is None:
                    continue
                
                grad = p.grad

                # --- 1. Decoupled Weight Decay ---
                if weight_decay != 0:
                    p.add_(p, alpha=-lr * weight_decay)

                param_state = self.state[p]

                # --- 2. State Initialization ---
                if 'step' not in param_state:
                    param_state['step'] = 0
                    
                    # First moment is always element-wise
                    param_state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                    
                    # Second moment is column-wise for matrices (dim >= 2)
                    if p.dim() >= 2:
                        # By setting the 0-th dimension to 1, we keep exactly one value per column
                        # For nn.Linear weight (out_features, in_features), shape becomes (1, in_features)
                        shape = list(p.shape)
                        shape[0] = 1
                        param_state['exp_avg_sq'] = torch.zeros(shape, dtype=p.dtype, device=p.device)
                    else:
                        param_state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

                exp_avg = param_state['exp_avg']
                exp_avg_sq = param_state['exp_avg_sq']
                
                param_state['step'] += 1
                step = param_state['step']

                # --- 3. First Moment Update ---
                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)

                # --- 4. Second Moment Update ---
                if p.dim() >= 2:
                    # Column-wise formulation: [v_t]_j = beta_2 * [v_{t-1}]_j + (1 - beta_2) * ||[g_t]_j||_2^2
                    # Summing the squares over the 0-th dimension achieves the squared L2 norm for each column.
                    # Using grad * grad.conj().real supports both real and complex tensors (identical to PyTorch backend).
                    col_sq_norm = torch.sum(grad * grad.conj(), dim=0, keepdim=True).real
                    exp_avg_sq.mul_(beta2).add_(col_sq_norm, alpha=1 - beta2)
                else:
                    # Fallback to standard element-wise Adam for 1D tensors (e.g., bias)
                    exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
                
                # --- 5. Bias Correction & Step ---
                bias_correction1 = 1 - beta1 ** step
                bias_correction2 = 1 - beta2 ** step
                
                # 核心改动：对于高维张量（dim >= 2），补充 \sqrt{d_out} 以对齐 Adam 的 scale
                if p.dim() >= 2:
                    d_out = p.shape[0]
                    step_size = (lr / bias_correction1) * math.sqrt(d_out)
                else:
                    step_size = lr / bias_correction1
                
                # Thanks to keepdim=True, exp_avg_sq shape (1, in_features) naturally broadcasts 
                # against exp_avg shape (out_features, in_features) during division.
                denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
                
                p.addcdiv_(exp_avg, denom, value=-step_size)

        return loss