"""
本文件定义了模型在 configure_optimizers 方法中使用的所有优化器策略。

核心设计思想:
1. OptimizerConfig 类: 一个集中的数据类，用于保存所有与优化器相关的超参数，便于管理和传递。
2. create_* 辅助函数: 每个函数负责创建一个特定类型的优化器实例（如 AdamW, SOAP 等）。
   它们接收一个参数列表和一个 OptimizerConfig 对象。
3. strategy_* 策略函数: 每个函数定义了一种将不同优化器应用于模型不同部分（如嵌入层、注意力层等）的策略。
   它们是最终被外部调用的接口。
4. OPTIMIZER_STRATEGIES 注册表: 一个字典，将策略名称映射到对应的策略函数，以实现动态选择。
"""

import torch
from dataclasses import dataclass, field
from typing import Tuple, Dict, Any, Union

from copy import deepcopy

from optimizer.combined_optimizer import CombinedOptimizer
from optimizer.normalized_sgd import NormalizedSGD
from optimizer.normalized_p_norm_descent import NormalizedSteepestDescent
from optimizer.normalized_matrix_norm_descent import MatrixNormSteepestDescent
from optimizer.activation_preconditioned_optimizer import ActivationPreconditionedOptimizer
from optimizer.activation_preconditioned_adam import ActivationPreconditionedAdam
from optimizer.activation_preconditioned_steepest_descent import ActivationPreconditionedSteepestDescent
from optimizer.activation_double_preconditioned_adam import ActivationDoublePreconditionedAdam
from optimizer.soap import SOAP
from optimizer.asgo import ASGO
from optimizer.eigen_block_basis_optimizer import EigenBlockOptimizer
from optimizer.columnwise_adamw import ColumnWiseAdamW

# ------------------- 1. Optimizer Configuration Class -------------------
@dataclass
class OptimizerConfig:
    """
    一个集中的数据类，用于封装所有优化器相关的超参数。
    这使得配置的传递和管理更加简洁和类型安全。
    """
    # --- 通用设置 ---
    learning_rate: float = 6e-4
    weight_decay: float = 0.1
    # betas 可以是标准的2元组，也可以是某些自定义优化器所需的3元组
    betas: Union[Tuple[float, float], Tuple[float, float, float]] = (0.9, 0.95)

    # --- SGD 类优化器专属设置 ---
    sgd_learning_rate: float = 6e-4
    sgd_momentum: float = 0.9

    # --- p-norm 规范化优化器专属设置 ---
    norm_p: Union[float, str] = 2.0
    scale: float = 0.2

    # --- 行为控制布尔标志 ---
    use_nesterov: bool = False
    use_grafting: bool = False

    # --- 激活预处理优化器专属设置 ---
    eps: float = 1e-8
    update_strategy: str = 'matrix-inf'

    # eigenblock_basis_optimizer 设置 ---
    eigen_tol: float = 1e-8 
    use_adam_in_basis: bool = False # 控制是否在基变换后使用 Adam

    # --- 由 configure_optimizers 方法动态填充的额外参数 ---
    # 例如，用于启用 AdamW 的 fused 实现
    extra_args: Dict[str, Any] = field(default_factory=dict)


# ------------------- 2. Parameter Grouping Helper -------------------
def _get_decay_nodecay_params(named_parameters):
    """
    Separates parameters into two groups: one with weight decay, one without.
    RULE: Apply weight decay to 2D+ parameters, but NOT to embeddings or LayerNorms/biases.
    """
    decay_params = []
    nodecay_params = []
    for pn, p in named_parameters:
        if p.dim() >= 2 and not pn.endswith('.wte.weight') and not pn.endswith('.wpe.weight'):
            decay_params.append(p)
        else:
            nodecay_params.append(p)
    return decay_params, nodecay_params


# ------------------- 3. `create_*` Optimizer Helper Functions -------------------
# +++ MODIFIED: 所有 create_* 函数现在都接收 OptimizerConfig 对象 +++

def create_adamw_optimizer(named_parameters, config: OptimizerConfig):
    """根据配置创建 AdamW 优化器，并正确处理权重衰减。"""
    decay_params, nodecay_params = _get_decay_nodecay_params(named_parameters)
    optim_groups = [
        {'params': decay_params, 'weight_decay': config.weight_decay},
        {'params': nodecay_params, 'weight_decay': 0.0}
    ]
    # 确保只传递标准的2个beta值给 torch.optim.AdamW
    adamw_betas = (config.betas[0], config.betas[1])
    print(f"  - AdamW: {len(decay_params)} decayed param tensors, {len(nodecay_params)} non-decayed.")
    return torch.optim.AdamW(optim_groups, lr=config.learning_rate, betas=adamw_betas, **config.extra_args)

def create_columnwise_adamw_optimizer(named_parameters, config: OptimizerConfig):
    """根据配置创建 Column-wise AdamW 优化器，并正确处理权重衰减。"""
    decay_params, nodecay_params = _get_decay_nodecay_params(named_parameters)
    optim_groups = [
        {'params': decay_params, 'weight_decay': config.weight_decay},
        {'params': nodecay_params, 'weight_decay': 0.0}
    ]
    # 确保只传递标准的2个beta值给 torch.optim.AdamW
    adamw_betas = (config.betas[0], config.betas[1])
    print(f"  - Column-wise AdamW: {len(decay_params)} decayed param tensors, {len(nodecay_params)} non-decayed.")
    return ColumnWiseAdamW(optim_groups, lr=config.learning_rate, betas=adamw_betas)

def create_sgd_optimizer(named_parameters, config: OptimizerConfig):
    """根据配置创建 p-norm 规范化的 SGD 优化器。"""
    decay_params, nodecay_params = _get_decay_nodecay_params(named_parameters)
    optim_groups = [{'params': decay_params, 'weight_decay': config.weight_decay}, {'params': nodecay_params, 'weight_decay': 0.0}]
    print(f"  - NormalizedSGD: {len(decay_params)} decayed param tensors, {len(nodecay_params)} non-decayed.")
    return NormalizedSteepestDescent(optim_groups, lr=config.sgd_learning_rate, momentum=config.sgd_momentum, norm_p=config.norm_p, scale=config.scale, use_nesterov=config.use_nesterov, use_grafting=config.use_grafting)

def create_matrix_sgd_optimizer(named_parameters, config: OptimizerConfig):
    """根据配置创建矩阵范数规范化的 SGD 优化器。"""
    decay_params, nodecay_params = _get_decay_nodecay_params(named_parameters)
    optim_groups = [{'params': decay_params, 'weight_decay': config.weight_decay}, {'params': nodecay_params, 'weight_decay': 0.0}]
    print(f"  - NormalizedMatrixSGD: {len(decay_params)} decayed param tensors, {len(nodecay_params)} non-decayed.")
    return MatrixNormSteepestDescent(optim_groups, lr=config.sgd_learning_rate, momentum=config.sgd_momentum, weight_decay=config.weight_decay, norm_p=config.norm_p, scale=config.scale, use_nesterov=config.use_nesterov)

def create_soap(named_parameters, config: OptimizerConfig):
    """根据配置创建 SOAP 优化器。"""
    decay_params, nodecay_params = _get_decay_nodecay_params(named_parameters)
    optim_groups = [{'params': decay_params, 'weight_decay': config.weight_decay}, {'params': nodecay_params, 'weight_decay': 0.0}]
    print(f"  - SOAP: {len(decay_params)} decayed param tensors, {len(nodecay_params)} non-decayed.")
    # 自定义优化器可以直接接收可能包含3个值的betas
    return SOAP(optim_groups, lr=config.sgd_learning_rate, betas=config.betas, weight_decay=config.weight_decay, eps=config.eps)

def create_asgo(named_parameters, config: OptimizerConfig):
    """根据配置创建 ASGO 优化器。"""
    decay_params, nodecay_params = _get_decay_nodecay_params(named_parameters)
    optim_groups = [{'params': decay_params, 'weight_decay': config.weight_decay}, {'params': nodecay_params, 'weight_decay': 0.0}]
    print(f"  - ASGO: {len(decay_params)} decayed param tensors, {len(nodecay_params)} non-decayed.")
    return ASGO(optim_groups, lr=config.sgd_learning_rate, betas=config.betas, weight_decay=config.weight_decay, eps=config.eps, use_grafting=config.use_grafting, use_nesterov=config.use_nesterov, scale=config.scale)

def create_activation_preconditioned_optimizer(named_parameters, config: OptimizerConfig):
    decay_params, nodecay_params = _get_decay_nodecay_params(named_parameters)
    optim_groups = [{'params': decay_params, 'weight_decay': config.weight_decay}, {'params': nodecay_params, 'weight_decay': 0.0}]
    print(f"  - ActivationPreconditionedOptimizer: {len(decay_params)} decayed param tensors, {len(nodecay_params)} non-decayed.")
    return ActivationPreconditionedOptimizer(optim_groups, lr=config.sgd_learning_rate, betas=config.betas, weight_decay=config.weight_decay, scale=config.scale, eps=config.eps, use_nesterov=config.use_nesterov, use_grafting=config.use_grafting)

def create_activation_preconditioned_adam(named_parameters, config: OptimizerConfig):
    decay_params, nodecay_params = _get_decay_nodecay_params(named_parameters)
    optim_groups = [{'params': decay_params, 'weight_decay': config.weight_decay}, {'params': nodecay_params, 'weight_decay': 0.0}]
    print(f"  - ActivationPreconditionedAdam: {len(decay_params)} decayed param tensors, {len(nodecay_params)} non-decayed.")
    return ActivationPreconditionedAdam(optim_groups, lr=config.sgd_learning_rate, betas=config.betas, weight_decay=config.weight_decay, eps=config.eps)

def create_activation_preconditioned_steepest_descent(named_parameters, config: OptimizerConfig):
    decay_params, nodecay_params = _get_decay_nodecay_params(named_parameters)
    optim_groups = [{'params': decay_params, 'weight_decay': config.weight_decay}, {'params': nodecay_params, 'weight_decay': 0.0}]
    print(f"  - ActivationPreconditionedSteepestDescent: {len(decay_params)} decayed param tensors, {len(nodecay_params)} non-decayed.")
    return ActivationPreconditionedSteepestDescent(optim_groups, lr=config.sgd_learning_rate, betas=config.betas, weight_decay=config.weight_decay, eps=config.eps, scale=config.scale, use_nesterov=config.use_nesterov, use_grafting=config.use_grafting, update_strategy=config.update_strategy)

def create_activation_double_preconditioned_adam(named_parameters, config: OptimizerConfig):
    decay_params, nodecay_params = _get_decay_nodecay_params(named_parameters)
    optim_groups = [{'params': decay_params, 'weight_decay': config.weight_decay}, {'params': nodecay_params, 'weight_decay': 0.0}]
    print(f"  - ActivationDoublePreconditionedAdam: {len(decay_params)} decayed param tensors, {len(nodecay_params)} non-decayed.")
    return ActivationDoublePreconditionedAdam(optim_groups, lr=config.sgd_learning_rate, betas=config.betas, weight_decay=config.weight_decay, eps=config.eps)

# !!!! create_eigen_block_optimizer !!!!
def create_eigen_block_optimizer(named_parameters, config: OptimizerConfig):
    """Creates Eigen-Block Optimizer (SignSGD or Adam variant)."""
    decay_params, nodecay_params = _get_decay_nodecay_params(named_parameters)
    optim_groups = [
        {'params': decay_params, 'weight_decay': config.weight_decay}, 
        {'params': nodecay_params, 'weight_decay': 0.0}
    ]
    # 如果 config 中包含 use_adam_in_basis 则使用，否则默认为 False
    use_adam = getattr(config, 'use_adam_in_basis', False)
    
    print(f"  - EigenBlockOptimizer (Adam={use_adam}): {len(decay_params)} decayed params, {len(nodecay_params)} non-decayed.")
    
    return EigenBlockOptimizer(optim_groups, 
                               lr=config.sgd_learning_rate, 
                               betas=config.betas, 
                               weight_decay=config.weight_decay, 
                               eigen_tol=config.eigen_tol, 
                               eps=config.eps,
                               use_adam=use_adam)

# ------------------- 4. `strategy_*` Optimizer Strategy Functions -------------------
# +++ MODIFIED: 调用 create_* 函数变得极其简洁 +++

def strategy_adamw(model, config: OptimizerConfig):
    """策略：对所有模型参数使用 AdamW 优化器。"""
    print("Using AdamW strategy for all parameters.")
    return create_adamw_optimizer(model.named_parameters(), config)

def strategy_columnwise_adamw(model, config: OptimizerConfig):
    """策略：对所有模型参数使用 Column-wise AdamW 优化器。"""
    print("Using Column-wise AdamW strategy for all parameters.")
    return create_columnwise_adamw_optimizer(model.named_parameters(), config)

def strategy_p_sgd(model, config: OptimizerConfig):
    """策略：对所有模型参数使用 p-norm 规范化的 SGD 优化器。"""
    print(f"Using p-normed-SGD strategy for all parameters with p = {config.norm_p}")
    return create_sgd_optimizer(model.named_parameters(), config)

def strategy_matrix_p_sgd(model, config: OptimizerConfig):
    """策略：对嵌入层和1D参数使用AdamW，对其他参数使用矩阵范数SGD。"""
    print(f"Using Matrix-p-Normed SGD with p = {config.norm_p}")
    # 1. 将参数按策略分组
    adam_params = [(n, p) for n, p in model.named_parameters() if ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    sgd_params = [(n, p) for n, p in model.named_parameters() if not ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    
    # 2. 为不同组的参数创建不同的优化器
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_matrix_sgd = create_matrix_sgd_optimizer(sgd_params, config)
    
    # 3. 使用 CombinedOptimizer 将它们合并
    return CombinedOptimizer([optimizer_adam, optimizer_matrix_sgd])

def strategy_matrix_p_sgd_adam_c_attn(model, config: OptimizerConfig):
    """策略：对嵌入层和1D参数使用AdamW，对其他参数使用矩阵范数SGD。"""
    print(f"Using Matrix-p-Normed SGD with p = {config.norm_p}")
    # 1. 将参数按策略分组
    adam_params = [(n, p) for n, p in model.named_parameters() if ('.wte.' in n or '.wpe.' in n or '.c_attn.' in n or p.dim() == 1) and p.requires_grad]
    sgd_params = [(n, p) for n, p in model.named_parameters() if not ('.wte.' in n or '.wpe.' in n or '.c_attn.' in n or p.dim() == 1) and p.requires_grad]
    
    # 2. 为不同组的参数创建不同的优化器
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_matrix_sgd = create_matrix_sgd_optimizer(sgd_params, config)
    
    # 3. 使用 CombinedOptimizer 将它们合并
    return CombinedOptimizer([optimizer_adam, optimizer_matrix_sgd])

def strategy_matrix_p_sgd_2_inf(model, config: OptimizerConfig):
    """策略：对嵌入层和1D参数使用AdamW，对其他参数使用矩阵范数SGD。"""
    print(f"Using Matrix-p-Normed SGD with p = {config.norm_p}")
    # 1. 将参数按策略分组
    adam_params = [(n, p) for n, p in model.named_parameters() if (p.dim() == 1) and p.requires_grad]
    sgd_params = [(n, p) for n, p in model.named_parameters() if not (p.dim() == 1) and p.requires_grad]
    
    # 2. 为不同组的参数创建不同的优化器
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    config_emb = deepcopy(config)
    config_emb.norm_p = '2,inf'
    optimizer_matrix_sgd = create_matrix_sgd_optimizer(sgd_params, config_emb)
    
    # 3. 使用 CombinedOptimizer 将它们合并
    return CombinedOptimizer([optimizer_adam, optimizer_matrix_sgd])

def strategy_soap(model, config: OptimizerConfig):
    """策略：对嵌入层和1D参数使用AdamW，对其他参数使用SOAP。"""
    adam_params = [(n, p) for n, p in model.named_parameters() if ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    other_params = [(n, p) for n, p in model.named_parameters() if not ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_actionvation = create_soap(other_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_actionvation])

def strategy_asgo(model, config: OptimizerConfig):
    """策略：对嵌入层和1D参数使用AdamW，对其他参数使用ASGO。"""
    adam_params = [(n, p) for n, p in model.named_parameters() if ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    other_params = [(n, p) for n, p in model.named_parameters() if not ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_asgo = create_asgo(other_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_asgo])

def strategy_activation_preconditioned_sgd(model, config: OptimizerConfig):
    adam_params = [(n, p) for n, p in model.named_parameters() if ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    other_params = [(n, p) for n, p in model.named_parameters() if not ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_actionvation = create_activation_preconditioned_optimizer(other_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_actionvation])

def strategy_activation_preconditioned_adam(model, config: OptimizerConfig):
    adam_params = [(n, p) for n, p in model.named_parameters() if ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    other_params = [(n, p) for n, p in model.named_parameters() if not ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_actionvation = create_activation_preconditioned_adam(other_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_actionvation])

def strategy_activation_preconditioned_steepest_descent(model, config: OptimizerConfig):
    adam_params = [(n, p) for n, p in model.named_parameters() if ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    other_params = [(n, p) for n, p in model.named_parameters() if not ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_actionvation = create_activation_preconditioned_steepest_descent(other_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_actionvation])

def strategy_activation_double_preconditioned_adam(model, config: OptimizerConfig):
    adam_params = [(n, p) for n, p in model.named_parameters() if ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    other_params = [(n, p) for n, p in model.named_parameters() if not ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_actionvation = create_activation_double_preconditioned_adam(other_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_actionvation])

def strategy_matrix_p_sgd_with_embedding(model, config: OptimizerConfig):
    print(f"Using Matrix-p-Normed SGD with embedding as well, with p = {config.norm_p}")
    adam_params = [(n, p) for n, p in model.named_parameters() if p.dim() == 1 and p.requires_grad]
    sgd_params = [(n, p) for n, p in model.named_parameters() if not (p.dim() == 1) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_matrix_sgd = create_matrix_sgd_optimizer(sgd_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_matrix_sgd])

# add matrix-p-sgd with embedding using normalized 2,infty steepest descent
def strategy_matrix_p_sgd_with_embedding_2_infty(model, config: OptimizerConfig):
    print(f"Using Matrix-p-Normed SGD with embedding as well, with p = {config.norm_p}")
    adam_params = [(n, p) for n, p in model.named_parameters() if p.dim() == 1 and p.requires_grad]
    emb_params = [(n, p) for n, p in model.named_parameters() if ('.wte.' in n or '.wpe.' in n) and p.requires_grad]
    sgd_params = [(n, p) for n, p in model.named_parameters() if not (p.dim() == 1 or '.wte.' in n or '.wpe.' in n) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_matrix_sgd = create_matrix_sgd_optimizer(sgd_params, config)
    # use 2,infty norm for embedding parameters
    config_emb = deepcopy(config)
    config_emb.norm_p = '2,inf'
    config_emb.scale = 0.5
    optimizer_matrix_sgd_emb = create_matrix_sgd_optimizer(emb_params, config_emb)
    return CombinedOptimizer([optimizer_adam, optimizer_matrix_sgd, optimizer_matrix_sgd_emb])

def strategy_adam_with_embedding_2_infty(model, config: OptimizerConfig):
    print(f"Using Adam with embedding as using Matrix-p-sgd with p = {config.norm_p}")
    adam_params = [(n, p) for n, p in model.named_parameters() if not ('.wte.' in n or '.wpe.' in n) and p.requires_grad]
    emb_params = [(n, p) for n, p in model.named_parameters() if ('.wte.' in n or '.wpe.' in n) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    # use 2,infty norm for embedding parameters
    config_emb = deepcopy(config)
    config_emb.norm_p = '2,inf'
    config_emb.scale = 0.5
    optimizer_matrix_sgd_emb = create_matrix_sgd_optimizer(emb_params, config_emb)
    return CombinedOptimizer([optimizer_adam, optimizer_matrix_sgd_emb])

# define strategy for mlp.c_proj use 2_infty, others use p-normed sgd or adamw
def strategy_matrix_p_sgd_with_mlp_c_proj_2_infty(model, config: OptimizerConfig):
    print(f"Using 1,2-Normed SGD in mlp.c_proj, with p = {config.norm_p}")
    adam_params = [(n, p) for n, p in model.named_parameters() if (p.dim() == 1 or '.wte.' in n or '.wpe.' in n) and p.requires_grad]
    mlp_c_proj_params = [(n, p) for n, p in model.named_parameters() if ('mlp.c_proj' in n) and p.requires_grad]
    sgd_params = [(n, p) for n, p in model.named_parameters() if not (p.dim() == 1 or '.wte.' in n or '.wpe.' in n or 'mlp.c_proj' in n) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_matrix_sgd = create_matrix_sgd_optimizer(sgd_params, config)
    # use 2,infty norm for embedding parameters
    config_emb = deepcopy(config)
    config_emb.norm_p = '2,inf'
    # config_emb.scale = 0.5
    optimizer_matrix_sgd_emb = create_matrix_sgd_optimizer(mlp_c_proj_params, config_emb)
    return CombinedOptimizer([optimizer_adam, optimizer_matrix_sgd, optimizer_matrix_sgd_emb])

def strategy_adam_with_mlp_c_proj_2_infty(model, config: OptimizerConfig):
    print(f"Using Adam with embedding as using Matrix-p-sgd with p = {config.norm_p}")
    adam_params = [(n, p) for n, p in model.named_parameters() if not ('mlp.c_proj' in n) and p.requires_grad]
    mlp_c_proj_params = [(n, p) for n, p in model.named_parameters() if ('mlp.c_proj' in n) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    # use 2,infty norm for embedding parameters
    config_emb = deepcopy(config)
    config_emb.norm_p = '2,inf'
    # config_emb.scale = 0.5
    optimizer_matrix_sgd_emb = create_matrix_sgd_optimizer(mlp_c_proj_params, config_emb)
    return CombinedOptimizer([optimizer_adam, optimizer_matrix_sgd_emb])

def strategy_mixed_attn_adam_others_sgd(model, config: OptimizerConfig):
    print("Using mixed strategy: AdamW for Attention, SGD for others.")
    adam_params = [(n, p) for n, p in model.named_parameters() if '.attn.' in n and p.requires_grad]
    sgd_params = [(n, p) for n, p in model.named_parameters() if '.attn.' not in n and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_other = create_sgd_optimizer(sgd_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_other])

def strategy_mixed_attn_emb_adam_others_sgd(model, config: OptimizerConfig):
    print("Using mixed strategy: AdamW for Attention and Embedding, SGD for others.")
    adam_params = [(n, p) for n, p in model.named_parameters() if ('.attn.' in n or '.wte.' in n or '.wpe.' in n) and p.requires_grad]
    sgd_params = [(n, p) for n, p in model.named_parameters() if not ('.attn.' in n or '.wte.' in n or '.wpe.' in n) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_other = create_sgd_optimizer(sgd_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_other])

def strategy_mixed_attn_emb_adam_others_matrix_sgd(model, config: OptimizerConfig):
    print("Using mixed strategy: AdamW for Attention and Embedding, SGD for others.")
    adam_params = [(n, p) for n, p in model.named_parameters() if ('.attn.' in n or '.wte.' in n or '.wpe.' in n) and ('proj' not in n) and p.requires_grad]
    sgd_params = [(n, p) for n, p in model.named_parameters() if not ('.attn.' in n or '.wte.' in n or '.wpe.' in n) and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_other = create_matrix_sgd_optimizer(sgd_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_other])

def strategy_mixed_attn_sgd_others_adam(model, config: OptimizerConfig):
    print("Using mixed strategy: SGD for Attention, AdamW for others.")
    sgd_params = [(n, p) for n, p in model.named_parameters() if '.attn.' in n and p.requires_grad]
    adam_params = [(n, p) for n, p in model.named_parameters() if '.attn.' not in n and p.requires_grad]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_other = create_sgd_optimizer(sgd_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_other])

def strategy_mixed_emb_adam_others_sgd(model, config: OptimizerConfig):
    print("Using mixed strategy: AdamW for embedding, SGD for others.")
    adam_params = [(n, p) for n, p in model.named_parameters() if p.requires_grad and ('.wte.' in n or '.wpe.' in n)]
    sgd_params = [(n, p) for n, p in model.named_parameters() if p.requires_grad and not ('.wte.' in n or '.wpe.' in n)]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_other = create_sgd_optimizer(sgd_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_other])

def strategy_mixed_emb_sgd_others_adam(model, config: OptimizerConfig):
    print("Using mixed strategy: AdamW for embedding, SGD for others.")
    sgd_params = [(n, p) for n, p in model.named_parameters() if p.requires_grad and ('.wte.' in n or '.wpe.' in n)]
    adam_params = [(n, p) for n, p in model.named_parameters() if p.requires_grad and not ('.wte.' in n or '.wpe.' in n)]
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_other = create_sgd_optimizer(sgd_params, config)
    return CombinedOptimizer([optimizer_adam, optimizer_other])

# strategies for eigen-block optimizer with signSGD for 2D weights, AdamW for others

def strategy_eigen_block_sign_sgd(model, config: OptimizerConfig):
    """
    Strategy: 
    - 2D weights: EigenBlockOptimizer with SignSGD (use_adam=False)
    - Others: AdamW
    """
    config.use_adam_in_basis = False # 强制 SignSGD 模式
    
    adam_params = [(n, p) for n, p in model.named_parameters() if ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    eigen_params = [(n, p) for n, p in model.named_parameters() if not ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_eigen = create_eigen_block_optimizer(eigen_params, config)
    
    return CombinedOptimizer([optimizer_adam, optimizer_eigen])

def strategy_eigen_block_adam(model, config: OptimizerConfig):
    """
    Strategy: 
    - 2D weights: EigenBlockOptimizer with Adam (use_adam=True)
    - Others: AdamW
    """
    config.use_adam_in_basis = True # 强制 Adam 模式
    
    adam_params = [(n, p) for n, p in model.named_parameters() if ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    eigen_params = [(n, p) for n, p in model.named_parameters() if not ('.wte.' in n or '.wpe.' in n or p.dim() == 1) and p.requires_grad]
    
    optimizer_adam = create_adamw_optimizer(adam_params, config)
    optimizer_eigen = create_eigen_block_optimizer(eigen_params, config)
    
    return CombinedOptimizer([optimizer_adam, optimizer_eigen])

# ------------------- 5. Optimizer Strategy Registry -------------------
OPTIMIZER_STRATEGIES = {
    "adamw": strategy_adamw,
    "columnwise-adamw": strategy_columnwise_adamw,
    "p-sgd": strategy_p_sgd,
    "matrix-p-sgd": strategy_matrix_p_sgd,
    "matrix-p-sgd-adam-c-attn": strategy_matrix_p_sgd_adam_c_attn,
    "matrix-p-sgd-2-inf": strategy_matrix_p_sgd_2_inf,
    "matrix-p-sgd-with-embedding": strategy_matrix_p_sgd_with_embedding,
    "matrix-p-sgd-with-embedding-2-inf": strategy_matrix_p_sgd_with_embedding_2_infty,
    "matrix-p-sgd-with-mlp-c-proj-2-inf": strategy_matrix_p_sgd_with_mlp_c_proj_2_infty,
    "adam-with-embedding-2-inf": strategy_adam_with_embedding_2_infty,
    "adam-with-mlp-c-proj-2-inf": strategy_adam_with_mlp_c_proj_2_infty,
    "attention-adamw-others-p-sgd": strategy_mixed_attn_adam_others_sgd,
    "attention-embedding-adamw-others-p-sgd": strategy_mixed_attn_emb_adam_others_sgd,
    "matrix-p-sgd-attn-adamw": strategy_mixed_attn_emb_adam_others_matrix_sgd,
    "attention-p-sgd-others-adamw": strategy_mixed_attn_sgd_others_adam,
    "embedding-adamw-others-p-sgd": strategy_mixed_emb_adam_others_sgd,
    "embedding-p-sgd-others-adamw": strategy_mixed_emb_sgd_others_adam,
    "activation-preconditioned-sgd": strategy_activation_preconditioned_sgd,
    "activation-preconditioned-adam": strategy_activation_preconditioned_adam,
    "activation-preconditioned-steepest-descent": strategy_activation_preconditioned_steepest_descent,
    "activation-double-preconditioned-adam": strategy_activation_double_preconditioned_adam,
    "soap": strategy_soap,
    "asgo": strategy_asgo,
    "eigen-block-signsgd": strategy_eigen_block_sign_sgd,
    "eigen-block-adam": strategy_eigen_block_adam,
}