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

InfoNCE (он же NT-Xent, CLIP loss) — стандарт для contrastive learning.

Как работает InfoNCE здесь (drone-fused vs satellite):
    У тебя батч из B пар (drone_i, satellite_i). Строим матрицу сходства B×B.
    Правильная пара — на диагонали (drone_0↔satellite_0, drone_1↔satellite_1, ...).
    Все остальные B-1 пар в строке — «негативные примеры».

    Loss = CrossEntropy по строкам (drone→satellite) + CrossEntropy по столбцам (satellite→drone).
    Это заставляет модель приближать правильные пары и отдалять неправильные.
"""
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 в обе стороны:
        - drone→satellite (d2s): для каждого слитого запроса дрона, какой спутник правильный?
        - satellite→drone (s2d): для каждого спутника, какой запрос дрона правильный?
    Усредняет оба направления.

    Названия acc_d2s/acc_s2d отражают направления слияния:
    drone→satellite и satellite→drone.
    (text-to-image / image-to-text), хотя сейчас "text" — это слитый drone-эмбеддинг,
    а "image" — спутниковый эмбеддинг.

    Направления взвешиваются несимметрично (протокол §6.3): основное —
    q2g (drone→satellite, d2s) с весом 0.6, вспомогательное — g2q (s2d)
    с весом 0.4.

    Args:
        label_smoothing: Сглаживание меток (протокол: 0.1).
        w_d2s: Вес направления drone→satellite (q2g). По умолчанию 0.6.
        w_s2d: Вес направления satellite→drone (g2q). По умолчанию 0.4.
    """

    def __init__(
        self,
        label_smoothing: float = 0.1,
        w_d2s: float = 0.6,
        w_s2d: float = 0.4,
    ) -> None:
        super().__init__()
        self.label_smoothing = label_smoothing
        self.w_d2s = w_d2s
        self.w_s2d = w_s2d

    def forward(self, logits: torch.Tensor) -> dict[str, torch.Tensor]:
        """
        Args:
            logits: Матрица сходства [B, B] = drone_emb @ satellite_emb.T / τ.
                    Правильные пары на диагонали.

        Returns:
            dict:
                loss:     Средний loss (скаляр).
                loss_d2s: Drone-to-satellite loss.
                loss_s2d: Satellite-to-drone loss.
                acc_d2s:  Accuracy drone→satellite (для мониторинга).
                acc_s2d:  Accuracy satellite→drone.
        """
        B = logits.shape[0]
        labels = torch.arange(B, device=logits.device)

        loss_d2s = F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)
        loss_s2d = F.cross_entropy(logits.T, labels, label_smoothing=self.label_smoothing)

        loss = self.w_d2s * loss_d2s + self.w_s2d * loss_s2d

        with torch.no_grad():
            acc_d2s = (logits.argmax(dim=1) == labels).float().mean()
            acc_s2d = (logits.T.argmax(dim=1) == labels).float().mean()

        return {
            "loss": loss,
            "loss_d2s": loss_d2s,
            "loss_s2d": loss_s2d,
            "acc_d2s": acc_d2s,
            "acc_s2d": acc_s2d,
        }


class WeightedInfoNCELoss(nn.Module):
    """Weighted InfoNCE из Game4Loc для GTA-UAV.

    В GTA-UAV не все пары одинаково «правильные» — есть positive (полное
    совпадение) и semi-positive (частичное перекрытие по IoU). Веса
    отражают степень совпадения.

    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:
            labels = torch.arange(B, device=logits.device)
            loss_d2s = F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)
            loss_s2d = F.cross_entropy(logits.T, labels, label_smoothing=self.label_smoothing)
        else:
            targets_d2s = F.softmax(weights, dim=1)
            targets_s2d = F.softmax(weights.T, dim=1)

            log_probs_d2s = F.log_softmax(logits, dim=1)
            log_probs_s2d = F.log_softmax(logits.T, dim=1)

            loss_d2s = F.kl_div(log_probs_d2s, targets_d2s, reduction="batchmean")
            loss_s2d = F.kl_div(log_probs_s2d, targets_s2d, reduction="batchmean")

        loss = (loss_d2s + loss_s2d) / 2.0

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

        return {
            "loss": loss,
            "loss_d2s": loss_d2s,
            "loss_s2d": loss_s2d,
            "acc_d2s": acc_d2s,
            "acc_s2d": acc_s2d,
        }
