"""
vlm_lab.py —— 视觉-语言模型的零件：图像编码器、文字编码器、CLIP / SigLIP 的损失、零样本分类与检索（第 13 讲）
=====================================================================

【这个文件要回答的问题】
    slides：VLM 的总体思路是把图像和文字都变成嵌入向量，再交给一个只有解码器的 Transformer 做"下一个（文字）token 预测"。
    图像编码器有两种常见选择：VQ-VAE（离散 token，第 9、12 讲）和 CLIP（用对比学习学一个连续的嵌入）。
    这个文件实现 CLIP 这一支：
        ImageEncoder   小 ViT（第 5 讲的结构）：图像 → 49 个块的向量 + 一个汇总向量
        TextEncoder    小 Transformer（非因果，像 BERT）：描述 → 一个汇总向量
        CLIP           两个编码器各接一个线性投影到同一个 d 维空间（多模态嵌入空间），加一个可学习的温度 τ
        clip_loss      slides 的"正确目标"：双向的交叉熵（InfoNCE）
        naive_loss     slides 的"错误（但直观）的目标"：Σ_i I_i·T_i − Σ_i Σ_{j≠i} I_i·T_j
        siglip_loss    SigLIP：把 softmax 换成逐对的 sigmoid
        train_clip / load_or_train_clip / embed_images / embed_texts / retrieval_accuracy / zero_shot

【数据】第 12 讲的带描述的形状世界（caption_world.py）：4 种说法随机用，让文字编码器见到不同的写法。
"""

import os
import random
import sys
import time

import torch
import torch.nn as nn
import torch.nn.functional as F

_HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture12", "lecture02", "lecture05"):
    _p = os.path.join(_HERE, "..", _d)
    if _p not in sys.path:
        sys.path.insert(0, _p)
from caption_world import CAPTION_VOCAB, CID, HOLDOUT, TEMPLATES, caption, make_captioned  # noqa: E402
from lm_models import Block  # noqa: E402
from vision_models import patchify  # noqa: E402

CKPT_DIR = os.path.join(_HERE, "checkpoints")
MAX_TXT = 12


class ImageEncoder(nn.Module):
    """
    小 ViT：28×28 的图切成 7×7 = 49 个 4×4 的块 → 线性嵌入 → 加上 [CLS] 和位置嵌入 → L 个非因果 Transformer 块。
    【返回】forward(x) -> (汇总向量 (N, d)，49 个块的向量 (N, 49, d))
        汇总向量用来做 CLIP 的对比学习；块的向量留给 03 号脚本的 VLM（"把整张图的细节交给语言模型"）。
    """

    def __init__(self, d=64, n_layers=2, n_heads=4, P=4):
        super().__init__()
        self.P = P
        self.embed = nn.Linear(P * P, d)
        self.cls = nn.Parameter(torch.zeros(1, 1, d))
        self.pos = nn.Parameter(torch.randn(1, (28 // P) ** 2 + 1, d) * 0.02)
        self.blocks = nn.ModuleList(Block(d, n_heads, 4 * d, norm="pre", causal=False) for _ in range(n_layers))
        self.norm = nn.LayerNorm(d)

    def forward(self, x):
        h = self.embed(patchify(x, self.P))
        h = torch.cat([self.cls.expand(len(x), -1, -1), h], dim=1) + self.pos
        for blk in self.blocks:
            h = blk(h)
        h = self.norm(h)
        return h[:, 0], h[:, 1:]


class TextEncoder(nn.Module):
    """
    小 Transformer 文字编码器（非因果）：词嵌入 + 位置嵌入 → L 个块 → 对真实的词（非 <pad>）取平均。
    【输入】ids (N, T)，右边补 <pad>（编号 0）
    """

    def __init__(self, d=64, n_layers=2, n_heads=4):
        super().__init__()
        self.tok = nn.Embedding(len(CAPTION_VOCAB), d)
        self.pos = nn.Embedding(MAX_TXT, d)
        self.blocks = nn.ModuleList(Block(d, n_heads, 4 * d, norm="pre", causal=False) for _ in range(n_layers))
        self.norm = nn.LayerNorm(d)

    def forward(self, ids):
        mask = (ids != 0).float()
        h = self.tok(ids) + self.pos(torch.arange(ids.size(1)))
        for blk in self.blocks:
            blk.attn.extra_mask = ((1 - mask) * -1e9)[:, None, None, :]      # 不看 <pad>
            h = blk(h)
            blk.attn.extra_mask = None
        h = self.norm(h)
        return (h * mask[..., None]).sum(1) / mask.sum(1, keepdim=True)


class CLIP(nn.Module):
    """
    CLIP（Radford et al., 2021）的缩小版：图像编码器和文字编码器各接一个线性投影，投到同一个 dim 维空间，再归一化成单位向量。
    相似度 = 余弦相似度 / τ，τ 是可学习的温度（存的是 log(1/τ)，初始 τ = 0.07，和原论文一样）。
    """

    def __init__(self, d=64, dim=32):
        super().__init__()
        self.img = ImageEncoder(d)
        self.txt = TextEncoder(d)
        self.img_proj = nn.Linear(d, dim, bias=False)
        self.txt_proj = nn.Linear(d, dim, bias=False)
        self.logit_scale = nn.Parameter(torch.tensor(float(torch.log(torch.tensor(1 / 0.07)))))
        self.bias = nn.Parameter(torch.tensor(-10.0))          # 只有 SigLIP 用（论文的初始值）

    def encode_image(self, x):
        return F.normalize(self.img_proj(self.img(x)[0]), dim=-1)

    def encode_text(self, ids):
        return F.normalize(self.txt_proj(self.txt(ids)), dim=-1)


def text_ids(captions):
    """描述（词列表的列表）→ (N, MAX_TXT) 的编号，右边补 <pad>。"""
    ids = torch.zeros(len(captions), MAX_TXT, dtype=torch.long)
    for i, c in enumerate(captions):
        ids[i, :len(c)] = torch.tensor([CID[w] for w in c])
    return ids


# =====================================================================
# 损失
# =====================================================================

def clip_loss(img_emb, txt_emb, logit_scale):
    """
    slides 的"正确目标"：N 张图和 N 句描述两两算相似度，得到 N×N 的矩阵 S（S_ij = I_i·T_j / τ）。
        图 → 文：第 i 行做 softmax，正确答案是第 i 列（p_IT(i 配 j | I_i, T_1..T_N)，"按列做 softmax"）
        文 → 图：第 j 列做 softmax，正确答案是第 j 行
    两个方向的交叉熵取平均（也就是最大化两个条件分布下正确配对的对数似然）。
    """
    S = logit_scale.exp() * img_emb @ txt_emb.t()
    labels = torch.arange(len(S))
    return 0.5 * (F.cross_entropy(S, labels) + F.cross_entropy(S.t(), labels))


def naive_loss(img_emb, txt_emb):
    """
    slides 的"错误（但直观）的目标"：最大化 Σ_i I_i·T_i − Σ_i Σ_{j≠i} I_i·T_j（这里取负号作为损失、除以 N）。
    【为什么"错"】它对相似度是线性的：把一对不匹配的推远，不管它们本来已经离得多远，奖励都一样——没有 softmax 那种
        "只盯着最容易混淆的那几个"的效果；而且 Σ_{j≠i} T_j 几乎就是"所有描述的平均"，第二项只是在把每张图推离这个平均方向。
    """
    S = img_emb @ txt_emb.t()
    N = len(S)
    pos = S.diag().sum()
    neg = S.sum() - pos
    return -(pos - neg) / N


def siglip_loss(img_emb, txt_emb, logit_scale, bias):
    """
    SigLIP（Zhai et al., 2023）：每一对 (i, j) 单独做一次二分类——配对的 y = +1，不配对的 y = −1：
        loss = −(1/N) Σ_i Σ_j log σ( y_ij · (t·I_i·T_j + b) )
    【和 CLIP 的区别】没有对整行 / 整列的 softmax 归一化：每一对的损失只依赖这一对，所以 N×N 的矩阵可以分块算、
        分散在很多机器上，不需要先把所有嵌入收集到一起求归一化常数（02 号脚本验证这一点）。
    """
    S = logit_scale.exp() * img_emb @ txt_emb.t() + bias
    y = 2 * torch.eye(len(S), dtype=S.dtype) - 1
    return -F.logsigmoid(y * S).sum() / len(S)


# =====================================================================
# 训练、评估
# =====================================================================

def train_clip(model, x, captions_fn, steps, loss="clip", batch=128, lr=1e-3, seed=0, log=None):
    """
    训练 CLIP。【参数】captions_fn(idx, rng) -> 这些图的描述（每次随机挑一种说法）；loss："clip" / "naive" / "siglip"
    【返回】最后 100 步的平均损失
    """
    torch.manual_seed(seed)
    gen = torch.Generator().manual_seed(seed)
    rng = random.Random(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    hist = []
    model.train()
    for step in range(1, steps + 1):
        for gr in opt.param_groups:
            gr["lr"] = lr * min(1.0, step / 100)
        idx = torch.randint(len(x), (batch,), generator=gen)
        I = model.encode_image(x[idx])
        T = model.encode_text(text_ids(captions_fn(idx.tolist(), rng)))
        if loss == "clip":
            ls = clip_loss(I, T, model.logit_scale)
        elif loss == "siglip":
            ls = siglip_loss(I, T, model.logit_scale, model.bias)
        else:
            ls = naive_loss(I, T)
        opt.zero_grad()
        ls.backward()
        opt.step()
        with torch.no_grad():
            model.logit_scale.clamp_(0, 4.6052)                 # τ ≥ 0.01，和原论文一样
        hist.append(ls.item())
        if log and step % max(steps // 4, 1) == 0:
            log(f"    step {step:5d}  最近 100 步的平均损失 {sum(hist[-100:]) / len(hist[-100:]):.4f}")
    model.eval()
    return sum(hist[-100:]) / len(hist[-100:])


def caption_sampler(attrs):
    """返回一个 captions_fn：给图的下标，随机挑一种说法生成描述。"""
    def fn(idx, rng):
        return [caption(attrs[i], rng.randrange(len(TEMPLATES))) for i in idx]
    return fn


ALL_ATTRS = [(s, z, v, h) for s in range(4) for z in range(2) for v in range(2) for h in range(2)]


@torch.no_grad()
def embed_images(model, x, batch=1000):
    return torch.cat([model.encode_image(x[i:i + batch]) for i in range(0, len(x), batch)])


@torch.no_grad()
def embed_texts(model, captions, batch=1000):
    return torch.cat([model.encode_text(text_ids(captions[i:i + batch])) for i in range(0, len(captions), batch)])


@torch.no_grad()
def retrieval_accuracy(model, x, attrs, template=0):
    """
    图 → 文检索：每张图在 32 种描述（某一种说法）里挑最相似的，看 4 个属性是否全对；也报告每个属性单独的准确率。
    【返回】(4 个全对的比例, [形状, 大小, 上下, 左右] 的准确率)
    """
    T = embed_texts(model, [caption(torch.tensor(a), template) for a in ALL_ATTRS])
    I = embed_images(model, x)
    pred = torch.tensor([ALL_ATTRS[j] for j in (I @ T.t()).argmax(1).tolist()])
    eq = (pred == attrs).float()
    return eq.all(1).float().mean().item(), eq.mean(0).tolist()


def load_or_train_clip(name, steps, loss="clip", holdout=True, log=print, batch=128):
    """
    训练（或读取缓存的）CLIP：8192 张带描述的图（seed = 1），holdout=True 时去掉第 12 讲那 4 种属性组合。
    【返回】(model, info)
    """
    cfg = {"steps": steps, "loss": loss, "holdout": holdout, "batch": batch, "d": 64, "dim": 32}
    path = os.path.join(CKPT_DIR, f"{name}.pt")
    # 【可复现】PyTorch 每次启动时的默认随机种子是随机的；构造模型（随机初始化）之前一定要先固定种子
    torch.manual_seed(0)
    model = CLIP()
    if steps >= 500 and os.path.exists(path):
        blob = torch.load(path, weights_only=True)
        if blob["config"] == cfg:
            model.load_state_dict(blob["state_dict"])
            model.eval()
            if log:
                log(f"  （读取缓存：checkpoints/{name}.pt）")
            return model, blob["info"]
    x, a, _ = make_captioned(8192, 1)
    if holdout:
        keep = torch.tensor([tuple(r.tolist()) not in HOLDOUT for r in a])
        x, a = x[keep], a[keep]
    t0 = time.perf_counter()
    final = train_clip(model, x, caption_sampler(a), steps, loss=loss, batch=batch, log=log)
    info = {"loss": final, "seconds": time.perf_counter() - t0, "n_train": len(x)}
    if steps >= 500:
        os.makedirs(CKPT_DIR, exist_ok=True)
        torch.save({"config": cfg, "state_dict": model.state_dict(), "info": info}, path)
        if log:
            log(f"  （已保存到 checkpoints/{name}.pt）")
    return model, info