"""Pairwise ranking losses for LLM routing.

These losses directly optimize the ranking constraint: score(winner) > score(loser).
More effective than pointwise losses when labels are binary (0/1) because they
create O(N*M) gradient signals per query where N correct, M incorrect.

Key insight: On a query with 3/7 correct LLMs:
- Pointwise loss: 7 gradients (one per LLM)
- Pairwise loss: 3*4=12 gradients (every correct/incorrect pair)
"""

from __future__ import annotations

import torch
import torch.nn.functional as F


def pairwise_margin_loss(
    scores: torch.Tensor,
    targets: torch.Tensor,
    margin: float = 0.5,
    threshold: float = 0.5,
) -> torch.Tensor:
    """Compute pairwise margin loss for ranking.

    For each query, enforces: score(winner) > score(loser) + margin
    where winner/loser are determined by targets > threshold.

    This is the MarginRankingLoss applied across all valid pairs.

    Args:
        scores: Model scores, shape (num_queries, num_llms).
        targets: Target rewards, shape (num_queries, num_llms).
        margin: Minimum score gap between winner and loser. Default 0.5.
        threshold: Threshold for binarizing targets. Default 0.5.

    Returns:
        Scalar loss value (mean over all valid pairs).
    """
    num_queries, num_llms = scores.shape
    device = scores.device

    # Binarize targets: 1 = winner, 0 = loser
    winners = (targets > threshold).float()
    losers = 1.0 - winners

    # For efficiency, compute all pairwise differences
    # scores_i - scores_j for all (i, j) pairs
    score_diff = scores.unsqueeze(2) - scores.unsqueeze(1)  # (Q, L, L)

    # Mask: 1 where targets_i > targets_j (winner vs loser)
    # This creates the "should be higher" constraint
    valid_pair_mask = winners.unsqueeze(2) * losers.unsqueeze(1)  # (Q, L, L)

    # Hinge loss: max(0, margin - (score_winner - score_loser))
    pair_loss = F.relu(margin - score_diff)  # (Q, L, L)

    # Apply mask and compute mean over valid pairs
    masked_loss = pair_loss * valid_pair_mask
    num_valid_pairs = valid_pair_mask.sum()

    if num_valid_pairs > 0:
        return masked_loss.sum() / num_valid_pairs
    else:
        return torch.tensor(0.0, device=device)


def pairwise_logistic_loss(
    scores: torch.Tensor,
    targets: torch.Tensor,
    threshold: float = 0.5,
    temperature: float = 1.0,
) -> torch.Tensor:
    """Compute pairwise logistic (Bradley-Terry) loss.

    For each (winner, loser) pair, models P(winner > loser) via logistic:
        P(i > j) = sigmoid((score_i - score_j) / temperature)

    Loss = -log P(winner > loser) for all valid pairs.

    This is smoother than margin loss and provides gradients everywhere.

    Args:
        scores: Model scores, shape (num_queries, num_llms).
        targets: Target rewards, shape (num_queries, num_llms).
        threshold: Threshold for binarizing targets. Default 0.5.
        temperature: Softness of preference. Lower = sharper. Default 1.0.

    Returns:
        Scalar loss value (mean over all valid pairs).
    """
    device = scores.device

    winners = (targets > threshold).float()
    losers = 1.0 - winners

    # Pairwise score differences
    score_diff = scores.unsqueeze(2) - scores.unsqueeze(1)  # (Q, L, L)

    # Valid pairs mask
    valid_pair_mask = winners.unsqueeze(2) * losers.unsqueeze(1)  # (Q, L, L)

    # Negative log-likelihood of correct ranking
    # -log(sigmoid(diff/T)) = -log(1/(1+exp(-diff/T))) = log(1+exp(-diff/T))
    # Use logsigmoid for numerical stability
    pair_nll = -F.logsigmoid(score_diff / temperature)  # (Q, L, L)

    masked_loss = pair_nll * valid_pair_mask
    num_valid_pairs = valid_pair_mask.sum()

    if num_valid_pairs > 0:
        return masked_loss.sum() / num_valid_pairs
    else:
        return torch.tensor(0.0, device=device)


def lambda_rank_loss(
    scores: torch.Tensor,
    targets: torch.Tensor,
    threshold: float = 0.5,
) -> torch.Tensor:
    """Compute LambdaRank-style loss with NDCG weighting.

    Like pairwise logistic, but weights each pair by the NDCG delta
    that would result from swapping the pair. This focuses learning
    on pairs that matter most for the ranking metric.

    For binary relevance: NDCG delta ≈ |log(rank_i+1) - log(rank_j+1)|

    Args:
        scores: Model scores, shape (num_queries, num_llms).
        targets: Target rewards, shape (num_queries, num_llms).
        threshold: Threshold for binarizing targets. Default 0.5.

    Returns:
        Scalar loss value.
    """
    device = scores.device
    num_queries, num_llms = scores.shape

    winners = (targets > threshold).float()
    losers = 1.0 - winners
    valid_pair_mask = winners.unsqueeze(2) * losers.unsqueeze(1)

    # Compute current ranks (1 = best, num_llms = worst)
    # Higher score = lower rank number
    _, order = scores.sort(dim=1, descending=True)
    ranks = torch.zeros_like(scores)
    ranks.scatter_(1, order, torch.arange(1, num_llms + 1, device=device).float().unsqueeze(0).expand(num_queries, -1))

    # NDCG delta for swapping positions i and j
    # Delta = |1/log(rank_i+1) - 1/log(rank_j+1)|
    log_ranks = torch.log2(ranks + 1)
    dcg_contribution = 1.0 / log_ranks
    delta = (dcg_contribution.unsqueeze(2) - dcg_contribution.unsqueeze(1)).abs()  # (Q, L, L)

    # Pairwise logistic loss weighted by delta
    score_diff = scores.unsqueeze(2) - scores.unsqueeze(1)
    pair_nll = -F.logsigmoid(score_diff)

    weighted_loss = pair_nll * delta * valid_pair_mask
    num_valid_pairs = valid_pair_mask.sum()

    if num_valid_pairs > 0:
        return weighted_loss.sum() / num_valid_pairs
    else:
        return torch.tensor(0.0, device=device)


def contrastive_ranking_loss(
    scores: torch.Tensor,
    targets: torch.Tensor,
    threshold: float = 0.5,
    temperature: float = 0.1,
) -> torch.Tensor:
    """InfoNCE-style contrastive loss for ranking.

    For each query:
    - Positive: highest-reward LLM
    - Negatives: all other LLMs

    Loss = -log(exp(score_pos/T) / sum(exp(score_i/T)))

    This naturally learns to separate the best from the rest.

    Args:
        scores: Model scores, shape (num_queries, num_llms).
        targets: Target rewards, shape (num_queries, num_llms).
        threshold: Not used (for API consistency). Positive = argmax(targets).
        temperature: Softmax temperature. Lower = harder negatives. Default 0.1.

    Returns:
        Scalar loss value.
    """
    device = scores.device
    num_queries, num_llms = scores.shape

    # Positive = best LLM per query
    best_idx = targets.argmax(dim=1)  # (Q,)

    # Compute log-softmax over LLMs
    log_probs = F.log_softmax(scores / temperature, dim=1)  # (Q, L)

    # Extract log-prob of the positive (best) LLM
    positive_log_prob = log_probs.gather(1, best_idx.unsqueeze(1)).squeeze(1)  # (Q,)

    # InfoNCE loss = -log_prob of positive = cross-entropy with argmax target
    return -positive_log_prob.mean()


# Hybrid loss combining listwise and pairwise
def hybrid_ranking_loss(
    scores: torch.Tensor,
    targets: torch.Tensor,
    margin: float = 0.5,
    pairwise_weight: float = 0.5,
    temperature: float = 0.1,
) -> torch.Tensor:
    """Hybrid loss: listwise KL + pairwise margin.

    Combines the benefits of both:
    - Listwise KL: Smooth, learns full distribution
    - Pairwise margin: Direct ranking constraint, more signal from binary labels

    Args:
        scores: Model scores, shape (num_queries, num_llms).
        targets: Target rewards (can be continuous), shape (num_queries, num_llms).
        margin: Pairwise margin. Default 0.5.
        pairwise_weight: Weight for pairwise term. Default 0.5.
        temperature: Softmax temperature for soft targets. Default 0.1.

    Returns:
        Combined loss value.
    """
    # Listwise: KL divergence to soft targets
    soft_targets = F.softmax(targets / temperature, dim=1).clamp_min(1e-8)
    log_probs = F.log_softmax(scores, dim=1)
    kl_loss = F.kl_div(log_probs, soft_targets, reduction='batchmean')

    # Pairwise: Margin loss on binary
    pair_loss = pairwise_margin_loss(scores, targets, margin=margin, threshold=0.5)

    return (1 - pairwise_weight) * kl_loss + pairwise_weight * pair_loss
