"""
caption_world.py —— 带文字描述的形状世界："a big circle at the top left"（第 12、13 讲）
=====================================================================

【这个文件要回答的问题】
    文本到图像、视觉-语言模型都需要"图像 + 文字"成对的数据。我们给第 5 讲的形状世界配上描述：
        每张图有 4 个属性：形状（circle / square / triangle / cross）、大小（small / big：半径是否超过 6.25 像素）、
        上下（top / bottom：中心的 y 是否小于 14）、左右（left / right：中心的 x 是否小于 14）——一共 4×2×2×2 = 32 种组合。
        描述就是把属性填进模板："a big circle at the top left"。
    生成的图像"听没听话"怎么判断？用第 6 讲的形状拟合（fit_shapes）从图里读出形状、中心、半径，再换算成 4 个属性，
    和描述比较——就像真实研究里用一个检测器 / 分类器去检查生成的图（比如 GenEval）。

【本文件提供】
    SHAPE_WORDS / SIZE_WORDS / VERT_WORDS / HORIZ_WORDS / CAPTION_VOCAB / CID     词表
    attributes(labels, params)            从形状世界的参数算出属性 (N, 4)
    caption(attr, template)               属性 → 描述（词列表）；TEMPLATES 里有几种说法（第 13 讲用）
    make_captioned(n, seed)               生成 (图像, 属性, 描述)
    read_attributes(images)               用形状拟合从图像里读出属性
    attribute_accuracy(pred, target)      逐属性的准确率 + 4 个属性全对的比例
"""

import os
import sys

import torch

_HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture06", "lecture05"):
    _p = os.path.join(_HERE, "..", _d)
    if _p not in sys.path:
        sys.path.insert(0, _p)
from shape_metrics import fit_shapes  # noqa: E402
from shapes_world import make_shapes  # noqa: E402

SHAPE_WORDS = ("circle", "square", "triangle", "cross")
SIZE_WORDS = ("small", "big")
VERT_WORDS = ("top", "bottom")
HORIZ_WORDS = ("left", "right")
ATTR_NAMES = ("形状", "大小", "上下", "左右")
# 文生图实验里训练时故意去掉的 4 种属性组合（形状, 大小, 上下, 左右），用来测"没见过的组合"
HOLDOUT = [(0, 1, 0, 0), (1, 0, 1, 1), (2, 1, 1, 0), (3, 0, 0, 1)]
BIG_R, MID = 6.25, 14.0

# 描述的几种说法（第 12 讲只用第 0 种；第 13 讲的 CLIP 用全部，让文字编码器见到不同的写法）
TEMPLATES = [
    lambda s, z, v, h: ["a", z, s, "at", "the", v, h],
    lambda s, z, v, h: ["a", s, ",", z, ",", v, h],
    lambda s, z, v, h: ["the", s, "is", z, "and", "at", "the", v, h],
    lambda s, z, v, h: [v, h, ":", "a", z, s],
]
CAPTION_VOCAB = ["<pad>", "<s>", "<img>", "a", "at", "the", ",", "is", "and", ":"] + \
    list(SHAPE_WORDS) + list(SIZE_WORDS) + list(VERT_WORDS) + list(HORIZ_WORDS)
CID = {w: i for i, w in enumerate(CAPTION_VOCAB)}


def attributes(labels, params):
    """
    【参数】labels (N,)：形状类别；params：make_shapes(..., return_params=True) 给出的 cx、cy、r
    【返回】(N, 4) 的 long 张量：[形状, 大小(0 小 1 大), 上下(0 上 1 下), 左右(0 左 1 右)]
    """
    return torch.stack([labels, (params["r"] > BIG_R).long(), (params["cy"] >= MID).long(),
                        (params["cx"] >= MID).long()], dim=1)


def caption(attr, template=0):
    """一行属性 → 描述（词列表）。"""
    s, z, v, h = (SHAPE_WORDS[int(attr[0])], SIZE_WORDS[int(attr[1])], VERT_WORDS[int(attr[2])],
                  HORIZ_WORDS[int(attr[3])])
    return TEMPLATES[template](s, z, v, h)


def make_captioned(n, seed, template=0):
    """
    生成 n 张带描述的图。
    【返回】(images (n, 1, 28, 28), attrs (n, 4), captions：词列表的列表)
    【边界】seed = 1 的图和第 6–9 讲生成模型的训练图完全相同（同一个 make_shapes 调用）。
    """
    x, y, p = make_shapes(n, seed=seed, return_params=True)
    a = attributes(y, p)
    return x, a, [caption(a[i], template) for i in range(n)]


def read_attributes(images, batch=500):
    """用第 6 讲的形状拟合读出每张图的属性 (N, 4)。约 2 秒 / 256 张。"""
    out = []
    for i in range(0, len(images), batch):
        _, k, p = fit_shapes(images[i:i + batch])
        out.append(torch.stack([k, (p[:, 2] > BIG_R).long(), (p[:, 1] >= MID).long(), (p[:, 0] >= MID).long()], 1))
    return torch.cat(out)


def attribute_accuracy(pred, target):
    """【返回】(每个属性的准确率列表, 4 个属性全对的比例)"""
    eq = (pred == target).float()
    return eq.mean(0).tolist(), eq.all(1).float().mean().item()