"""
=============================================================================
第 5 讲 · 代码 3：优化输入，而不是参数 —— 一个最小的风格迁移，以及"网络心中的夜晚"
=============================================================================

对应讲义：lectures/lecture05_deep_learning_intuition.md 第 4 节（案例 3：Neural Style Transfer）

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
到目前为止，"训练"永远是：数据固定，用梯度下降改**参数**。神经风格迁移（Gatys et al., 2015）反过来：
**网络参数固定，用梯度下降改图像的像素。** 课件上的原话是：我们不是在通过最小化 L 学参数，而是在学一张图。

    演示 1：什么是"内容"，什么是"风格"？用一个固定的小网络提取特征：
            内容 = 某一层的激活值（"哪里有什么"）；风格 = 激活值的 Gram 矩阵（"有哪些局部图案"，不管在哪里）。
            把房子平移 6 个像素：内容表示变了很多，风格表示几乎不变。
    演示 2：风格迁移。从一张噪声图出发，只改像素，最小化 J = J_content + w · J_style：
            只用内容损失、只用风格损失、两者都用（风格权重 w = 1 / 10 / 100）。网络参数从头到尾没变过。
    演示 3：同样的"对输入求梯度"，用在一个训练好的白天 / 黑夜分类器上：
            (a) 网络心中"最像夜晚"和"最像白天"的图长什么样？
            (b) 把一张白天的照片每个像素只改一点点（肉眼几乎看不出），能不能让网络说它是夜晚？（第 18 讲预告）

-----------------------------------------------------------------------------
【我们的"预训练网络"】
-----------------------------------------------------------------------------
真正的风格迁移用一个在 ImageNet 上训练好的深层 CNN（Gatys 用的是 VGG-19）。这里没有它，
所以用一个**固定的、随机初始化的**单层卷积（8 个 3×3 滤波器 + ReLU）代替 —— 研究表明（Ustyuzhaninov et al., 2016），
随机滤波器的特征统计就已经能抓住不少纹理信息。卷积的细节留给第 12 讲，这里你只需要知道：
每个滤波器在图上每个位置看一个 3×3 的小方块，算一个加权和，输出一张"这个图案在哪里出现"的特征图。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 03_optimize_the_input.py       # 约 2 秒

需要 numpy；装了 matplotlib 会额外保存 figures/lecture05_style_transfer.png（没装就只打印字符画）。
依赖本目录的 case_tools.py、daynight_world.py，以及第 1 讲的 tiny_nn.py（演示 3 用它训练分类器）、console.py。
同一台机器上重复运行，输出逐字相同。
"""

import os

import numpy as np

from case_tools import HERE, ArrayAdam, ascii_image, luminance, pad, print_header, print_panels, sigmoid
from daynight_world import make_day_night, to_X
from tiny_nn import Standardizer, accuracy, predict, train

N = 24                 # 图像大小 24×24（灰度）
N_FILTERS = 8          # 8 个 3×3 滤波器
POOL = 4               # 内容表示：把特征图按 4×4 的块平均 → 6×6


# =============================================================================
# 两张图：内容图（一座房子）和风格图（斜条纹）
# =============================================================================
def make_images():
    """
    【返回】(content, style)，两张 (24, 24) 的灰度图，值在 [0, 1]
    【numpy 小知识】np.mgrid[0:N, 0:N] 返回两个 (N, N) 的矩阵：每个像素的行号 rr 和列号 cc。
            用它们写条件（比如 |cc − 11.5| ≤ rr − 3 是一个三角形的屋顶）就能"画"出形状。
    """
    rr, cc = np.mgrid[0:N, 0:N]
    content = np.zeros((N, N))
    content[12:21, 6:18] = 1.0                                          # 房子的墙
    content[(rr >= 4) & (rr < 12) & (np.abs(cc - 11.5) <= (rr - 3))] = 1.0   # 屋顶（三角形）
    content[15:21, 10:14] = 0.0                                         # 门
    style = (((rr + cc) // 2) % 2 == 0).astype(float)                  # 45° 的斜条纹，周期 4 像素
    return content, style


# =============================================================================
# 固定的"预训练网络"：一层 3×3 卷积 + ReLU；内容 = 池化后的特征，风格 = Gram 矩阵
# =============================================================================
FILTERS = np.random.default_rng(0).normal(0, 1, (N_FILTERS, 9))      # (8, 9)：8 个 3×3 滤波器，各拉平成 9 个数


def im2col(img):
    """
    把图像里每个位置周围的 3×3 小方块取出来排成一行（图像四周先补一圈 0，这样输出还是 24×24）。

    【参数】img：(N, N)
    【返回】(N·N, 9)：第 k 行是第 k 个位置（按行优先）的 3×3 邻域
    【DL 联系】卷积 = 在每个位置做"小方块 · 滤波器"的点积。把所有小方块排成矩阵后，
            卷积就变成了一次矩阵乘法 FILTERS @ patchesᵀ —— 这正是很多深度学习库实现卷积的方式（第 12 讲）。
    """
    padded = np.pad(img, 1)                                             # 四周补一圈 0 → (N + 2, N + 2)
    windows = np.lib.stride_tricks.sliding_window_view(padded, (3, 3))  # (N, N, 3, 3)
    return windows.reshape(N * N, 9)


def col2im(d_cols):
    """
    im2col 的"反向"：把每个小方块上的梯度加回到它来自的像素上（一个像素属于 9 个小方块，梯度要累加）。

    【参数】d_cols：(N·N, 9)，损失对每个小方块的梯度
    【返回】(N, N)，损失对每个像素的梯度
    【Python 小知识】两重 for 循环只有 3 × 3 = 9 次，每次处理一整张 (N, N) 的切片，所以很快。
    """
    d_padded = np.zeros((N + 2, N + 2))
    d = d_cols.reshape(N, N, 3, 3)
    for i in range(3):
        for j in range(3):
            d_padded[i:i + N, j:j + N] += d[:, :, i, j]
    return d_padded[1:-1, 1:-1]                                         # 去掉补的那一圈


def features(img):
    """
    前向传播：图像 → 特征图 → (内容表示, 风格表示)。

    【公式】
        F = ReLU(FILTERS · patchesᵀ)                     (8, 576)：8 张 24×24 的特征图，各拉平
        内容 C(x) = 把每张特征图按 4×4 的块平均         (8, 6, 6)：粗粒度的"哪里有什么"
        风格 S(x) = Gram 矩阵 F Fᵀ / 576                (8, 8)：  第 (i, j) 个元素 = 特征 i 和特征 j
                                                                  在所有位置上"同时出现"的平均程度
    【直觉】Gram 矩阵把"位置"求和掉了：只记录"哪些局部图案经常一起出现"，不记录它们出现在哪里。
            所以它描述的是纹理 / 笔触 / 风格，而不是物体的布局。内容表示保留了 6×6 的粗略位置。
    【返回】(cols, Z, F, content, gram)：前三个是反向传播要用的中间结果
    """
    cols = im2col(img)
    Z = FILTERS @ cols.T
    F = np.maximum(Z, 0.0)
    content = F.reshape(N_FILTERS, N // POOL, POOL, N // POOL, POOL).mean(axis=(2, 4))
    gram = F @ F.T / F.shape[1]
    return cols, Z, F, content, gram


def style_transfer_loss(img, target_content, target_gram, w_content, w_style):
    """
    风格迁移的损失和它对**像素**的梯度（网络参数 FILTERS 不动）。

    【公式】J = w_c · mean((C(x) − C_target)²) + w_s · mean((S(x) − S_target)²)
    【反向传播】从损失一路往回传到像素：
        内容项：dC → 每个 4×4 块里的 16 个位置平分（平均池化的反向）→ dF
        风格项：dS → dF = (dS + dSᵀ) F / 576          （Gram 矩阵 F Fᵀ 的导数）
        两项的 dF 相加 → dZ = dF · 1{Z > 0}（ReLU）→ d_cols = dZᵀ FILTERS → col2im → d_img
    【参数】img：(N, N)；target_content：(8, 6, 6)；target_gram：(8, 8)；w_content、w_style：两项的权重
    【返回】(J, J_content, J_style, d_img)
    """
    cols, Z, F, content, gram = features(img)
    J_c = float(np.mean((content - target_content) ** 2))
    J_s = float(np.mean((gram - target_gram) ** 2))
    d_content = w_content * 2 * (content - target_content) / content.size
    dF = np.repeat(np.repeat(d_content, POOL, axis=1), POOL, axis=2).reshape(N_FILTERS, -1) / POOL ** 2
    d_gram = w_style * 2 * (gram - target_gram) / gram.size
    dF = dF + (d_gram + d_gram.T) @ F / F.shape[1]
    dZ = dF * (Z > 0)
    d_img = col2im(dZ.T @ FILTERS)
    return w_content * J_c + w_style * J_s, J_c, J_s, d_img


def pixel_gradient_check(content, style):
    """数值梯度检验：随机图上 6 个像素的手写梯度 vs (J(x + h) − J(x − h)) / 2h，返回最大相对误差。"""
    rng = np.random.default_rng(1)
    x = rng.random((N, N))
    tc, tg = features(content)[3], features(style)[4]
    _, _, _, g = style_transfer_loss(x, tc, tg, 1.0, 1.0)
    errs = []
    for _ in range(6):
        i, j = rng.integers(0, N, 2)
        xp, xm = x.copy(), x.copy()
        xp[i, j] += 1e-6
        xm[i, j] -= 1e-6
        num = (style_transfer_loss(xp, tc, tg, 1, 1)[0] - style_transfer_loss(xm, tc, tg, 1, 1)[0]) / 2e-6
        errs.append(abs(num - g[i, j]) / (abs(num) + abs(g[i, j]) + 1e-12))
    return max(errs)


def run_style_transfer(content, style, w_content, w_style, steps=800, lr=0.02, seed=2):
    """
    从一张灰色噪声图出发，用 Adam 改像素，最小化风格迁移的损失。

    【思路】这和训练网络的循环一模一样（前向 → 损失 → 反向 → 更新），只是被更新的变量从"参数"换成了"像素"。
            每一步更新后把像素裁剪回 [0, 1]（像素值必须是合法的亮度），这叫"投影梯度下降"。
    【参数】content、style：两张目标图；w_content、w_style：权重；steps、lr：步数和学习率；seed：初始噪声的种子
    【返回】(生成的图 (N, N), 记录 list[(步数, J_content, J_style)])
    【边界】所有配置用同一张初始噪声图（同一个 seed），差别只来自损失函数。
    """
    tc, tg = features(content)[3], features(style)[4]
    img = np.clip(0.5 + 0.1 * np.random.default_rng(seed).normal(size=(N, N)), 0.0, 1.0)
    opt = ArrayAdam(img.shape, lr=lr)
    log = []
    for step in range(steps + 1):
        _, J_c, J_s, g = style_transfer_loss(img, tc, tg, w_content, w_style)
        if step in (0, 100, steps):
            log.append((step, J_c, J_s))
        if step == steps:
            break
        opt.step(img, g)
        np.clip(img, 0.0, 1.0, out=img)                 # out=img：结果直接写回 img，不新建数组
    return img, log


# =============================================================================
# 演示 1：内容 vs 风格
# =============================================================================
def rel_dist(a, b):
    """相对距离 ‖a − b‖ / ‖a‖：0 表示一模一样，1 左右表示差得和 a 本身一样大。"""
    return float(np.linalg.norm(a - b) / np.linalg.norm(a))


def demo_representations(content, style):
    print_header("演示 1：什么是\"内容\"，什么是\"风格\"？")
    shifted = np.roll(content, 6, axis=1)                    # 把房子整体右移 6 个像素
    _, _, _, c0, g0 = features(content)
    _, _, _, c1, g1 = features(shifted)
    _, _, _, c2, g2 = features(style)
    print_panels([ascii_image(content, N), ascii_image(shifted, N), ascii_image(style, N)],
                 ["房子（内容图）", "右移 6 像素的房子", "斜条纹（风格图）"])
    print("\n" + pad("比较", 30) + pad("内容表示的相对距离", 22, ">") + pad("风格表示（Gram）的相对距离", 30, ">"))
    rows = [("房子 vs 右移 6 像素的房子", rel_dist(c0, c1), rel_dist(g0, g1)),
            ("房子 vs 斜条纹", rel_dist(c0, c2), rel_dist(g0, g2))]
    for name, dc, ds in rows:
        print(pad(name, 30) + f"{dc:>22.3f}{ds:>30.3f}")
    print(f"""
观察 Observations：
  · 同一座房子平移 6 个像素：内容表示变了 {rows[0][1]:.2f}（"哪里有什么"变了），风格表示只变了 {rows[0][2]:.3f}——
    Gram 矩阵把所有位置加在一起，**不在乎图案出现在哪里**，只在乎出现了哪些图案、哪些图案经常一起出现。
  · 房子和斜条纹：两种表示都差得很远（{rows[1][1]:.2f}、{rows[1][2]:.2f}）——它们的局部图案（大块的平面 + 直边 vs 细密的斜线）完全不同。
  · 所以风格迁移的损失可以这样设计：让生成图的**内容表示**接近内容图（保住房子的布局），
    让它的 **Gram 矩阵**接近风格图（换上条纹的"笔触"）。这两个表示都来自**同一个**固定的网络。""")


# =============================================================================
# 演示 2：风格迁移
# =============================================================================
def demo_style_transfer(content, style):
    print_header("演示 2：风格迁移 —— 网络参数固定，用梯度下降改像素")
    print(f"像素梯度的数值检验：最大相对误差 {pixel_gradient_check(content, style):.2e}（很小 = 对像素的反向传播写对了）\n")
    filters_before = FILTERS.copy()
    configs = [("只用内容损失", 1.0, 0.0), ("只用风格损失", 0.0, 1.0),
               ("内容 + 风格（w = 1）", 1.0, 1.0), ("内容 + 风格（w = 10）", 1.0, 10.0),
               ("内容 + 风格（w = 100）", 1.0, 100.0)]
    results = []
    print(pad("损失函数", 26) + pad("J_content：开始 → 结束", 26, ">") + pad("J_style：开始 → 结束", 26, ">"))
    for name, wc, ws in configs:
        img, log = run_style_transfer(content, style, wc, ws)
        results.append((name, img))
        (_, c0, s0), (_, c1, s1) = log[0], log[-1]
        print(pad(name, 26) + pad(f"{c0:.4f} → {c1:.4f}", 26, ">") + pad(f"{s0:.4f} → {s1:.4f}", 26, ">"))
    unchanged = np.array_equal(filters_before, FILTERS)
    print(f"\n网络参数（8 个滤波器）在全部 5 次优化前后是否完全相同：{unchanged}")
    print("\n生成的图（800 步 Adam，学习率 0.02，都从同一张灰色噪声图出发）：")
    print_panels([ascii_image(content, N), ascii_image(style, N)] + [ascii_image(results[k][1], N) for k in (0, 1)],
                 ["内容图", "风格图", "只用内容损失", "只用风格损失"])
    print()
    print_panels([ascii_image(results[k][1], N) for k in (2, 3, 4)], [results[k][0] for k in (2, 3, 4)])
    print("""
观察 Observations：
  · **只用内容损失**：从噪声出发，重建出了房子的形状（边缘有些毛糙 —— 内容表示是 6×6 的粗粒度，细节它不管）。
  · **只用风格损失**：生成了满屏的斜条纹，但**没有房子** —— Gram 矩阵不记录位置，它只要求"斜线图案出现的统计"对得上。
    条纹的方向、粗细都对，但具体在哪里断开、在哪里拐弯是随机的（取决于初始噪声）。
  · **两者都用**：风格权重 w 决定了在"像房子"和"像条纹"之间怎样取舍。w = 1 时基本是房子；w = 10 时房子的轮廓
    还在、到处都长出了条纹；w = 100 时几乎只剩条纹。Gatys 的论文里也展示了不同权重比例下的同一组结果 ——
    **没有"正确"的权重，损失函数里的权重就是你对"想要什么"的定义。**
  · 自始至终网络参数一个都没变（上面的检查是 True）。梯度下降并不关心它在改什么：
    给它一个变量和"损失对这个变量的梯度"，它就去改那个变量。训练网络时改参数，风格迁移时改像素。""")
    return results


def save_style_figure(content, style, results, path):
    """（可选）把演示 2 的图画成一张 PNG。没装 matplotlib 时返回 False，不报错。"""
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        return False
    plt.rcParams["font.sans-serif"] = ["PingFang SC", "Heiti SC", "Arial Unicode MS", "Noto Sans CJK SC",
                                       "Microsoft YaHei", "DejaVu Sans"]
    panels = [("内容图", content), ("风格图", style)] + list(results)
    fig, axes = plt.subplots(1, len(panels), figsize=(2.0 * len(panels), 2.4), dpi=150)
    fig.patch.set_facecolor("#fcfcfb")
    for ax, (title, img) in zip(axes, panels):
        ax.imshow(img, cmap="gray", vmin=0, vmax=1, interpolation="nearest")
        ax.set_title(title, fontsize=8, color="#0b0b0b")
        ax.axis("off")
    fig.tight_layout()
    os.makedirs(os.path.dirname(path), exist_ok=True)
    fig.savefig(path, facecolor=fig.get_facecolor())
    plt.close(fig)
    return True


# =============================================================================
# 演示 3：对一个训练好的分类器"优化输入"
# =============================================================================
def logit_and_input_grad(params, sc, X):
    """
    一层隐藏层的网络（tiny_nn 训练的）输出的 logit z，以及 z 对**原始像素** X 的梯度。

    【公式】先标准化 X_s = (X − μ) / σ，然后 Z1 = W1 X_s + b1，A1 = ReLU(Z1)，z = W2 A1 + b2。
            对输入求导，链式法则一路传到最前面：
                dz/dA1 = W2ᵀ，  dz/dZ1 = W2ᵀ · 1{Z1 > 0}，  dz/dX_s = W1ᵀ dz/dZ1，  dz/dX = dz/dX_s / σ
    【参数】params：{"W1", "b1", "W2", "b2"}；sc：训练时用的 Standardizer；X：(192, m) 原始像素（8×8×3 拉平）
    【返回】(z (1, m), dz/dX (192, m))
    【DL 联系】反向传播通常停在第一层的 dW1；这里多走一步，算到 dX。可解释性（saliency map，第 21 讲）、
            对抗样本（第 18 讲）、风格迁移用的都是这一步。
    """
    Xs = (X - sc.mu) / sc.sd
    Z1 = params["W1"] @ Xs + params["b1"]
    z = params["W2"] @ np.maximum(Z1, 0.0) + params["b2"]
    dZ1 = params["W2"].T * (Z1 > 0)                    # (16, 1) * (16, m) → (16, m)：广播
    return z, (params["W1"].T @ dZ1) / sc.sd


def demo_classifier_input(seed=301):
    print_header("演示 3：同一个技巧用在分类器上 —— 网络心中的\"夜晚\"长什么样？")
    imgs, y, _ = make_day_night(5000, np.random.default_rng(seed))
    X, Y = to_X(imgs, 8), y[None, :].astype(float)
    sc = Standardizer().fit(X[:, :4000])
    params, _, _ = train(sc.transform(X[:, :4000]), Y[:, :4000], [192, 16, 1], np.random.default_rng(seed + 1),
                         lr=1e-3, steps=3000, batch_size=128, weight_decay=1e-4,
                         X_val=sc.transform(X[:, 4000:]), Y_val=Y[:, 4000:], eval_every=100)
    acc = accuracy(params, sc.transform(X[:, 4000:]), Y[:, 4000:])
    print(f"先训练一个白天 / 黑夜分类器：8×8 输入，网络 [192, 16, 1]，4000 张训练图；1000 张验证图上的准确率 {acc:.3f}")

    # 数值检验对输入的梯度
    x0 = X[:, :1].copy()
    _, g = logit_and_input_grad(params, sc, x0)
    k = 37
    xp, xm = x0.copy(), x0.copy()
    xp[k] += 1e-6
    xm[k] -= 1e-6
    num = (logit_and_input_grad(params, sc, xp)[0] - logit_and_input_grad(params, sc, xm)[0]).item() / 2e-6
    print(f"对输入的梯度，数值检验（第 {k} 个像素）：解析 {g[k, 0]:.6f}，数值 {num:.6f}\n")

    # (a) 从一张灰图出发，分别往"更像夜晚"和"更像白天"的方向改像素
    panels, titles, notes = [], [], []
    for target, sign in (("夜晚", 1.0), ("白天", -1.0)):
        x = np.full((192, 1), 0.5)
        opt = ArrayAdam(x.shape, lr=0.01)
        for _ in range(300):
            _, g = logit_and_input_grad(params, sc, x)
            opt.step(x, -sign * g)                      # 最小化 −sign·z = 最大化 sign·z
            np.clip(x, 0.0, 1.0, out=x)
        z, _ = logit_and_input_grad(params, sc, x)
        img = x.reshape(8, 8, 3)
        panels.append(ascii_image(luminance(img), 16))
        titles.append(f"最像{target}（P(夜晚) = {sigmoid(z.item()):.3f}）")
        lum = luminance(img)
        notes.append((target, lum[:2].mean(), lum[4:].mean(), lum.std(), img[:2].reshape(-1, 3).mean(axis=0)))
    print("(a) 从一张全灰的图出发，300 步 Adam 改像素（8×8 放大显示）：")
    print_panels(panels, titles)
    for target, sky, ground, contrast, rgb in notes:
        print(f"    最像{target}的图：天空（最上面两行）亮度 {sky:.2f}，地面（下半部分）亮度 {ground:.2f}，"
              f"所有像素亮度的标准差 {contrast:.2f}；天空的平均 RGB = ({rgb[0]:.2f}, {rgb[1]:.2f}, {rgb[2]:.2f})")

    # (b) 对抗样本预告：沿着梯度的符号方向，每个像素只改 ε
    day = np.flatnonzero((y[4000:] == 0) & (predict(params, sc.transform(X[:, 4000:]))[0] < 0.5)) + 4000
    _, g = logit_and_input_grad(params, sc, X[:, day])
    print(f"\n(b) 取验证集里 {len(day)} 张被正确判为\"白天\"的图，每个像素沿梯度的符号方向改 ε（x' = x + ε · sign(∂z/∂x)）：")
    print(pad("ε（像素值 0–1）", 18) + pad("约等于 0–255 里的", 20, ">") + pad("被翻成\"夜晚\"的比例", 22, ">"))
    flips = {}
    for eps in (0.01, 0.02, 0.05):
        X_adv = np.clip(X[:, day] + eps * np.sign(g), 0.0, 1.0)
        flips[eps] = float(np.mean(predict(params, sc.transform(X_adv))[0] > 0.5))
        print(pad(f"{eps}", 18) + pad(f"{eps * 255:.1f}", 20, ">") + f"{flips[eps]:>22.1%}")
    print(f"""
观察 Observations：
  · (a) "最像夜晚"和"最像白天"两张图最大的区别在天空：{notes[0][1]:.2f} vs {notes[1][1]:.2f} —— "天空暗不暗"正是这个世界里
    最强的线索（代码 01 演示 2：晴天、晴朗的夜晚在任何分辨率下都能分对）。
    但两张图都满是亮暗相间的斑点（亮度标准差 {notes[0][3]:.2f} 和 {notes[1][3]:.2f}），"最像白天"的天空还是黄白色的 ——
    **不像任何一张真实的照片**，网络却给出了接近 0 或 1 的概率。分类器只学会了"往哪边推概率"，
    没有学会"一张真实的照片长什么样"；不加约束地把一个输出推到最大，得到的往往就是这种图。
    所以用这种方法做可视化时，通常要加正则项让图保持"自然"（第 21 讲）。
  · (b) 每个像素只改 0.02（0–255 里约 5 个灰度级，肉眼几乎看不出），就有 {flips[0.02]:.0%} 的白天照片被判成夜晚；
    改 0.05 时 {flips[0.05]:.0%}。这就是**对抗样本**（adversarial examples）：对输入求梯度，既能用来"画画"，也能用来"攻击"。
    第 18 讲会解释它为什么存在、怎样防御。
  · 三个应用（风格迁移、可视化、对抗样本）的数学是同一个：**固定参数，对输入求梯度，改输入。**
    它们的区别只在于损失函数写的是什么 —— 又一次，损失函数定义了任务。""")


def main():
    content, style = make_images()
    demo_representations(content, style)
    results = demo_style_transfer(content, style)
    path = os.path.join(HERE, "figures", "lecture05_style_transfer.png")
    if save_style_figure(content, style, results, path):
        print("\n（已保存图片 figures/lecture05_style_transfer.png）")
    demo_classifier_input()


if __name__ == "__main__":
    main()