from typing import Any, Dict

import torch
from torch import nn
import torch.distributed as dist

from virtex.modules.label_smoothing import CrossEntropyLossWithLabelSmoothing
from virtex.modules.textual_heads import TextualHead
from virtex.modules.visual_backbones import VisualBackbone


class ImageTextContrastiveModel(nn.Module):
    def __init__(
        self,
        visual: VisualBackbone,
        textual: TextualHead,
        label_smoothing: float = 0.0
    ):
        super().__init__()
        self.visual = visual
        self.textual = textual
        self.padding_idx = self.textual.padding_idx

        self.visual_projection = nn.Linear(
            self.visual.visual_feature_size,
            self.textual.textual_feature_size,
            bias=False,
        )
        self.logit_scale = nn.Parameter(torch.log(torch.tensor(1/0.07)))
        self.loss = CrossEntropyLossWithLabelSmoothing(
            label_smoothing, ignore_index=self.padding_idx
        )

    def forward(self, batch: Dict[str, torch.Tensor]) -> Dict[str, Any]:

        # Check if logit_scale needs to be clipped from last iteration.
        self.logit_scale.data = torch.clamp(self.logit_scale.data, 0, 3.912)
        # 50 times

        # shape: (batch_size, channels, height, width)
        visual_features = self.visual(batch["image"])
        batch_size = visual_features.size(0)

        # shape: (batch_size, channels)
        visual_features = visual_features.mean(dim=[2, 3]).view(batch_size, -1)

        # shape: (batch_size, textual_feature_size)
        visual_features = self.visual_projection(visual_features)

        caption_tokens = batch["caption_tokens"]
        caption_lengths = batch["caption_lengths"]

        # shape: (batch_size, max_caption_length, hidden_size)
        textual_features = self.textual(caption_tokens, caption_lengths)

        # Take features from the first time-step (as BERT-* models do).
        # shape: (batch_size, hidden_size)
        textual_features = textual_features[:, 0, :]

        # Normalize visual and textual features.
        # shape: (batch_size, textual_feature_size)
        visual_features = visual_features / visual_features.norm(dim=-1, keepdim=True)
        textual_features = textual_features / textual_features.norm(
            dim=-1, keepdim=True
        )
        # Gather textual features from all processes into one large tensor to
        # increase negative samples for contrastive learning.
        gathered_textual_features = [
            torch.zeros_like(textual_features) for _ in range(dist.get_world_size())
        ]
        dist.all_gather(gathered_textual_features, textual_features)

        # Shift features of current rank to zeroth index for easy implementation.
        gathered_textual_features[0], gathered_textual_features[dist.get_rank()] = (
            gathered_textual_features[dist.get_rank()],
            gathered_textual_features[0],
        )
        # shape: (batch_size * world_size, textual_feature_size)
        gathered_textual_features = torch.cat(gathered_textual_features, dim=0)

        # Calculate pairwise cosine similarity as logits.
        logit_scale = self.logit_scale.exp()
        visual_logits = logit_scale * visual_features @ gathered_textual_features.t()

        # Targets are an identity matrix (image [i] should match with caption [i])
        visual_loss = self.loss(
            visual_logits, torch.arange(visual_logits.size(0)).to(visual_logits.device)
        )

        # Do the same thing for visual features.
        gathered_visual_features = [
            torch.zeros_like(visual_features) for _ in range(dist.get_world_size())
        ]
        dist.all_gather(gathered_visual_features, visual_features)

        gathered_visual_features[0], gathered_visual_features[dist.get_rank()] = (
            gathered_visual_features[dist.get_rank()],
            gathered_visual_features[0],
        )
        # shape: (batch_size * world_size, textual_feature_size)
        gathered_visual_features = torch.cat(gathered_visual_features, dim=0)

        # Calculate pairwise cosine similarity as logits.
        logit_scale = self.logit_scale.exp()
        textual_logits = logit_scale * textual_features @ gathered_visual_features.t()

        # Targets are an identity matrix (image [i] should match with caption [i])
        textual_loss = self.loss(
            textual_logits,
            torch.arange(textual_logits.size(0)).to(textual_logits.device),
        )
        loss = 0.5 * (visual_loss + textual_loss)
        output_dict: Dict[str, Any] = {
            "loss": loss,
            # Single scalar per batch for logging in training script.
            "loss_components": {"contrastive": loss.clone().detach()},
        }

        return output_dict