"""
=============================================================================
第 5 讲 · 代码 2：人脸验证 —— 为什么要学一个"编码"，以及 triplet loss
=============================================================================

对应讲义：lectures/lecture05_deep_learning_intuition.md 第 3 节（案例 2：Face Verification）

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
场景（来自 CS230 课件）：学校想在食堂、体育馆、游泳池门口用人脸验证学生证。门口的摄像头拍一张照片，
和学生证上的照片比较，输出 y = 1（是本人）或 y = 0（不是）。

    演示 1：最简单的办法 —— 直接比较两张照片的像素距离，小于阈值就算同一个人。为什么不行？
            看看像素的变化有多少来自"是谁"，多少来自光照、姿态、表情。
    演示 2：学一个编码器 f(x)，把一张脸变成一个 16 维的向量（encoding / embedding），用 triplet loss 训练：
                L = max(0, ‖f(A) − f(P)‖² − ‖f(A) − f(N)‖² + α)
            A = anchor（某个人的一张照片），P = positive（同一个人的另一张），N = negative（另一个人）。
            先做梯度检验，再训练；然后在**训练时从没见过的 100 个新同学**上测：
            验证（是不是同一个人）和 one-shot 识别（每人只有一张注册照片，认出是谁）。
    演示 3：损失函数的每一项都有用吗？去掉 margin（α = 0）；只"拉近同一个人"、不"推开不同的人"。
    演示 4：没有任何标签时：用编码做人脸聚类（k-means），把手机相册里同一个人的照片归到一起。

-----------------------------------------------------------------------------
【一个"知道真相"的玩具人脸世界】
-----------------------------------------------------------------------------
每个人有一个 6 维的"身份向量" z（脸型、五官……），每张照片还有一个 6 维的"干扰向量" u（光照、姿态、表情……），
每张照片都重新抽。照片的 100 个"像素"由一个固定的、随机的非线性函数生成：
        x = C · tanh(A z + B u) + 噪声
干扰的强度（B 的尺度 1.5）比身份（A 的尺度 1.0）大 —— 这是第一次写这个世界时定下的数，没有为了结果调过。
模型只看得到 x，看不到 z 和 u。训练集：300 个人 × 每人 8 张；测试：另外 100 个从没见过的人 × 每人 6 张。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 02_face_verification_triplet.py     # 约 5 秒

需要 numpy。依赖本目录的 case_tools.py，以及第 1 讲的 tiny_nn.py（只用它的 Adam 优化器）、console.py。
同一台机器上重复运行，输出逐字相同。
"""

import numpy as np

from case_tools import pad, print_header
from tiny_nn import Adam

Q, R, H_GEN, D = 6, 6, 48, 100       # 身份维数、干扰维数、生成器的隐藏宽度、像素个数
NUISANCE_SCALE = 1.5                  # 干扰（光照、姿态……）相对身份的强度
N_TRAIN_PEOPLE, PHOTOS_TRAIN = 300, 8
N_TEST_PEOPLE, PHOTOS_TEST = 100, 6
H_ENC, E_DIM = 64, 16                 # 编码器：100 → 64（ReLU）→ 16 → 归一化到单位长度
MARGIN = 0.2                          # triplet loss 的 margin α（FaceNet 论文用的就是 0.2）


# =============================================================================
# 玩具人脸世界
# =============================================================================
def make_generator(seed=0):
    """
    固定的"照相机"：把 (身份 z, 干扰 u) 变成 100 个像素的随机非线性函数。

    【返回】dict：A (48, 6)、B (48, 6)、C (100, 48)
    【思路】A、B 的尺度决定了身份和干扰各自对像素的影响有多大；B 乘了 NUISANCE_SCALE = 1.5。
            除以 √(输入维数) 是为了让 A z、B u 的每个分量方差大约是 1（和第 1 讲 He / Xavier 初始化同一个道理）。
    """
    rng = np.random.default_rng(seed)
    return {"A": rng.normal(0, 1 / np.sqrt(Q), (H_GEN, Q)),
            "B": rng.normal(0, NUISANCE_SCALE / np.sqrt(R), (H_GEN, R)),
            "C": rng.normal(0, 1 / np.sqrt(H_GEN), (D, H_GEN))}


def take_photos(G, Z, rng, nuisance=True):
    """
    给每一列身份向量拍一张照片。

    【参数】G：make_generator 的返回值；Z：(6, m) 身份向量（同一个人可以出现多次 = 拍多张）；
            nuisance：False 时干扰向量全为 0（"标准光照、正脸"，只用在演示 1 的方差分解里）
    【返回】X：(100, m)，每一列是一张照片
    """
    U = rng.normal(size=(R, Z.shape[1])) if nuisance else np.zeros((R, Z.shape[1]))
    return G["C"] @ np.tanh(G["A"] @ Z + G["B"] @ U) + 0.05 * rng.normal(size=(D, Z.shape[1]))


def make_people(G, n_people, photos_each, rng):
    """
    生成 n_people 个人、每人 photos_each 张照片。

    【返回】X (100, n·k)、ids (n·k,)（第几个人）、Z (6, n)（每个人真实的身份向量，只用来算"上限"）
    【numpy 小知识】np.repeat([0, 1, 2], 2) → [0, 0, 1, 1, 2, 2]：每个人的照片排在一起。
    """
    Z = rng.normal(size=(Q, n_people))
    ids = np.repeat(np.arange(n_people), photos_each)
    return take_photos(G, Z[:, ids], rng), ids, Z


def make_world():
    """生成训练人群（300 人）和测试人群（100 个新同学），并用训练集的统计量标准化像素。"""
    G = make_generator(0)
    X_tr, id_tr, Z_tr = make_people(G, N_TRAIN_PEOPLE, PHOTOS_TRAIN, np.random.default_rng(1))
    X_te, id_te, Z_te = make_people(G, N_TEST_PEOPLE, PHOTOS_TEST, np.random.default_rng(2))
    mu, sd = X_tr.mean(axis=1, keepdims=True), X_tr.std()
    return {"G": G, "X_tr": (X_tr - mu) / sd, "id_tr": id_tr, "Z_tr": Z_tr,
            "X_te": (X_te - mu) / sd, "id_te": id_te, "Z_te": Z_te, "mu": mu, "sd": sd}


# =============================================================================
# 距离与评估
# =============================================================================
def pair_sq_dists(E1, E2):
    """
    两组向量之间所有配对的平方欧氏距离。

    【思路】‖a − b‖² = ‖a‖² + ‖b‖² − 2 aᵀb，一次矩阵乘法算出所有配对，不用两重循环。
    【参数】E1：(d, m1)；E2：(d, m2)（每一列是一个向量）
    【返回】(m1, m2)，第 (i, j) 个元素是 E1 第 i 列和 E2 第 j 列的平方距离
    【边界】浮点误差可能让结果是一个很小的负数（−1e−15），用 np.maximum(·, 0) 截掉。
    """
    sq = np.maximum((E1 ** 2).sum(axis=0)[:, None] + (E2 ** 2).sum(axis=0)[None, :] - 2 * E1.T @ E2, 0.0)
    return sq


def sample_pairs(ids, rng, n=2000):
    """
    随机抽 n 对"同一个人"和 n 对"不同的人"，返回四个下标数组 (a_same, b_same, a_diff, b_diff)。

    【思路】同一个人：先抽一张照片 a，再从同一个人的其他照片里抽 b；不同的人：从其他人的照片里抽 b。
    【边界】要求每个人至少有 2 张照片。
    """
    m = len(ids)
    a = rng.integers(0, m, n)
    same_b = np.array([rng.choice(np.flatnonzero((ids == ids[i]) & (np.arange(m) != i))) for i in a])
    a2 = rng.integers(0, m, n)
    diff_b = np.array([rng.choice(np.flatnonzero(ids != ids[i])) for i in a2])
    return a, same_b, a2, diff_b


def pair_distances(E, pairs):
    """给定编码 E (d, m) 和 sample_pairs 的结果，返回 (同一个人的距离 (n,), 不同人的距离 (n,))。"""
    a, b, a2, b2 = pairs
    return ((E[:, a] - E[:, b]) ** 2).sum(axis=0), ((E[:, a2] - E[:, b2]) ** 2).sum(axis=0)


def best_threshold(d_same, d_diff):
    """
    在一组"同一个人 / 不同人"的距离上，找让平衡准确率（两类各自准确率的平均）最高的阈值。

    【思路】把所有距离排序，逐个当作候选阈值：距离 ≤ 阈值判"同一个人"。
    【返回】(阈值, 这个阈值下的平衡准确率)
    """
    cand = np.sort(np.concatenate([d_same, d_diff]))
    accs = [(np.mean(d_same <= t) + np.mean(d_diff > t)) / 2 for t in cand]
    k = int(np.argmax(accs))
    return float(cand[k]), float(accs[k])


def evaluate(E_tr, id_tr, E_te, id_te):
    """
    一套完整的评估（像素、编码、真实身份向量都用它，比较才公平）：

        1. 阈值：在**训练人群**的配对上选（这些人模型训练时见过）；
        2. 验证准确率：把这个阈值原封不动地用到**新同学**的配对上（平衡准确率）；
        3. one-shot 识别：每个新同学的第 1 张照片当"学生证照片"（注册库），其余照片去库里找最近的，
           top-1 准确率（100 个人里猜中的概率是 1%）。
        4. 新同学上"不同人的平均距离 / 同一个人的平均距离"：越大说明编码把人分得越开。

    【参数】E_tr：(d, m_tr) 训练人群的表示；E_te：(d, m_te) 新同学的表示；id_tr、id_te：每张照片是谁
    【返回】dict：threshold、verify_train、verify_test、identify、ratio
    """
    d_s, d_d = pair_distances(E_tr, sample_pairs(id_tr, np.random.default_rng(10)))
    thr, acc_tr = best_threshold(d_s, d_d)
    d_s, d_d = pair_distances(E_te, sample_pairs(id_te, np.random.default_rng(11)))
    acc_te = (np.mean(d_s <= thr) + np.mean(d_d > thr)) / 2
    first = np.array([np.flatnonzero(id_te == p)[0] for p in range(id_te.max() + 1)])     # 每人第 1 张：注册照片
    queries = np.setdiff1d(np.arange(len(id_te)), first)                                   # 其余照片：门口拍的
    nearest = np.argmin(pair_sq_dists(E_te[:, queries], E_te[:, first]), axis=1)
    ratio = float(d_d.mean() / d_s.mean()) if d_s.mean() > 0 else float("inf")    # 真实 z：同一个人距离为 0 → 无穷大
    return {"threshold": thr, "verify_train": acc_tr, "verify_test": float(acc_te),
            "identify": float(np.mean(nearest == id_te[queries])), "ratio": ratio}


# =============================================================================
# 编码器与 triplet loss（前向、反向都手写）
# =============================================================================
def init_encoder(rng, d_in=D, h=H_ENC, e=E_DIM):
    """
    编码器 f(x) = normalize(W2 · ReLU(W1 x + b1) + b2)：一层隐藏层，输出被归一化到单位长度。

    【参数】rng；d_in、h、e：输入、隐藏、编码的维数
    【返回】参数字典 {"W1": (h, d_in), "b1": (h, 1), "W2": (e, h), "b2": (e, 1)}（He 初始化，和第 1 讲一样）
    """
    return {"W1": rng.normal(0, np.sqrt(2 / d_in), (h, d_in)), "b1": np.zeros((h, 1)),
            "W2": rng.normal(0, np.sqrt(1 / h), (e, h)), "b2": np.zeros((e, 1))}


def encode(p, X):
    """
    前向传播：X (d_in, m) → 编码 E (e, m)，每一列长度为 1。

    【思路】为什么要归一化到单位长度（FaceNet 也这么做）？如果不归一化，网络可以把所有编码一起放大 c 倍，
            距离就放大 c² 倍，固定的 margin α 相对来说就越来越小 —— 损失可以靠"放大"而不是"分得更开"降下来。
            归一化之后所有编码都在单位球面上，任意两点的平方距离在 [0, 4] 之间，α = 0.2 才有固定的意义。
    【返回】(E, cache)：cache 里存反向传播要用的中间结果
    """
    Z1 = p["W1"] @ X + p["b1"]
    A1 = np.maximum(Z1, 0.0)
    F = p["W2"] @ A1 + p["b2"]
    norm = np.sqrt((F ** 2).sum(axis=0, keepdims=True)) + 1e-12
    return F / norm, (X, Z1, A1, F, norm)


def encoder_backward(p, cache, dE):
    """
    反向传播：已知损失对编码 E 的梯度 dE (e, m)，求损失对 W1、b1、W2、b2 的梯度。

    【思路】链式法则，从后往前三步：
        1. 穿过归一化 e = F / ‖F‖：   dF = (dE − e · (eᵀ dE)) / ‖F‖
           （直觉：沿着 e 本身方向的那部分梯度只会改变长度，而长度被归一化抵消了，所以要减掉）
        2. 穿过第二层 F = W2 A1 + b2：  dW2 = dF A1ᵀ，db2 = Σ dF，dA1 = W2ᵀ dF
        3. 穿过 ReLU 和第一层：          dZ1 = dA1 · 1{Z1 > 0}，dW1 = dZ1 Xᵀ，db1 = Σ dZ1
    【返回】grads 字典，键是 "dW1" 等（和 tiny_nn.Adam 的约定一样）
    【DL 联系】第 2–4 讲会把第 2、3 步推导清楚；第 1 步是本讲新加的一块"积木"。
            框架（PyTorch，第 9 讲起）会自动做这些，但手写一次能看清 triplet loss 的梯度到底推了谁。
    """
    X, Z1, A1, F, norm = cache
    e = F / norm
    dF = (dE - e * (e * dE).sum(axis=0, keepdims=True)) / norm
    grads = {"dW2": dF @ A1.T, "db2": dF.sum(axis=1, keepdims=True)}
    dZ1 = (p["W2"].T @ dF) * (Z1 > 0)
    grads["dW1"] = dZ1 @ X.T
    grads["db1"] = dZ1.sum(axis=1, keepdims=True)
    return grads


def triplet_loss_and_grads(Ea, Ep, En, alpha=MARGIN, mode="triplet"):
    """
    triplet loss 和它对三组编码的梯度。

    【公式】对每个三元组 i：  ℓ_i = max(0, ‖a − p‖² − ‖a − n‖² + α)，  J = (1/B) Σ ℓ_i
            只有 ℓ_i > 0 的三元组（"还没分开够"的）有梯度：
                ∂ℓ/∂a = 2(n − p)，   ∂ℓ/∂p = −2(a − p)，   ∂ℓ/∂n = 2(a − n)
            直觉：梯度下降沿着负梯度走 —— p 被拉向 a，n 被推离 a，a 往"离 p 近、离 n 远"的方向挪。
    【参数】Ea、Ep、En：(e, B) 三组编码；alpha：margin；
            mode："triplet"（完整的 triplet loss）或 "pull"（只有 ‖a − p‖² 一项，演示 3 的反面例子）
    【返回】(J, 有效三元组的比例, dEa, dEp, dEn)
    """
    B = Ea.shape[1]
    d_ap = ((Ea - Ep) ** 2).sum(axis=0)
    if mode == "pull":
        return float(d_ap.mean()), 1.0, 2 * (Ea - Ep) / B, -2 * (Ea - Ep) / B, np.zeros_like(En)
    d_an = ((Ea - En) ** 2).sum(axis=0)
    loss = np.maximum(d_ap - d_an + alpha, 0.0)
    active = (loss > 0).astype(float)                   # (B,)：哪些三元组还有梯度
    dEa = 2 * (En - Ep) * active / B
    dEp = -2 * (Ea - Ep) * active / B
    dEn = 2 * (Ea - En) * active / B
    return float(loss.mean()), float(active.mean()), dEa, dEp, dEn


def sample_triplets(ids, photos_each, rng, B):
    """
    随机抽 B 个三元组 (anchor, positive, negative) 的下标。

    【思路】照片按人排好（第 k 个人的照片下标是 k·photos_each 到 (k+1)·photos_each − 1），所以：
            positive = 同一个人的另一张：在"第几张"上加 1 到 photos_each − 1 再取余；
            negative = 另一个人：在"第几个人"上加 1 到 n_people − 1 再取余，张数随便抽。
            这是最朴素的"随机三元组"。FaceNet 用的是更难的"semi-hard"三元组（讲义 §3.4）。
    【返回】三个 (B,) 的下标数组
    """
    a = rng.integers(0, len(ids), B)
    person, k = a // photos_each, a % photos_each
    p = person * photos_each + (k + rng.integers(1, photos_each, B)) % photos_each
    n = random_negatives(a, ids, photos_each, rng, (B,))
    return a, p, n


def random_negatives(a, ids, photos_each, rng, shape):
    """
    给每个 anchor 随机抽"另一个人"的照片作为负样本。

    【参数】a：(B,) anchor 下标；shape：(B,) 每个 anchor 抽 1 个，或 (B, K) 每个 anchor 抽 K 个候选
    【返回】形状为 shape 的下标数组，保证和对应的 anchor 不是同一个人
    【numpy 小知识】a[:, None] 把 (B,) 变成 (B, 1)，和 (B, K) 广播，每一行用同一个 anchor。
    """
    n_people = ids.max() + 1
    person = a // photos_each
    if len(shape) == 2:
        person = person[:, None]
    other = (person + rng.integers(1, n_people, shape)) % n_people      # 加 1 到 n_people − 1 再取余：一定是别人
    return other * photos_each + rng.integers(0, photos_each, shape)


def pick_semihard(params, X, a, p, cand):
    """
    FaceNet 的 semi-hard 负样本挖掘：给每个 anchor 从 K 个候选负样本里挑一个"有用"的。

    【思路】随机三元组大多太容易（负样本本来就离得很远，损失为 0，没有梯度）。用当前的编码器算一下距离，
            在"比正样本远、但又没远出 margin"的候选里挑最近的那个（semi-hard）；
            如果一个都没有，就挑最近的那个负样本（hard）。
    【参数】params：当前编码器；X：(D, m) 全部训练照片；a、p：(B,) anchor 和 positive 的下标；
            cand：(B, K) 每个 anchor 的 K 个候选负样本（都和 anchor 不是同一个人）
    【返回】(B,) 选中的负样本下标
    【numpy 小知识】np.where(条件, A, B) 逐元素地"条件成立取 A，否则取 B"；
            cand[np.arange(B), j] 从每一行里各取一个元素（第 i 行取第 j[i] 列）。
    """
    B, K = cand.shape
    Ea, Ep = encode(params, X[:, a])[0], encode(params, X[:, p])[0]
    Ec = encode(params, X[:, cand.ravel()])[0].reshape(-1, B, K)           # (e, B, K)
    d_ap = ((Ea - Ep) ** 2).sum(axis=0)                                    # (B,)
    d_ac = ((Ea[:, :, None] - Ec) ** 2).sum(axis=0)                        # (B, K)
    semi = d_ac > d_ap[:, None]                                            # 比正样本远的候选
    j = np.where(semi.any(axis=1), np.argmin(np.where(semi, d_ac, np.inf), axis=1), np.argmin(d_ac, axis=1))
    return cand[np.arange(B), j]


def train_encoder(world, seed=3, steps=3000, B=128, alpha=MARGIN, mode="triplet", mining="random", log_at=()):
    """
    用 triplet loss 训练编码器。

    【思路】每一步：抽 B 个三元组 → 三组照片拼在一起做一次前向 → 算损失和对编码的梯度 → 一次反向 → Adam 更新。
    【参数】world：make_world() 的返回值；seed：初始化和抽样的种子；steps、B：步数和批大小；
            alpha、mode：见 triplet_loss_and_grads；mining："random"（随机负样本）或 "semihard"（见 pick_semihard，
            每个 anchor 从 10 个候选里挑）；log_at：在哪些步打印训练损失
    【返回】(训练好的参数, 最后 100 步的平均损失)
    【Python 小知识】np.concatenate([a, b, c], axis=1) 把三个 (d, B) 的矩阵左右拼成 (d, 3B)。
    """
    rng = np.random.default_rng(seed)
    params = init_encoder(rng)
    opt = Adam(params, lr=1e-3)
    X, ids = world["X_tr"], world["id_tr"]
    recent = []
    for step in range(steps + 1):
        a, p, n = sample_triplets(ids, PHOTOS_TRAIN, rng, B)
        if mining == "semihard":
            cand = random_negatives(a, ids, PHOTOS_TRAIN, rng, (B, 10))     # 每个 anchor 10 个候选负样本
            n = pick_semihard(params, X, a, p, cand)
        E, cache = encode(params, X[:, np.concatenate([a, p, n])])
        J, frac, dEa, dEp, dEn = triplet_loss_and_grads(E[:, :B], E[:, B:2 * B], E[:, 2 * B:], alpha, mode)
        recent = (recent + [J])[-100:]
        if step in log_at:
            print(f"    第 {step:>4} 步：这一批的平均损失 {J:.4f}，仍有梯度的三元组 {frac:.1%}")
        if step == steps:
            break
        grads = encoder_backward(params, cache, np.concatenate([dEa, dEp, dEn], axis=1))
        opt.step(params, grads)
    return params, float(np.mean(recent))


def gradient_check(seed=0):
    """
    数值梯度检验：在一个很小的编码器（5 → 4 → 3）和 6 个三元组上，比较手写的梯度和 (J(θ+h) − J(θ−h)) / 2h。

    【返回】相对误差 ‖数值 − 解析‖ / (‖数值‖ + ‖解析‖)，1e−7 左右或更小说明反向传播写对了（第 7 讲）。
    【numpy 小知识】np.ndindex(shape) 依次给出数组的每一个下标元组 (i, j)。
    """
    rng = np.random.default_rng(seed)
    p = init_encoder(rng, d_in=5, h=4, e=3)
    Xa, Xp, Xn = (rng.normal(size=(5, 6)) for _ in range(3))
    X = np.concatenate([Xa, Xp, Xn], axis=1)

    def J_of(params):
        E, _ = encode(params, X)
        return triplet_loss_and_grads(E[:, :6], E[:, 6:12], E[:, 12:], alpha=1.0)[0]

    E, cache = encode(p, X)
    _, _, dEa, dEp, dEn = triplet_loss_and_grads(E[:, :6], E[:, 6:12], E[:, 12:], alpha=1.0)
    grads = encoder_backward(p, cache, np.concatenate([dEa, dEp, dEn], axis=1))
    num, ana, h = [], [], 1e-6
    for k in p:
        for idx in np.ndindex(p[k].shape):
            old = p[k][idx]
            p[k][idx] = old + h
            jp = J_of(p)
            p[k][idx] = old - h
            jm = J_of(p)
            p[k][idx] = old
            num.append((jp - jm) / (2 * h))
            ana.append(grads["d" + k][idx])
    num, ana = np.array(num), np.array(ana)
    return float(np.linalg.norm(num - ana) / (np.linalg.norm(num) + np.linalg.norm(ana)))


# =============================================================================
# 演示 1：像素距离为什么不行
# =============================================================================
def demo_pixels(world):
    print_header("演示 1：直接比较像素 —— 同一个人的两张照片，像素差多少？")
    G = world["G"]
    rng = np.random.default_rng(101)
    # 方差分解：200 个人 × 每人 20 张。"同一个人的平均照片"之间的差异 = 身份带来的；同一个人照片之间的差异 = 干扰带来的
    Z = rng.normal(size=(Q, 200))
    ids = np.repeat(np.arange(200), 20)
    X = (take_photos(G, Z[:, ids], rng) - world["mu"]) / world["sd"]
    person_mean = X.reshape(D, 200, 20).mean(axis=2)                 # (100, 200)：每个人 20 张照片的平均
    within = X.reshape(D, 200, 20) - person_mean[:, :, None]
    total_var = X.var(axis=1).sum()
    within_var = (within ** 2).mean(axis=(1, 2)).sum()
    share_id = 1 - within_var / total_var
    print(f"像素的总方差里，来自\"是谁\"（身份）的只占 {share_id:.1%}，来自光照、姿态、表情（干扰）的占 {1 - share_id:.1%}。\n")

    d_s, d_d = pair_distances(world["X_te"], sample_pairs(world["id_te"], np.random.default_rng(11)))
    print("在 100 个新同学上随机抽 2000 对\"同一个人\"和 2000 对\"不同的人\"，比较像素的平方距离：")
    print(pad("", 18) + pad("平均", 10, ">") + pad("中位数", 10, ">") + pad("25% 分位", 10, ">") + pad("75% 分位", 10, ">"))
    for name, d in (("同一个人", d_s), ("不同的人", d_d)):
        q25, q50, q75 = np.percentile(d, [25, 50, 75])
        print(pad(name, 18) + f"{d.mean():>10.1f}{q50:>10.1f}{q25:>10.1f}{q75:>10.1f}")
    above = float(np.mean(d_s > np.median(d_d)))
    res = evaluate(world["X_tr"], world["id_tr"], world["X_te"], world["id_te"])
    print(f"\n{above:.1%} 的\"同一个人\"配对，像素距离比\"不同的人\"配对的中位数还大。")
    print(f"用最好的阈值（在训练人群上选）判断是否同一个人：新同学上的验证准确率 {res['verify_test']:.3f}（瞎猜是 0.5）")
    print(f"one-shot 识别（100 个新同学，每人一张注册照片）：top-1 准确率 {res['identify']:.3f}（瞎猜是 0.010）")
    print(f"""
观察 Observations：
  · 像素的变化里只有 {share_id:.0%} 来自"是谁"，其余都来自光照、姿态、表情这些和身份无关的因素。
    所以"同一个人、不同光照"的两张照片，在像素空间里可能比"不同的人、相同光照"的两张离得还远。
  · 两种配对的距离分布大面积重叠：{above:.0%} 的"同一个人"配对比一半的"不同人"配对还远。
    无论阈值放在哪里，都会把很多本人拒之门外、或者把很多别人放进来 —— 验证准确率只有 {res['verify_test']:.2f}。
  · 课件里列的三个原因（光照不同、化妆 / 留胡子、学生证照片过时）说的都是这件事：
    **像素距离衡量的是"两张图长得像不像"，而我们要的是"是不是同一个人"。** 需要一个只保留身份、扔掉干扰的表示。""")
    return res


# =============================================================================
# 演示 2：用 triplet loss 训练编码器
# =============================================================================
def demo_triplet(world, res_pixels):
    print_header("演示 2：用 triplet loss 学一个 16 维的人脸编码 f(x)")
    print(f"梯度检验（5 → 4 → 3 的小编码器，6 个三元组）：相对误差 {gradient_check():.2e}（1e-7 左右或更小 = 反向传播写对了）\n")
    print(f"训练：{N_TRAIN_PEOPLE} 个人 × 每人 {PHOTOS_TRAIN} 张，随机三元组，每批 128 个，3000 步，Adam，margin α = {MARGIN}")
    params, _ = train_encoder(world, log_at=(0, 500, 1000, 2000, 3000))
    rand_params = init_encoder(np.random.default_rng(3))         # 同一个种子的"训练前"编码器：对照组
    rows = [
        ("像素距离（100 维）", res_pixels),
        ("随机初始化的编码器（训练前）", evaluate(encode(rand_params, world["X_tr"])[0], world["id_tr"],
                                              encode(rand_params, world["X_te"])[0], world["id_te"])),
        ("triplet loss 训练后的编码器", evaluate(encode(params, world["X_tr"])[0], world["id_tr"],
                                             encode(params, world["X_te"])[0], world["id_te"])),
        ("上限：真实身份向量 z（6 维）", evaluate(world["Z_tr"][:, world["id_tr"]], world["id_tr"],
                                             world["Z_te"][:, world["id_te"]], world["id_te"])),
    ]
    print("\n" + pad("表示", 30) + pad("训练人群验证", 14, ">") + pad("新同学验证", 12, ">")
          + pad("新同学识别 top-1", 18, ">") + pad("不同人/同一人距离", 20, ">"))
    for name, r in rows:
        ratio = f"{r['ratio']:.2f}" if np.isfinite(r["ratio"]) else "∞"
        print(pad(name, 30) + f"{r['verify_train']:>14.3f}{r['verify_test']:>12.3f}{r['identify']:>18.3f}" + pad(ratio, 20, ">"))
    enc = rows[2][1]
    print(f"""
观察 Observations：
  · 训练后的编码器在**训练时从没见过的 100 个新同学**上：验证准确率 {enc['verify_test']:.3f}（像素 {res_pixels['verify_test']:.3f}），
    one-shot 识别 {enc['identify']:.3f}（像素 {res_pixels['identify']:.3f}）。每个新同学只有一张注册照片，也**没有重新训练任何东西**——
    新同学入学时，只要把学生证照片过一遍编码器、存下 16 个数就行。
  · 对照组：随机初始化的编码器（{rows[1][1]['verify_test']:.3f} / {rows[1][1]['identify']:.3f}）比像素还差一点 —— 随机的非线性变换不会自己"扔掉光照"，
    是 triplet loss 教会了它哪些变化该忽略。
  · "不同人 / 同一人"的平均距离之比：像素只有 {res_pixels['ratio']:.2f}，编码 {enc['ratio']:.2f}。编码把同一个人的照片聚到一起、把不同的人推开。
  · 训练人群上的验证准确率（{enc['verify_train']:.3f}）比新同学（{enc['verify_test']:.3f}）高一些：编码器对见过的人更熟。
    差距不大，说明它学到的是"怎样从照片里读出身份"这种**通用的本领**，而不是记住了 300 个人。
  · 离上限（真实身份向量，{rows[3][1]['identify']:.3f}）还有距离：要"完全扔掉干扰"需要更多的人、更大的网络
    （FaceNet 的训练集有约 800 万个不同的人）。""")
    return params, rows[2][1]


# =============================================================================
# 演示 3：损失函数的每一项都有用吗？
# =============================================================================
def demo_ablation(world, main_params):
    print_header("演示 3：改一改损失函数和三元组的选法 —— 每一部分有什么用？")
    print(pad("损失函数 / 三元组", 56) + pad("新同学验证", 12, ">") + pad("识别 top-1", 12, ">")
          + pad("编码的分散程度", 16, ">") + pad("最后仍有梯度", 14, ">"))
    variants = [("完整的 triplet loss（α = 0.2，随机三元组，演示 2）", None),
                ("同上，但用 semi-hard 负样本", dict(mining="semihard")),
                ("去掉 margin（α = 0）", dict(alpha=0.0)),
                ("只拉近同一个人：L = ‖f(A) − f(P)‖²", dict(mode="pull"))]
    out = []
    for name, kw in variants:
        params = main_params if kw is None else train_encoder(world, **kw)[0]
        frac = last_active_fraction(world, params, **(kw or {}))
        E_te = encode(params, world["X_te"])[0]
        r = evaluate(encode(params, world["X_tr"])[0], world["id_tr"], E_te, world["id_te"])
        spread = float(E_te.std(axis=1).mean())        # 每一维在所有照片上的标准差，再对 16 维取平均
        out.append((r, spread, frac))
        print(pad(name, 56) + f"{r['verify_test']:>12.3f}{r['identify']:>12.3f}{spread:>16.3f}{frac:>14.1%}")
    (r0, s0, f0), (r1, _, f1), (r2, s2, _), (r3, s3, _) = out
    print(f"""
观察 Observations：
  · 随机三元组很快就"太容易"了：训练结束时只有 {f0:.1%} 的随机三元组还有梯度（演示 2 的训练记录里，500 步时就已经几乎为 0）——
    随便抽的"另一个人"通常本来就离得很远。semi-hard 挖掘专门挑"比正样本远、但还没远出 margin"的负样本，
    训练结束时仍有 {f1:.1%} 的三元组在提供梯度。效果：识别 {r0['identify']:.3f} → {r1['identify']:.3f}，验证 {r0['verify_test']:.3f} → {r1['verify_test']:.3f}。
    在这个只有 300 个人的玩具世界里差别不大（验证甚至略低一点，在随机波动范围内）；
    FaceNet 的论文则强调，在几百万人的数据集上，怎样挑三元组对训练能不能快速收敛至关重要。
  · 去掉 margin（α = 0）：只要求"负样本比正样本远"，远一点点就满足了、梯度就没了。
    验证 {r2['verify_test']:.3f}、识别 {r2['identify']:.3f}（完整版 {r0['verify_test']:.3f} / {r0['identify']:.3f}），编码的分散程度也小得多（{s2:.3f}）。
    识别要在 100 个人里挑最近的那个，需要人和人之间**留出余量** —— 这就是 margin α 的作用。
  · 只拉近同一个人、不推开不同的人：网络找到了一个"作弊"的解 —— **把所有照片都映射到几乎同一个点**，
    损失几乎为 0，但编码的分散程度从 {s0:.3f} 塌缩到 {s3:.3f}，验证 {r3['verify_test']:.3f}、识别 {r3['identify']:.3f}，和瞎猜差不多。
    这叫**表示坍缩（representation collapse）**：损失函数没有禁止的捷径，梯度下降一定会找到。
  · 教训（也是这一讲的主题）：**损失函数定义了任务。** 你写下的每一项都在告诉网络"什么是好"；
    少了一项，网络就会钻那个空子。对比学习（contrastive learning）里大量的工作，都是在和"坍缩"作斗争（讲义 §3.6）。""")


def last_active_fraction(world, params, alpha=MARGIN, mode="triplet", mining="random"):
    """
    训练结束后，用同样的方式抽 20 批三元组，看有多少比例的三元组损失 > 0（还能提供梯度）。

    【返回】float，0–1。"只拉近"（mode = "pull"）没有 hinge，所有三元组都有梯度，返回 1。
    """
    rng = np.random.default_rng(999)
    X, ids = world["X_tr"], world["id_tr"]
    fracs = []
    for _ in range(20):
        a, p, n = sample_triplets(ids, PHOTOS_TRAIN, rng, 128)
        if mining == "semihard":
            n = pick_semihard(params, X, a, p, random_negatives(a, ids, PHOTOS_TRAIN, rng, (128, 10)))
        E = encode(params, X[:, np.concatenate([a, p, n])])[0]
        fracs.append(triplet_loss_and_grads(E[:, :128], E[:, 128:256], E[:, 256:], alpha, mode)[1])
    return float(np.mean(fracs))


# =============================================================================
# 演示 4：人脸聚类
# =============================================================================
def kmeans(E, k, rng, restarts=20, iters=50):
    """
    k-means 聚类：把 m 个点分成 k 簇，让每个点到自己簇中心的平方距离之和最小。

    【思路】重复：(1) 每个点归到最近的中心；(2) 每个中心挪到它那一簇的平均位置。
            结果依赖初始中心，所以随机初始化 restarts 次，留下总距离最小的那次。
    【参数】E：(d, m)；k：簇数；rng；restarts、iters：重启次数、每次的迭代次数
    【返回】labels：(m,)，每个点属于第几簇
    【边界】某一簇变空时，把它的中心随机放到一个点上。
    """
    m = E.shape[1]
    best, best_cost = None, np.inf
    for _ in range(restarts):
        centers = E[:, rng.choice(m, k, replace=False)].copy()
        for _ in range(iters):
            labels = np.argmin(pair_sq_dists(E, centers), axis=1)
            for j in range(k):
                members = E[:, labels == j]
                centers[:, j] = members.mean(axis=1) if members.shape[1] else E[:, rng.integers(m)]
        cost = pair_sq_dists(E, centers)[np.arange(m), labels].sum()
        if cost < best_cost:
            best, best_cost = labels, cost
    return best


def purity(labels, ids):
    """纯度：每一簇按它里面最多的那个人算"对"，对的照片占全部照片的比例。1 = 每簇都只有一个人。"""
    return float(sum(np.bincount(ids[labels == j]).max() for j in np.unique(labels)) / len(ids))


def demo_cluster(world, params):
    print_header("演示 4：没有标签的时候 —— 把相册里同一个人的照片聚到一起（k-means）")
    n_people = 10
    keep = world["id_te"] < n_people                         # 新同学里的前 10 个人，每人 6 张，共 60 张
    ids = world["id_te"][keep]
    reps = [("像素", world["X_te"][:, keep]),
            ("triplet 编码", encode(params, world["X_te"][:, keep])[0]),
            ("上限：真实身份向量 z", world["Z_te"][:, ids])]
    print(f"{n_people} 个新同学 × 每人 {PHOTOS_TEST} 张 = {len(ids)} 张照片，k-means 分成 {n_people} 簇（20 次随机重启）\n")
    print(pad("表示", 24) + pad("纯度", 8, ">") + "   每一簇里各有几张照片")
    scores = {}
    for name, E in reps:
        labels = kmeans(E, n_people, np.random.default_rng(401))
        scores[name] = purity(labels, ids)
        sizes = sorted(np.bincount(labels, minlength=n_people).tolist(), reverse=True)
        print(pad(name, 24) + f"{scores[name]:>8.3f}   {sizes}")
    print(f"""
观察 Observations：
  · 用像素聚类，纯度只有 {scores['像素']:.3f}：簇是按"光照、姿态相似"聚起来的，而不是按人。
  · 用 triplet 编码聚类，纯度 {scores['triplet 编码']:.3f}。同一个编码器、没有任何新标签、没有重新训练，
    就能把相册里的照片按人分组 —— 课件里的"人脸聚类"就是这样做的（真实系统还会先做人脸检测，把脸裁出来）。
  · 这说明一个好的编码是**通用的**：验证（比较两个编码的距离）、识别（在库里找最近的编码，K 近邻）、
    聚类（对编码做 k-means）用的是同一个 f(x)，只是后面接的"小算法"不同。""")


def main():
    world = make_world()
    res_pixels = demo_pixels(world)
    params, _ = demo_triplet(world, res_pixels)
    demo_ablation(world, params)
    demo_cluster(world, params)


if __name__ == "__main__":
    main()