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


class ImageEncoder(nn.Module):
    """A small CNN stand-in for a full vision encoder."""

    def __init__(self, embed_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.proj = nn.Linear(32, embed_dim)

    def forward(self, x):
        h = self.net(x).flatten(1)
        return F.normalize(self.proj(h), dim=-1)


class TextEncoder(nn.Module):
    """A small embedding-bag stand-in for a full text Transformer."""

    def __init__(self, vocab_size, embed_dim=64):
        super().__init__()
        self.embed = nn.EmbeddingBag(vocab_size, embed_dim, mode="mean")
        self.proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, token_ids, offsets):
        h = self.embed(token_ids, offsets)
        return F.normalize(self.proj(h), dim=-1)


def clip_contrastive_loss(image_embeds, text_embeds, temperature=0.07):
    logits = image_embeds @ text_embeds.T / temperature
    n = logits.size(0)
    labels = torch.arange(n)
    loss_i2t = F.cross_entropy(logits, labels)
    loss_t2i = F.cross_entropy(logits.T, labels)
    return (loss_i2t + loss_t2i) / 2
