"""
第 02 章配套代码：反向传播从零实现，训练一个 2-2-1 MLP 学会 XOR
Runnable with: numpy only.  python3 02_backprop_xor.py

这正是单层感知机做不到、而多层网络 + 反向传播能做到的事（呼应第 01 章）。
完整演示前向传播、链式法则反向求梯度、梯度下降更新。
"""
import numpy as np

rng = np.random.default_rng(1)


def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-z))


def sigmoid_grad(a):
    # 注意：输入是 sigmoid 的输出 a，导数 a*(1-a)
    return a * (1.0 - a)


# XOR 数据
X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=float)
Y = np.array([[0], [1], [1], [0]], dtype=float)

# 2-2-1 网络参数
W1 = rng.normal(0, 1, (2, 2)); b1 = np.zeros((1, 2))
W2 = rng.normal(0, 1, (2, 1)); b2 = np.zeros((1, 1))

lr = 0.5
for epoch in range(20000):
    # ---- 前向 ----
    z1 = X @ W1 + b1          # (4,2)
    a1 = sigmoid(z1)          # 隐层激活
    z2 = a1 @ W2 + b2         # (4,1)
    a2 = sigmoid(z2)          # 输出

    # ---- 损失：均方误差 ----
    loss = np.mean((a2 - Y) ** 2)

    # ---- 反向（链式法则）----
    # dL/da2 = 2(a2-Y)/N ;  da2/dz2 = sigmoid'(a2)
    dz2 = (a2 - Y) * sigmoid_grad(a2)        # (4,1)  delta_output
    dW2 = a1.T @ dz2                          # (2,1)
    db2 = dz2.sum(0, keepdims=True)
    # 误差反传到隐层：dz1 = (dz2 W2^T) * sigmoid'(a1)
    dz1 = (dz2 @ W2.T) * sigmoid_grad(a1)    # (4,2)  delta_hidden
    dW1 = X.T @ dz1                           # (2,2)
    db1 = dz1.sum(0, keepdims=True)

    # ---- 梯度下降 ----
    W2 -= lr * dW2; b2 -= lr * db2
    W1 -= lr * dW1; b1 -= lr * db1

    if epoch % 4000 == 0:
        print(f"epoch {epoch:5d}  loss={loss:.5f}")

# 最终预测
a1 = sigmoid(X @ W1 + b1)
pred = sigmoid(a1 @ W2 + b2)
print("\n最终预测 (应接近 0,1,1,0):")
for xi, pi in zip(X, pred.ravel()):
    print(f"  XOR{tuple(xi.astype(int))} -> {pi:.3f}")
