"""
优化器与学习率调度器 — Person C 负责实现

功能要求:
1. build_optimizer: 根据配置创建优化器
2. build_scheduler: 根据配置创建学习率调度器
3. InverseSqrtScheduler: 自定义 inverse square root 调度器

技术要点:
- AdamW 是 Transformer 训练的标准优化器
- Cosine with warmup 是目前最流行的调度策略
- Inverse sqrt 是经典 Transformer 论文使用的调度策略
"""

from __future__ import annotations

import math
from typing import Optional

import torch
from torch.optim import Adam, AdamW
from torch.optim.lr_scheduler import LambdaLR


def build_optimizer(model, config: dict) -> torch.optim.Optimizer:
    """
    根据配置创建优化器。

    TODO [Person C]: 实现以下逻辑:
    1. 从 config 中读取 optimizer type, lr, weight_decay, betas, eps
    2. 根据 type 创建 Adam / AdamW / Adafactor
    3. (可选) 对不同参数组设置不同学习率:
       - embedding 层可以用较小的 lr
       - LayerNorm 的 bias 不加 weight_decay
    """
    raise NotImplementedError("TODO: Person C 实现 build_optimizer")


def build_scheduler(
    optimizer: torch.optim.Optimizer,
    config: dict,
    num_training_steps: Optional[int] = None,
) -> torch.optim.lr_scheduler._LRScheduler:
    """
    根据配置创建学习率调度器。

    TODO [Person C]: 实现以下逻辑:
    1. 从 config 中读取 scheduler type, warmup_steps, min_lr
    2. type == "cosine_with_warmup":
       使用 get_cosine_schedule_with_warmup (transformers 库)
    3. type == "inverse_sqrt":
       实现经典的 lr = d_model^(-0.5) * min(step^(-0.5), step * warmup^(-1.5))
    4. type == "linear":
       使用 get_linear_schedule_with_warmup

    参考: Attention Is All You Need, Section 5.3
    """
    raise NotImplementedError("TODO: Person C 实现 build_scheduler")


class InverseSqrtScheduler(LambdaLR):
    """
    Inverse Square Root 学习率调度器。

    lr = base_lr * min(step^{-0.5}, step * warmup_steps^{-1.5})

    这是原始 Transformer 论文使用的调度策略。

    TODO [Person C]:
    1. 实现 lr_lambda 函数
    2. warmup 阶段线性增长
    3. warmup 后按 step^{-0.5} 衰减
    """

    def __init__(self, optimizer, warmup_steps: int = 4000):
        raise NotImplementedError("TODO: Person C 实现 InverseSqrtScheduler.__init__")