"""
=============================================================================
CS230 · 第 2 讲共享模块：从零实现逻辑回归（= 没有隐藏层的神经网络）
=============================================================================

对应讲义：lectures/lecture02_logistic_regression_nn_mindset.md §2–§9

-----------------------------------------------------------------------------
【这个模块里有什么】
-----------------------------------------------------------------------------
函数的名字和接口模仿 Coursera 编程作业 "Logistic Regression with a Neural Network Mindset"，
这样你做作业时可以一一对照：

    sigmoid                  σ(z) = 1 / (1 + e^{−z})
    initialize_with_zeros    w = 0（形状 (n_x, 1)），b = 0
    compute_cost             交叉熵代价 J = −(1/m) Σ [y log a + (1 − y) log(1 − a)]
    propagate                **向量化**的前向 + 反向：返回梯度 dw、db 和代价 J
    propagate_two_loops      同样的计算，用两重 for 循环（对样本、对特征）写 —— Coursera 视频里的第一版
    propagate_one_loop       去掉对特征的循环，只剩对样本的循环
    optimize                 梯度下降：重复 "propagate → 更新参数" num_iterations 次
    predict                  概率 > 0.5 判为 1
    model                    把上面几步串起来：初始化 → 训练 → 在训练集和测试集上预测

整个模型只有一个"神经元"：

        x (n_x 个像素) ──→ z = wᵀx + b ──→ a = σ(z) ──→ 预测 "是正类" 的概率 ŷ = a

【记号】（Andrew Ng 的约定，见讲义 §0.4）
    w：(n_x, 1) **列向量**；b：一个实数（Python float）
    X：(n_x, m)，每一列一个样本；Y：(1, m)
    Z = wᵀX + b：(1, m)；A = σ(Z)：(1, m)
    代码里 dw 表示 ∂J/∂w（形状和 w 相同），db 表示 ∂J/∂b，dZ 表示 ∂J/∂Z 乘以 m（见 propagate 的说明）。
    第 1 讲的 tiny_nn.py 把同一个模型写成 layer_dims = [n_x, 1]：它的 W1 形状是 (1, n_x)，正好等于这里的 wᵀ。

【注释标签】
    【思路】为什么这样写   【参数】【返回】输入输出（含形状）   【边界】什么时候会出问题
    【DL 联系】和深度学习的哪个概念有关、后面哪一讲深入
    【numpy 小知识】/【Python 小知识】第一次出现的写法

直接运行本文件会做一个小自检（梯度检验 + 三种写法结果一致，< 1 秒）：
    python3 logreg.py
"""

import numpy as np


# =============================================================================
# 1. sigmoid 与代价函数
# =============================================================================
def sigmoid(z):
    """
    sigmoid 函数 σ(z) = 1 / (1 + e^{−z})：把任意实数压到 (0, 1) 之间。

    【思路】直接写 1 / (1 + np.exp(-z)) 在 z 很负（比如 −1000）时 np.exp(1000) 会溢出成 inf 并发出警告。
            用恒等式 σ(z) = 0.5 · (1 + tanh(z / 2)) 计算，tanh 不会溢出，数学上完全相同（代码 01 演示 1 验证）。
    【参数】z：任意形状的数组（或一个数）
    【返回】形状相同的数组，值在 [0, 1]（|z| 大于约 37 时，浮点数精度让结果恰好等于 0 或 1）
    【DL 联系】逻辑回归和二分类网络的输出层都用它（第 3 讲会把它和 tanh、ReLU 放在一起比较）。
    """
    return 0.5 * (1.0 + np.tanh(0.5 * z))


def compute_cost(A, Y):
    """
    交叉熵代价（cross-entropy cost）：J = −(1/m) Σ_i [ y(i) log a(i) + (1 − y(i)) log(1 − a(i)) ]

    【思路】每个样本的损失（loss）L(a, y) = −[y log a + (1 − y) log(1 − a)]：
                y = 1 时 L = −log a      → a 越接近 1 损失越接近 0，a → 0 时损失 → ∞
                y = 0 时 L = −log(1 − a) → a 越接近 0 损失越接近 0
            代价（cost）J 是 m 个样本损失的平均。为什么用它、为什么不用平方误差，见讲义 §3。
    【参数】A：(1, m) 预测概率；Y：(1, m) 标签（0 或 1）
    【返回】float
    【边界】a 恰好等于 0 或 1 时 log 会得到 −inf。这里先把 a 裁剪到 [1e−12, 1 − 1e−12]，
            所以单个样本的损失最多约 27.6（= −log 1e−12）。
    【numpy 小知识】np.clip(A, lo, hi) 把小于 lo 的值改成 lo、大于 hi 的改成 hi；np.mean 求所有元素的平均。
    """
    P = np.clip(A, 1e-12, 1.0 - 1e-12)
    return float(-np.mean(Y * np.log(P) + (1.0 - Y) * np.log(1.0 - P)))


def initialize_with_zeros(dim):
    """
    把 w 初始化成全 0 的列向量、b 初始化成 0。

    【思路】逻辑回归的代价函数是凸的（讲义 §3.4、§4.4），没有"坏的局部最优"，从哪里出发都会朝同样的方向下降，
            所以从 0 出发最简单（代码 04 演示 5 比较了全 0 和随机两种起点）。
            **注意**：有隐藏层的网络不能这样做 —— 同一层的所有单元会永远一模一样（"对称性"问题，第 3 讲）。
    【参数】dim：w 的长度 n_x
    【返回】(w, b)：w 形状 (dim, 1)，b = 0.0
    【numpy 小知识】np.zeros((dim, 1)) 要传一个元组 (dim, 1)，得到二维的列向量；
            写成 np.zeros(dim) 会得到形状 (dim,) 的 rank-1 数组（第 1 讲 §3.5 讲过它的坑）。
    """
    w = np.zeros((dim, 1))
    b = 0.0
    return w, b


# =============================================================================
# 2. 前向 + 反向传播：三种写法
# =============================================================================
def propagate(w, b, X, Y):
    """
    **向量化**的前向传播与反向传播：一次算出所有 m 个样本的代价和梯度，没有任何 for 循环。

    【思路】讲义 §6–§7 的推导结果：
            前向：Z = wᵀX + b          形状 (1, m)，b 被广播到每一列
                  A = σ(Z)              形状 (1, m)
                  J = 交叉熵代价
            反向：dZ = A − Y            形状 (1, m)，第 i 列是第 i 个样本的 ∂L/∂z(i)
                  dw = (1/m) X dZᵀ      形状 (n_x, 1)，= 每个样本的 x(i)·dz(i) 的平均
                  db = (1/m) Σ_i dZ     一个数，= 每个样本的 dz(i) 的平均
    【参数】
        w : (n_x, 1)；b : float
        X : (n_x, m)；Y : (1, m)
    【返回】
        grads : dict，{"dw": (n_x, 1) 数组, "db": float}
        cost  : float
    【边界】用 assert 检查 dw 的形状和 w 相同 —— 形状不对时立刻报错，而不是悄悄广播出错误的结果。
    【numpy 小知识】
        w.T @ X       ：矩阵乘法（等价于 np.dot(w.T, X)）。(1, n_x) @ (n_x, m) → (1, m)
        X @ dZ.T      ：(n_x, m) @ (m, 1) → (n_x, 1)。这一步同时完成了"对每个样本乘 x(i)"和"对样本求和"
        np.sum(dZ)    ：所有元素求和，得到一个数
    【DL 联系】第 3、4 讲的每一层都是这三行公式的推广：dW[l] = (1/m) dZ[l] A[l−1]ᵀ，db[l] = (1/m) Σ dZ[l]。
    """
    m = X.shape[1]
    # 前向传播 forward
    Z = w.T @ X + b                         # (1, m)
    A = sigmoid(Z)                          # (1, m)
    cost = compute_cost(A, Y)
    # 反向传播 backward
    dZ = A - Y                              # (1, m)
    dw = X @ dZ.T / m                       # (n_x, 1)
    db = float(np.sum(dZ)) / m              # 一个数
    assert dw.shape == w.shape, f"dw 形状 {dw.shape} 和 w 形状 {w.shape} 不一致"
    return {"dw": dw, "db": db}, cost


def propagate_two_loops(w, b, X, Y):
    """
    和 propagate 完全相同的计算，但用**两重 for 循环**写：外层对样本 i，内层对特征 j。

    【思路】这是 Coursera 视频 "Gradient Descent on m Examples" 里的第一版（视频里 n_x = 2，写成 dw1、dw2；
            这里推广到任意 n_x，用内层循环代替 dw1、dw2、…）：
                J = 0, dw = 0, db = 0
                对每个样本 i：
                    z = Σ_j w_j x_j(i) + b             （内层循环 1：对特征求和）
                    a = σ(z)
                    J += −[y log a + (1 − y) log(1 − a)]
                    dz = a − y
                    对每个特征 j：dw_j += x_j(i) · dz   （内层循环 2）
                    db += dz
                J /= m, dw /= m, db /= m
            它和向量化版本的数学完全一样，只是慢得多（代码 03 演示 1 计时）。
    【参数】【返回】同 propagate
    【边界】a 同样裁剪到 [1e−12, 1 − 1e−12] 再取 log，和 compute_cost 保持一致。
    【Python 小知识】range(m) 产生 0, 1, ..., m−1；X[j, i] 取第 j 行第 i 列的一个数；
            import math 后 math.log 是 Python 自带的对数函数（对单个数比 np.log 快）。
    """
    import math
    n_x, m = X.shape
    J = 0.0
    dw = np.zeros((n_x, 1))
    db = 0.0
    for i in range(m):
        z = b
        for j in range(n_x):
            z += w[j, 0] * X[j, i]
        a = float(sigmoid(z))
        p = min(max(a, 1e-12), 1.0 - 1e-12)
        J += -(Y[0, i] * math.log(p) + (1 - Y[0, i]) * math.log(1.0 - p))
        dz = a - Y[0, i]
        for j in range(n_x):
            dw[j, 0] += X[j, i] * dz
        db += dz
    return {"dw": dw / m, "db": float(db) / m}, float(J) / m


def propagate_one_loop(w, b, X, Y):
    """
    去掉了对特征的内层循环、只保留对样本的循环的版本（Coursera 视频 "Vectorization" 里的中间一步）。

    【思路】内层的 "对每个特征 j：dw_j += x_j(i)·dz" 就是向量运算 dw += x(i)·dz（x(i) 是第 i 列）；
            内层的 "z = Σ_j w_j x_j(i) + b" 就是内积 wᵀx(i) + b。每一步仍然只处理一个样本。
    【参数】【返回】同 propagate
    【numpy 小知识】X[:, [i]] 取第 i 列并**保留二维形状** (n_x, 1)；X[:, i] 会变成一维的 (n_x,)（第 1 讲 §3.4）。
    """
    import math
    n_x, m = X.shape
    J = 0.0
    dw = np.zeros((n_x, 1))
    db = 0.0
    for i in range(m):
        x_i = X[:, [i]]                         # (n_x, 1)
        z = float((w.T @ x_i)[0, 0]) + b        # (1, n_x) @ (n_x, 1) → (1, 1)，取出那个数
        a = float(sigmoid(z))
        p = min(max(a, 1e-12), 1.0 - 1e-12)
        J += -(Y[0, i] * math.log(p) + (1 - Y[0, i]) * math.log(1.0 - p))
        dz = a - Y[0, i]
        dw += x_i * dz                          # 一次加上整列：n_x 个特征同时更新
        db += dz
    return {"dw": dw / m, "db": float(db) / m}, float(J) / m


# =============================================================================
# 3. 梯度下降
# =============================================================================
def optimize(w, b, X, Y, num_iterations=2000, learning_rate=0.005, propagate_fn=propagate):
    """
    梯度下降：重复 num_iterations 次 "算梯度 → 沿负梯度方向走一步"。

    【思路】
            for i in 1..num_iterations：
                grads, cost = propagate(w, b, X, Y)
                w ← w − α · dw
                b ← b − α · db
            每一步记录**更新前**的代价，最后再算一次更新后的代价，
            所以 costs[k] = 走了 k 步之后的代价，k = 0, 1, ..., num_iterations（共 num_iterations + 1 个数）。
            （Coursera 作业每 100 步记一次；这里每步都记，方便检查"代价有没有不降反升"。）
    【参数】
        w, b           : 初始参数（不会被修改，函数内部先复制一份）
        X, Y           : 训练数据 (n_x, m)、(1, m)
        num_iterations : 更新多少步
        learning_rate  : 学习率 α
        propagate_fn   : 用哪种写法算梯度（propagate / propagate_one_loop / propagate_two_loops），默认向量化
    【返回】
        params : {"w": (n_x, 1), "b": float}
        grads  : 最后一步的梯度 {"dw", "db"}
        costs  : list[float]，长度 num_iterations + 1
    【边界】梯度出现 inf / nan 时抛出 FloatingPointError（矩阵乘法在 BLAS 里产生的 inf / nan 不会触发 numpy 的
            浮点警告，所以要显式检查 —— 第 1 讲 README 记录过这个坑）。
    【Python 小知识】w.copy()：复制一份数组。如果不复制，w -= ... 会原地修改调用者传进来的数组。
            函数也可以当参数传进来（propagate_fn），调用时写 propagate_fn(w, b, X, Y)。
    """
    w = w.copy()
    b = float(b)
    costs = []
    grads = None
    for _ in range(num_iterations):
        grads, cost = propagate_fn(w, b, X, Y)
        if not (np.all(np.isfinite(grads["dw"])) and np.isfinite(grads["db"])):
            raise FloatingPointError("gradient is not finite")
        costs.append(cost)
        w = w - learning_rate * grads["dw"]
        b = b - learning_rate * grads["db"]
    costs.append(compute_cost(sigmoid(w.T @ X + b), Y))
    return {"w": w, "b": b}, grads, costs


def predict(w, b, X):
    """
    用学到的参数做预测：概率 a > 0.5 判为 1，否则判为 0。

    【参数】w：(n_x, 1)；b：float；X：(n_x, m)
    【返回】(1, m) 的 0.0 / 1.0 数组
    【思路】a > 0.5 ⇔ z = wᵀx + b > 0。所以逻辑回归的**决策边界**是超平面 wᵀx + b = 0 —— 一个线性分类器。
    """
    A = sigmoid(w.T @ X + b)
    return (A > 0.5).astype(float)


def accuracy(w, b, X, Y):
    """准确率：预测和标签相同的比例（0 到 1 之间的 float）。"""
    return float(np.mean(predict(w, b, X) == Y))


def model(X_train, Y_train, X_test, Y_test, num_iterations=2000, learning_rate=0.005, w0=None, b0=0.0):
    """
    把所有步骤串起来：初始化 → 梯度下降 → 预测（接口模仿 Coursera 作业的 model()）。

    【参数】
        X_train, Y_train, X_test, Y_test : 训练集和测试集
        num_iterations, learning_rate    : 梯度下降的步数和学习率
        w0, b0                           : 可选的初始参数；w0 为 None 时用 initialize_with_zeros
    【返回】dict：
        "costs"（每一步的代价）、"w"、"b"、"train_acc"、"test_acc"、"learning_rate"、"num_iterations"
    """
    if w0 is None:
        w0, b0 = initialize_with_zeros(X_train.shape[0])
    params, _, costs = optimize(w0, b0, X_train, Y_train, num_iterations, learning_rate)
    w, b = params["w"], params["b"]
    return {"costs": costs, "w": w, "b": b,
            "train_acc": accuracy(w, b, X_train, Y_train), "test_acc": accuracy(w, b, X_test, Y_test),
            "learning_rate": learning_rate, "num_iterations": num_iterations}


# =============================================================================
# 4. 数值梯度（双边差分）：用来验证手推的梯度
# =============================================================================
def numerical_gradient(w, b, X, Y, eps=1e-6):
    """
    用双边差分（two-sided difference）近似 ∂J/∂w_j 和 ∂J/∂b：
        ∂J/∂θ ≈ [J(θ + ε) − J(θ − ε)] / (2ε)

    【思路】"把一个参数往上推一点、往下推一点，看代价变了多少"。它不需要任何推导，只要会算 J，
            所以可以用来**检查**手推的梯度对不对。缺点是每个参数要算两次 J，参数多时非常慢 ——
            所以只用来检查，不用来训练。系统的"梯度检验"在第 7 讲。
    【参数】w：(n_x, 1)；b：float；X、Y：数据；eps：推动的幅度 ε
    【返回】(dw_num, db_num)：形状 (n_x, 1) 的数组和一个 float
    【Python 小知识】def J(w_, b_): ... 写在函数里面，是一个"内部函数"：它能直接用外层函数的 X、Y。
            w.copy() 复制一份再修改，原来的 w 不受影响。
    """
    def J(w_, b_):
        return compute_cost(sigmoid(w_.T @ X + b_), Y)

    dw_num = np.zeros_like(w)                  # np.zeros_like(w)：和 w 形状相同的全 0 数组
    for j in range(w.shape[0]):
        w_plus, w_minus = w.copy(), w.copy()
        w_plus[j, 0] += eps
        w_minus[j, 0] -= eps
        dw_num[j, 0] = (J(w_plus, b) - J(w_minus, b)) / (2 * eps)
    db_num = (J(w, b + eps) - J(w, b - eps)) / (2 * eps)
    return dw_num, float(db_num)


def relative_error(a, b):
    """
    两组数的相对误差 ‖a − b‖ / (‖a‖ + ‖b‖)。约 1e−7 或更小说明两者一致（第 7 讲梯度检验用同样的标准）。

    【边界】a、b 都是 0 时分母为 0，返回 0.0。
    """
    a = np.ravel(np.asarray(a, dtype=float))
    b = np.ravel(np.asarray(b, dtype=float))
    denom = np.linalg.norm(a) + np.linalg.norm(b)
    return 0.0 if denom == 0 else float(np.linalg.norm(a - b) / denom)


# =============================================================================
# 自检
# =============================================================================
if __name__ == "__main__":
    rng = np.random.default_rng(0)
    X = rng.normal(size=(5, 40))
    Y = (rng.random((1, 40)) < sigmoid(X[0:1] - X[1:2])).astype(float)
    w = rng.normal(0, 0.5, (5, 1))
    b = 0.3
    g, c = propagate(w, b, X, Y)
    g1, c1 = propagate_one_loop(w, b, X, Y)
    g2, c2 = propagate_two_loops(w, b, X, Y)
    dw_num, db_num = numerical_gradient(w, b, X, Y)
    print(f"代价：向量化 {c:.12f}，一重循环 {c1:.12f}，两重循环 {c2:.12f}")
    print(f"三种写法的 dw 最大差异：{max(np.max(np.abs(g['dw'] - g1['dw'])), np.max(np.abs(g['dw'] - g2['dw']))):.1e}")
    print(f"手推梯度 vs 双边差分的相对误差：dw {relative_error(g['dw'], dw_num):.1e}，"
          f"db {relative_error(g['db'], db_num):.1e}   （约 1e-7 或更小 = 一致）")
    d = model(X, Y, X, Y, num_iterations=500, learning_rate=0.5)
    print(f"训练 500 步：代价 {d['costs'][0]:.4f} → {d['costs'][-1]:.4f}，训练准确率 {d['train_acc']:.3f}")