# Challenge 1: Zero-Shot Text-to-Image Retrieval
# Goal: Given a list of query captions and a catalog of images, rank the images based on their cosine
# similarities to each caption and evaluate performance using Recall@k metrics.
# Write a retrieval function that takes a query text embedding and a matrix of image embeddings,
# sorts the database images, and checks if the correct matching image index is within the top k ranks.
def evaluate_recall_at_k(text_embeddings, image_embeddings, k=1):
    # text_embeddings shape: [N, D]
    # image_embeddings shape: [N, D]
    num_queries = text_embeddings.size(0)
    correct_matches = 0

    # Step 1: Compute N x N Similarity Matrix
    sim = torch.matmul(text_embeddings, image_embeddings.t()) # [N, N]

    # Step 2: Retrieve Top-K Indices along Image axis for each Text Query
    top_k_indices = torch.topk(sim, k, dim=1).indices # [N, k]

    for i in range(num_queries):
        # The matching image is at index i
        if i in top_k_indices[i]:
            correct_matches += 1

    recall = correct_matches / num_queries
    return recall

'''
Challenge 2: Cross-Modal Cross-Attention Fusion Layer
Goal: Replace simple mean pooling with a dynamic Attention Mechanism. Create a PyTorch module
where sequence-based text representations serve as 'Queries', and spatio-temporal image patch
sequences serve as 'Keys' and 'Values'. This allows the model to attend to specific image coordinates
when interpreting specific words (e.g., matching the word 'dog' to a spatial canine coordinate patch).
'''
class CrossModalAttentionBlock(nn.Module):
    def __init__(self, embed_dim=256, num_heads=4):
        super().__init__()
        self.mha = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, text_seq, image_patches):
        # text_seq shape:      [Batch, SeqLength, EmbedDim] (Queries)
        # image_patches shape: [Batch, NumPatches, EmbedDim] (Keys / Values)

        # Execute dynamic cross-attention
        attn_out, weights = self.mha(
            query=text_seq,
            key=image_patches,
            value=image_patches
        )
        # Skip Connection and Layer Norm
        output = self.norm(text_seq + attn_out)
        return output, weights