"""第 40 章：从像素和两个词开始，训练一个小型图文检索模型。"""

from argparse import ArgumentParser
from copy import deepcopy
from hashlib import sha256
from pathlib import Path
import json
import platform

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F


COLORS = ["red", "green", "blue"]
SHAPES = ["circle", "square", "triangle"]
CAPTIONS = [f"{color} {shape}" for color in COLORS for shape in SHAPES]
RGB = np.array([[.85, .18, .16], [.15, .68, .30], [.15, .35, .88]])
HELD_OUT = [2, 7]  # red triangle、blue square 不进入训练或验证。
SEEN = [i for i in range(9) if i not in HELD_OUT]
SIZE, PATCH, WIDTH, DIM = 32, 8, 32, 16
SEED, EPOCHS, TEMPERATURE = 2026, 40, .10
REPORT = Path(__file__).with_name("data") / "cross-modal-alignment.json"


def make_data(split):
    counts = {"train": 48, "validation": 16, "test": 24}
    seeds = {"train": 101, "validation": 202, "test": 303}
    rng = np.random.default_rng(seeds[split])
    labels = range(9) if split == "test" else SEEN
    yy, xx = np.mgrid[:SIZE, :SIZE]
    images, targets, ids = [], [], []
    for label in labels:
        for index in range(counts[split]):
            cx, cy = rng.uniform(12, 20, 2)
            radius, angle = rng.uniform(6, 10), rng.uniform(-.3, .3)
            dx, dy = xx - cx, yy - cy
            u = np.cos(angle) * dx + np.sin(angle) * dy
            v = -np.sin(angle) * dx + np.cos(angle) * dy
            if label % 3 == 0:
                mask = u*u + v*v <= radius*radius
            elif label % 3 == 1:
                mask = np.maximum(abs(u), abs(v)) <= radius * .88
            else:
                mask = (v >= -radius) & (v <= radius) & (abs(u) <= (v+radius)/2)
            background = rng.uniform(.86, .99)
            color = np.clip(RGB[label//3] + rng.normal(0, .025, 3), 0, 1)
            image = np.full((SIZE, SIZE, 3), background)
            image[mask] = color
            image = np.clip(image + rng.normal(0, .01, image.shape), 0, 1)
            images.append(image.transpose(2, 0, 1).astype("float32"))
            targets.append(label)
            ids.append(f"{split}-{label}-{index:02d}")
    return torch.from_numpy(np.stack(images)), torch.tensor(targets), ids


def text_ids(labels):
    # 六个词；词表只表达颜色与形状，不支持自由句子。
    return torch.stack([labels // 3, 3 + labels % 3], dim=-1)


class TinyImageEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.patch = nn.Conv2d(3, WIDTH, PATCH, stride=PATCH)
        self.cls = nn.Parameter(torch.randn(1, 1, WIDTH) * .02)
        self.position = nn.Parameter(torch.randn(1, 17, WIDTH) * .02)
        self.blocks = nn.ModuleList([
            nn.TransformerEncoderLayer(WIDTH, 4, WIDTH*2, dropout=0,
                                       activation="gelu", batch_first=True,
                                       norm_first=True) for _ in range(2)
        ])
        self.norm = nn.LayerNorm(WIDTH)
        self.projection = nn.Linear(WIDTH, DIM, bias=False)

    def forward(self, images):
        x = self.patch(images * 2 - 1).flatten(2).transpose(1, 2)
        x = torch.cat([self.cls.expand(len(x), -1, -1), x], dim=1)
        x = x + self.position
        for block in self.blocks:
            x = block(x)  # 图像可以同时看全部块，没有语言生成的因果掩码。
        return self.projection(self.norm(x[:, 0]))


class DualEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.image = TinyImageEncoder()
        self.words = nn.Embedding(6, WIDTH)
        self.text_projection = nn.Linear(WIDTH, DIM, bias=False)

    def encode_image(self, images):
        return F.normalize(self.image(images), dim=-1)

    def encode_text(self, ids):
        return F.normalize(self.text_projection(self.words(ids).mean(1)), dim=-1)

    def forward(self, images, ids):
        return self.encode_image(images) @ self.encode_text(ids).T / TEMPERATURE


def contrastive_loss(logits):
    if logits.ndim != 2 or logits.shape[0] != logits.shape[1] or len(logits) < 2:
        raise ValueError("需要至少两个一一配对的样本，得分矩阵必须为方阵")
    target = torch.arange(len(logits), device=logits.device)
    return (F.cross_entropy(logits, target) +
            F.cross_entropy(logits.T, target)) / 2


@torch.inference_mode()
def scores(model, images, labels):
    model.eval()
    vectors = torch.cat([model.encode_image(x) for x in images.split(64)])
    return vectors @ model.encode_text(text_ids(labels)).T


def retrieval_metrics(similarity, image_labels, caption_labels):
    # 多张图可与同一句描述相关；按语义标签判断，不能只认矩阵对角线。
    i_order = similarity.argsort(dim=1, descending=True, stable=True)
    t_order = similarity.T.argsort(dim=1, descending=True, stable=True)
    result = {}
    for name, order, gallery, query in [
        ("i2t", i_order, caption_labels, image_labels),
        ("t2i", t_order, image_labels, caption_labels),
    ]:
        relevant = gallery[order] == query[:, None]
        for k in [1, 3]:
            result[f"{name}_hit{k}"] = relevant[:, :k].any(1).float().mean().item()
    return result


def evaluate(model, data):
    images, labels, ids = data
    captions = torch.arange(9)
    similarity = scores(model, images, captions)
    prediction = similarity.argmax(1)
    known = torch.tensor([int(x) in SEEN for x in labels])
    result = retrieval_metrics(similarity, labels, captions)
    result.update({
        "seen_i2t_hit1": (prediction[known] == labels[known]).float().mean().item(),
        "heldout_i2t_hit1": (prediction[~known] == labels[~known]).float().mean().item(),
        "per_class_correct": [int(((prediction == label) & (labels == label)).sum())
                              for label in range(9)],
        "ids": ids, "labels": labels.tolist(), "predictions": prediction.tolist(),
        "similarities": similarity.tolist(),
    })
    return result


def tensor_digest(tensor):
    return sha256(tensor.contiguous().numpy().tobytes()).hexdigest()


def search_images(model, query, images, k=3):
    if query not in CAPTIONS:
        raise ValueError("本例只接受九种颜色与形状组合，例如 blue square")
    similarity = scores(model, images, torch.tensor([CAPTIONS.index(query)]))[:, 0]
    order = similarity.argsort(descending=True, stable=True)[:k]
    return [(int(index), float(similarity[index])) for index in order]


def run():
    torch.set_num_threads(1)
    torch.manual_seed(SEED)
    torch.use_deterministic_algorithms(True)
    # 关闭推理专用 fast path，使训练与评估使用同一注意力实现路径。
    torch.backends.mha.set_fastpath_enabled(False)
    train, validation, test = [make_data(s) for s in ["train", "validation", "test"]]
    model = DualEncoder()
    initial = evaluate(model, test)
    optimizer = torch.optim.AdamW(model.parameters(), lr=.001, weight_decay=.01)
    generator = torch.Generator().manual_seed(SEED + 1)
    image_bank = train[0].reshape(7, 48, 3, SIZE, SIZE)
    labels = torch.tensor(SEEN)
    best, best_loss, best_epoch = None, float("inf"), 0
    history = []
    for epoch in range(1, EPOCHS + 1):
        model.train()
        orders = torch.stack([torch.randperm(48, generator=generator) for _ in SEEN])
        total = 0.
        for step in range(48):
            # 每批七对，每种语义只出现一次，避免同义描述被当成负样本。
            images = image_bank[torch.arange(7), orders[:, step]]
            loss = contrastive_loss(model(images, text_ids(labels)))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total += loss.item()
        if epoch == 1 or epoch % 5 == 0:
            val_logits = scores(model, validation[0], labels) / TEMPERATURE
            val_targets = torch.searchsorted(labels, validation[1])
            val_loss = F.cross_entropy(val_logits, val_targets).item()
            history.append({"epoch": epoch, "train_loss": total/48,
                            "validation_i2t_ce": val_loss})
            if val_loss < best_loss:
                best, best_loss, best_epoch = deepcopy(model.state_dict()), val_loss, epoch
    model.load_state_dict(best)
    final = evaluate(model, test)
    state_hash = sha256()
    for name, tensor in sorted(best.items()):
        state_hash.update(name.encode())
        state_hash.update(tensor.contiguous().numpy().tobytes())
    report = {
        "config": {"seed": SEED, "epochs": EPOCHS, "temperature": TEMPERATURE,
                   "size": SIZE, "patch": PATCH, "width": WIDTH, "dimension": DIM,
                   "batch_size": 7, "learning_rate": .001, "weight_decay": .01,
                   "held_out": HELD_OUT, "selection": "minimum seen validation i2t CE"},
        "environment": {"python": platform.python_version(), "torch": torch.__version__,
                        "numpy": np.__version__, "device": "cpu", "threads": 1},
        "source_sha256": sha256(Path(__file__).read_bytes()).hexdigest(),
        "data_sha256": {name: tensor_digest(data[0]) for name, data in
                        zip(["train", "validation", "test"], [train, validation, test])},
        "model_sha256": state_hash.hexdigest(),
        "parameter_count": sum(p.numel() for p in model.parameters()),
        "selected_epoch": best_epoch, "history": history,
        "initial": initial, "trained": final,
    }
    return model, report


def print_summary(report):
    print("合成图文：训练 336 张，验证 112 张，测试 216 张；候选描述 9 条。")
    print("留出组合：red triangle、blue square；训练与验证均不使用。")
    print(f"参数 {report['parameter_count']}；按验证损失选择第 {report['selected_epoch']} 轮。")
    for key, name in [("initial", "随机初始化"), ("trained", "对比训练后")]:
        row = report[key]
        print(f"{name}：图找文 Hit@1={row['i2t_hit1']:.3f}；"
              f"文找图 Hit@1={row['t2i_hit1']:.3f}；Hit@3={row['t2i_hit3']:.3f}")
        print(f"  已见组合={row['seen_i2t_hit1']:.3f}；"
              f"未见组合={row['heldout_i2t_hit1']:.3f}")
    print("每类图找文正确数（各 24 张）：", report["trained"]["per_class_correct"])
    print("这些结果只描述合成图库；模型不支持自由句子，也不会生成回答。")


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path, help="另存本次训练的完整报告")
    parser.add_argument("--checkpoint", type=Path, help="另存模型权重")
    parser.add_argument("--query", choices=CAPTIONS, help="训练后用两个词检索测试图库")
    args = parser.parse_args()
    model, report = run()
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n")
    if args.checkpoint:
        args.checkpoint.parent.mkdir(parents=True, exist_ok=True)
        torch.save(model.state_dict(), args.checkpoint)
        restored = DualEncoder()
        restored.load_state_dict(torch.load(args.checkpoint, weights_only=True))
        torch.testing.assert_close(scores(model, make_data("test")[0], torch.arange(9)),
                                   scores(restored, make_data("test")[0], torch.arange(9)),
                                   rtol=0, atol=0)
    print_summary(report)
    if args.query:
        images, labels, ids = make_data("test")
        print(f"查询：{args.query}")
        for index, similarity in search_images(model, args.query, images):
            print(f"  {ids[index]}；实际 {CAPTIONS[int(labels[index])]}；余弦 {similarity:.4f}")


if __name__ == "__main__":
    main()

1. $14\times14=196$ 块，每块有 $16\times16\times3=768$ 个数；加入 CLS 后为 197 个位置。
2. 块数为四倍，图块间分数元素数为 16 倍。投影、其他算子、硬件利用率和内存访问也会影响实际时间；加入 CLS 后的精确矩阵边长还需另算。
3. $2-2\times0.6=0.8$。
4. 概率为 $e^2/(e^2+1)\approx0.8808$，损失为 $-\log(0.8808)\approx0.1269$。这里的“分数”已经包含温度缩放。
5. 相同描述给出相同文本向量，两列相同；目标却要求把其中一列当正例、另一列当负例。可以避免批内重复语义，也可以改成承认多个正例的训练目标。
6. Hit@5 为 1，Precision@5 为 $3/5$，Recall@5 为 $3/20$。
7. 分别沿该文字对应的列、该图片对应的行排序。两个方向在不同候选集合里竞争，不要求互相选中；以实际记录为准。
8. 保持数据种子不变，只改变模型初始化与批次排列的训练种子。分别列出结果，再报告均值与范围。这检查的是训练随机性，不等于更换数据分布的稳健性测试。
9. 可以画出“水平位置—正确描述相似度”曲线，同时记录最高分候选。相似度降低与预测变化是两种不同现象，都值得观察。
10. 图片中需要多个对象与可核对的位置关系；文本需要保留词序，训练候选要包含左右关系相反但词汇相同的描述。数据划分还要防止重复布局泄漏，评价不能只检查颜色和形状是否出现。