from __future__ import annotations

"""Weighted InfoNCE loss for GTA-UAV cross-view geo-localization.

Adapted from Game4Loc (https://github.com/Yux1angJi/GTA-UAV).
Uses per-sample label smoothing based on positive_weights (IoU/distance)
to handle partial overlap between drone and satellite crops.

Standard InfoNCE assumes strict 1-to-1 pairs and treats all non-diagonal
entries as negatives. In GTA-UAV, multiple satellite crops can validly
match one drone image (partial IoU overlap), causing false negatives.
WeightedInfoNCE softens this with adaptive label smoothing per sample.
"""

import math

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


@gin.configurable
class WeightedInfoNCELoss(nn.Module):
    """Weighted InfoNCE with adaptive per-sample label smoothing.

    For each sample i, eps_i = 1 - (1 - base_smoothing) / (1 + exp(-k * w_i))
    where w_i is the positive weight (e.g. IoU with matched satellite crop).
    Higher weight → lower eps → sharper target (strong positive).
    Lower weight → higher eps → softer target (weak/semi-positive).

    Args:
        temperature_init: Initial temperature (or learnable logit_scale).
        learnable_temperature: If True, temperature is learnable (CLIP-style).
        label_smoothing: Base label smoothing (used when no weights provided).
        k: Sigmoid steepness for weight → eps mapping.
        tau_min: Min clamp for learnable temperature.
        tau_max: Max clamp for learnable temperature.
    """

    def __init__(
        self,
        temperature_init: float = 0.07,
        learnable_temperature: bool = True,
        label_smoothing: float = 0.1,
        k: float = 5.0,
        tau_min: float = 0.01,
        tau_max: float = 0.5,
    ) -> None:
        super().__init__()
        self.label_smoothing = label_smoothing
        self.k = k
        self.tau_min = tau_min
        self.tau_max = tau_max
        self.learnable_temperature = learnable_temperature

        if learnable_temperature:
            self.logit_scale = nn.Parameter(
                torch.tensor(math.log(1.0 / temperature_init))
            )
        else:
            self.logit_scale = None
            self.temperature = temperature_init

    @property
    def current_temperature(self) -> float:
        if self.logit_scale is not None:
            tau = 1.0 / self.logit_scale.exp().clamp(
                min=1.0 / self.tau_max, max=1.0 / self.tau_min,
            ).item()
            return tau
        return self.temperature

    def _compute_eps(self, positive_weights: torch.Tensor | None, n: int) -> torch.Tensor | list[float]:
        """Compute per-sample label smoothing from positive weights."""
        if positive_weights is not None:
            # Higher weight → lower eps (sharper, stronger positive).
            return 1.0 - (1.0 - self.label_smoothing) / (1.0 + torch.exp(-self.k * positive_weights))
        return [self.label_smoothing] * n

    def _weighted_loss(
        self,
        sim_matrix: torch.Tensor,
        eps_all: torch.Tensor | list[float],
    ) -> torch.Tensor:
        """Weighted InfoNCE: per-sample interpolation between hard and uniform targets.

        For each row i:
          L_i = (1-eps_i) * [-sim[i,i] + logsumexp(sim[i,:])]
              + eps_i * [-mean(sim[i,:]) + logsumexp(sim[i,:])]
        """
        n = sim_matrix.shape[0]
        total_loss = torch.tensor(0.0, device=sim_matrix.device)
        for i in range(n):
            eps = eps_all[i] if isinstance(eps_all, list) else eps_all[i]
            logsumexp = torch.logsumexp(sim_matrix[i, :], dim=0)
            total_loss += (1 - eps) * (-sim_matrix[i, i] + logsumexp)
            total_loss += eps * (-sim_matrix[i, :].mean() + logsumexp)
        return total_loss / n

    def forward(
        self,
        embeddings: dict[str, torch.Tensor],
        epoch: int = 0,
        total_epochs: int = 1,
        positive_weights: torch.Tensor | None = None,
        queue_negatives: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        """Compute weighted InfoNCE loss.

        Args:
            embeddings: Dict with 'query' [B,D], 'gallery' [B,D], 'gate_q', 'gate_g'.
            positive_weights: Per-sample weight [B] (e.g. IoU with matched sat crop).
            queue_negatives: Extra negatives [Q,D] from memory bank (not used with weighted loss).
        """
        query = embeddings["query"].float()
        gallery = embeddings["gallery"].float()

        # Temperature.
        if self.learnable_temperature:
            clamped = self.logit_scale.float().clamp(
                min=math.log(1.0 / self.tau_max),
                max=math.log(1.0 / self.tau_min),
            )
            logit_scale = clamped.exp()
            tau = 1.0 / logit_scale
        else:
            logit_scale = 1.0 / self.temperature
            tau = self.temperature

        sim_q2g = logit_scale * query @ gallery.t()
        sim_g2q = sim_q2g.t()

        eps = self._compute_eps(positive_weights, query.shape[0])

        loss_q2g = self._weighted_loss(sim_q2g, eps)
        loss_g2q = self._weighted_loss(sim_g2q, eps)
        total = (loss_q2g + loss_g2q) / 2

        gate_q = embeddings.get("gate_q", 1.0)
        gate_g = embeddings.get("gate_g", 1.0)

        return {
            "total": total,
            "temperature": tau if isinstance(tau, torch.Tensor) else torch.tensor(tau, device=total.device),
            "gate_q": torch.tensor(gate_q, device=total.device) if not isinstance(gate_q, torch.Tensor) else gate_q.detach(),
            "gate_g": torch.tensor(gate_g, device=total.device) if not isinstance(gate_g, torch.Tensor) else gate_g.detach(),
        }