"""优化器。"""
import numpy as np


class Optimizer:
    def __init__(self, params, lr):
        self.params = list(params)
        self.lr = lr

    def zero_grad(self):
        for p in self.params:
            p.grad = None

    def step(self):
        raise NotImplementedError


class SGD(Optimizer):
    def __init__(self, params, lr=0.01, momentum=0.0):
        super().__init__(params, lr)
        self.momentum = momentum
        self.v = [np.zeros_like(p.data) for p in self.params]

    def step(self):
        for p, v in zip(self.params, self.v):
            if p.grad is None:
                continue
            v *= self.momentum
            v += p.grad
            p.data -= self.lr * v


class AdamW(Optimizer):
    """Adam + 解耦权重衰减。weight_decay=0 时就是普通 Adam。"""

    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0):
        super().__init__(params, lr)
        self.b1, self.b2 = betas
        self.eps = eps
        self.weight_decay = weight_decay
        self.t = 0
        self.m = [np.zeros_like(p.data) for p in self.params]
        self.v = [np.zeros_like(p.data) for p in self.params]

    def step(self):
        self.t += 1
        c1, c2 = 1 - self.b1 ** self.t, 1 - self.b2 ** self.t
        for p, m, v in zip(self.params, self.m, self.v):
            if p.grad is None:
                continue
            g = p.grad
            m *= self.b1
            m += (1 - self.b1) * g
            v *= self.b2
            v += (1 - self.b2) * g * g
            if self.weight_decay:
                p.data -= self.lr * self.weight_decay * p.data
            p.data -= self.lr * (m / c1) / (np.sqrt(v / c2) + self.eps)

    def state_dict(self):
        """优化器状态（步数 + 一阶/二阶动量）。和 model.state_dict() 一起存，才能精确续训。"""
        state = {"t": np.array(self.t), "lr": np.array(self.lr)}
        for i, (m, v) in enumerate(zip(self.m, self.v)):
            state[f"m.{i}"], state[f"v.{i}"] = m.copy(), v.copy()
        return state

    def load_state_dict(self, state):
        assert sum(k.startswith("m.") for k in state) == len(self.params), "参数个数和存档对不上"
        self.t, self.lr = int(state["t"]), float(state["lr"])
        for i in range(len(self.params)):
            self.m[i][...] = state[f"m.{i}"]
            self.v[i][...] = state[f"v.{i}"]


Adam = AdamW


def clip_grad_norm(params, max_norm):
    """按全局 L2 范数裁剪梯度，返回裁剪前的范数。"""
    grads = [p.grad for p in params if p.grad is not None]
    total = float(np.sqrt(sum((g.astype(np.float64) ** 2).sum() for g in grads)))
    if total > max_norm:
        scale = max_norm / (total + 1e-6)
        for p in params:
            if p.grad is not None:
                p.grad = p.grad * scale
    return total