"""Contrastive loss functions для dual-encoder retrieval.

InfoNCE (он же NT-Xent, CLIP loss) — стандарт для contrastive learning.
Также включён вариант weighted-InfoNCE из Game4Loc для GTA-UAV.

Как работает InfoNCE:
    У тебя батч из B пар (text_i, image_i). Строим матрицу сходства B×B.
    Правильная пара — на диагонали (text_0↔image_0, text_1↔image_1, ...).
    Все остальные B-1 пар в строке — «негативные примеры».
    
    Loss = CrossEntropy по строкам (text→image) + CrossEntropy по столбцам (image→text).
    
    Это заставляет модель:
    - Приближать правильные пары (высокое сходство)
    - Отдалять неправильные (низкое сходство)
    
    Чем больше B, тем больше негативов, тем сильнее сигнал.
    Поэтому batch_size=64 важен.
"""
from __future__ import annotations

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


class InfoNCELoss(nn.Module):
    """Symmetric InfoNCE loss (как в CLIP).
    
    Считает loss в обе стороны:
        - text→image: для каждого текста, какое изображение правильное?
        - image→text: для каждого изображения, какой текст правильный?
    Усредняет оба направления.
    
    Args:
        label_smoothing: Сглаживание меток. 0.0 = жёсткие метки (стандарт CLIP).
            Небольшое значение (0.1) может помочь при шумных описаниях.
    """

    def __init__(self, label_smoothing: float = 0.0) -> None:
        super().__init__()
        self.label_smoothing = label_smoothing

    def forward(self, logits: torch.Tensor) -> dict[str, torch.Tensor]:
        """
        Args:
            logits: Матрица сходства [B, B] = text_emb @ image_emb.T / τ.
                    Правильные пары на диагонали.
        
        Returns:
            dict:
                loss:     Средний loss (скаляр).
                loss_t2i: Text-to-image loss.
                loss_i2t: Image-to-text loss.
                acc_t2i:  Accuracy text→image (для мониторинга).
                acc_i2t:  Accuracy image→text.
        """
        B = logits.shape[0]
        # Метки: правильная пара для i-го текста — i-е изображение
        labels = torch.arange(B, device=logits.device)

        # Text → Image: строки матрицы
        loss_t2i = F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)
        # Image → Text: столбцы матрицы (транспонируем)
        loss_i2t = F.cross_entropy(logits.T, labels, label_smoothing=self.label_smoothing)

        loss = (loss_t2i + loss_i2t) / 2.0

        # Accuracy для мониторинга (не для backprop)
        with torch.no_grad():
            acc_t2i = (logits.argmax(dim=1) == labels).float().mean()
            acc_i2t = (logits.T.argmax(dim=1) == labels).float().mean()

        return {
            "loss": loss,
            "loss_t2i": loss_t2i,
            "loss_i2t": loss_i2t,
            "acc_t2i": acc_t2i,
            "acc_i2t": acc_i2t,
        }


class WeightedInfoNCELoss(nn.Module):
    """Weighted InfoNCE из Game4Loc для GTA-UAV.
    
    В GTA-UAV не все пары одинаково «правильные» — есть positive (полное
    совпадение) и semi-positive (частичное перекрытие по IoU). Веса
    отражают степень совпадения.
    
    Вместо жёстких меток [0,0,1,0,...] используются мягкие [0.1, 0.0, 0.7, 0.2,...],
    где вес пропорционален IoU перекрытия drone-satellite пар.
    
    Args:
        label_smoothing: Дополнительное сглаживание.
    """

    def __init__(self, label_smoothing: float = 0.0) -> None:
        super().__init__()
        self.label_smoothing = label_smoothing

    def forward(
        self,
        logits: torch.Tensor,
        weights: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        """
        Args:
            logits:  [B, B] матрица сходства.
            weights: [B, B] матрица весов (soft labels). Если None → обычный InfoNCE.
        """
        B = logits.shape[0]

        if weights is None:
            # Fallback на обычный InfoNCE
            labels = torch.arange(B, device=logits.device)
            loss_t2i = F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)
            loss_i2t = F.cross_entropy(logits.T, labels, label_smoothing=self.label_smoothing)
        else:
            # Нормализуем веса в распределение вероятностей по каждой строке
            targets_t2i = F.softmax(weights, dim=1)   # для text→image
            targets_i2t = F.softmax(weights.T, dim=1)  # для image→text

            # KL-divergence вместо cross-entropy (для soft labels)
            log_probs_t2i = F.log_softmax(logits, dim=1)
            log_probs_i2t = F.log_softmax(logits.T, dim=1)

            loss_t2i = F.kl_div(log_probs_t2i, targets_t2i, reduction="batchmean")
            loss_i2t = F.kl_div(log_probs_i2t, targets_i2t, reduction="batchmean")

        loss = (loss_t2i + loss_i2t) / 2.0

        with torch.no_grad():
            labels = torch.arange(B, device=logits.device)
            acc_t2i = (logits.argmax(dim=1) == labels).float().mean()
            acc_i2t = (logits.T.argmax(dim=1) == labels).float().mean()

        return {
            "loss": loss,
            "loss_t2i": loss_t2i,
            "loss_i2t": loss_i2t,
            "acc_t2i": acc_t2i,
            "acc_i2t": acc_i2t,
        }