"""附录：从 Python 列表到 NumPy 数组和 PyTorch 梯度，CPU 小例子。"""

import numpy as np
import torch


def experiment():
    pages = [10, 4, 8]
    x = np.array([[10, 2], [4, 6], [8, 0]], dtype=np.float64)
    w = np.array([3, 5], dtype=np.float64)
    target = np.array([62, 60, 45], dtype=np.float64)
    prediction = x @ w + 20
    residual = prediction - target
    wrong = prediction[:, None] - target
    a = np.array([1, 2, 3])
    view = a[:2]
    independent = a[:2].copy()
    view[0] = 9

    xt, yt = torch.tensor(x), torch.tensor(target[:, None])
    model = torch.nn.Linear(2, 1, dtype=torch.float64)
    with torch.no_grad():
        model.weight.copy_(torch.tensor([[3.0, 5.0]], dtype=torch.float64))
        model.bias.fill_(20)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    loss = ((model(xt) - yt) ** 2).mean()
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    first = model.weight.grad.detach().clone()
    # 重新前向得到新计算图，但故意不清理参数上的旧梯度。
    ((model(xt) - yt) ** 2).mean().backward()
    accumulated = model.weight.grad.detach().clone()
    optimizer.zero_grad(set_to_none=True)
    ((model(xt) - yt) ** 2).mean().backward()
    bias_gradient = model.bias.grad.item()
    optimizer.step()
    model.eval()
    with torch.no_grad():
        after = ((model(xt) - yt) ** 2).mean().item()
    manual = 2 * x.T @ residual / len(x)
    np.testing.assert_allclose(first.numpy().ravel(), manual)
    torch.testing.assert_close(accumulated, 2 * first)
    return dict(list_twice=pages * 2, array_twice=(np.array(pages) * 2).tolist(),
                prediction=prediction, residual=residual, wrong=wrong,
                original=a, independent=independent, gradient=first.numpy().ravel(),
                accumulated=accumulated.numpy().ravel(), bias_gradient=bias_gradient,
                before=loss.item(), after=after)


def main():
    r = experiment()
    print("列表乘 2：", r["list_twice"])
    print("数组乘 2：", r["array_twice"])
    print("三份预测：", r["prediction"].tolist())
    print("正确残差形状：", r["residual"].shape)
    print("误广播残差形状：", r["wrong"].shape)
    print("修改切片后原数组：", r["original"].tolist())
    print("独立副本：", r["independent"].tolist())
    print("权重梯度：", [round(v, 6) for v in r["gradient"].tolist()])
    print("未清零再反传：", [round(v, 6) for v in r["accumulated"].tolist()])
    print(f'偏置梯度：{r["bias_gradient"]:.6f}')
    print(f'一步 SGD 的 MSE：{r["before"]:.6f} → {r["after"]:.6f}')


if __name__ == "__main__":
    main()