def make_synthetic_image(shape, color, size=16):
    """Renders a crude colored-shape 'image' as a tensor, purely
    to give the pipeline something concrete and free of external
    data dependencies to train and classify."""
    img = torch.zeros(3, size, size)
    color_idx = {"red": 0, "green": 1, "blue": 2}[color]
    if shape == "square":
        img[color_idx, 4:12, 4:12] = 1.0
    elif shape == "circle":
        yy, xx = torch.meshgrid(
            torch.arange(size), torch.arange(size), indexing="ij"
        )
        center = size / 2
        mask = ((xx - center) ** 2 + (yy - center) ** 2) < (size / 3) ** 2
        img[color_idx][mask] = 1.0
    return img


def tokenize(text, stoi):
    return [stoi[w] for w in text.lower().split() if w in stoi]


if __name__ == "__main__":
    torch.manual_seed(0)
    shapes = ["square", "circle"]
    colors = ["red", "green", "blue"]
    captions = [f"a {c} {s}" for s in shapes for c in colors]
    vocab = sorted(set(w for c in captions for w in c.split()))
    stoi = {w: i for i, w in enumerate(vocab)}

    image_encoder = ImageEncoder(embed_dim=32)
    text_encoder = TextEncoder(len(vocab), embed_dim=32)
    optimizer = torch.optim.Adam(
        list(image_encoder.parameters()) + list(text_encoder.parameters()),
        lr=1e-3,
    )

    images, token_lists = [], []
    for s in shapes:
        for c in colors:
            images.append(make_synthetic_image(s, c))
            token_lists.append(tokenize(f"a {c} {s}", stoi))

    images = torch.stack(images)
    flat_tokens = torch.tensor(
        [t for tokens in token_lists for t in tokens]
    )
    offsets = torch.tensor(
        [0] + list(torch.cumsum(
            torch.tensor([len(t) for t in token_lists]), dim=0
        )[:-1])
    )

    for epoch in range(300):
        img_emb = image_encoder(images)
        txt_emb = text_encoder(flat_tokens, offsets)
        loss = clip_contrastive_loss(img_emb, txt_emb)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if epoch % 100 == 0:
            print(f"epoch {epoch}  contrastive loss {loss.item():.4f}")

    # Zero-shot classification: classify a new red circle image
    # against every candidate caption, none seen paired at
    # test time (though all captions were seen during training
    # here, given the tiny synthetic vocabulary).
    test_image = make_synthetic_image("circle", "red").unsqueeze(0)
    with torch.no_grad():
        test_embed = image_encoder(test_image)
        candidate_tokens = torch.tensor(
            [t for tokens in token_lists for t in tokens]
        )
        candidate_offsets = offsets
        candidate_embeds = text_encoder(candidate_tokens, candidate_offsets)
        sims = (test_embed @ candidate_embeds.T).squeeze(0)
        best_idx = int(sims.argmax())
    print("\nTrue label: a red circle")
    print("Predicted label:", captions[best_idx])
