"""Zero-cost signal extraction from routing logits.

This module provides functions to extract uncertainty, difficulty, and confidence
signals directly from model predictions, without requiring additional trainable
parameters or auxiliary heads.

Key insight: These signals emerge naturally from the softmax distribution. Training
separate heads introduces gradient interference that hurts primary task performance
(observed: -0.9% frac_oracle with aux heads vs. without).

Usage:
    logits = model(query_emb, llm_emb, ...)  # Shape: (num_queries, num_llms)
    uncertainty = compute_uncertainty_from_logits(logits)
    difficulty = compute_difficulty_from_logits(logits)
"""

from __future__ import annotations

import math

import torch
import torch.nn.functional as F


def compute_uncertainty_from_logits(
    logits: torch.Tensor,
    temperature: float = 1.0,
) -> torch.Tensor:
    """Extract uncertainty from routing logits via normalized entropy.

    Uncertainty measures how confident the router is in its prediction.
    High entropy in the softmax distribution indicates the model cannot
    clearly distinguish which LLM is best for this query.

    Args:
        logits: Raw routing scores with shape (num_queries, num_llms).
        temperature: Temperature for softmax scaling (use learned temperature
            from calibration for best results). Default 1.0.

    Returns:
        Uncertainty scores in [0, 1], shape (num_queries,).
        0 = perfectly certain (one-hot distribution).
        1 = maximally uncertain (uniform distribution).
    """
    if logits.dim() != 2:
        raise ValueError(
            f"Expected logits with shape (num_queries, num_llms), got {logits.shape}"
        )

    num_llms = logits.size(-1)
    if num_llms < 2:
        raise ValueError(f"Need at least 2 LLMs for meaningful uncertainty, got {num_llms}")

    # Scale by temperature and compute softmax
    probs = F.softmax(logits / temperature, dim=-1)

    # Entropy: -sum(p * log(p)), with numerical stability
    # Using torch.xlogy handles p=0 case correctly (0 * log(0) = 0)
    entropy = -torch.sum(torch.xlogy(probs, probs), dim=-1)

    # Normalize to [0, 1]: max entropy = log(num_llms) when uniform
    max_entropy = math.log(num_llms)
    normalized_uncertainty = entropy / max_entropy

    return normalized_uncertainty


def compute_difficulty_from_logits(
    logits: torch.Tensor,
    temperature: float = 1.0,
) -> torch.Tensor:
    """Extract difficulty from routing logits via inverse max probability.

    Difficulty measures how hard a query is to route. When the maximum
    probability is low (all LLMs score similarly), the query is intrinsically
    difficult because there's no clear winner.

    This is related to uncertainty but has different semantics:
    - Uncertainty: model's confidence in its prediction
    - Difficulty: intrinsic property of the query (some queries have no clear best LLM)

    Args:
        logits: Raw routing scores with shape (num_queries, num_llms).
        temperature: Temperature for softmax scaling. Default 1.0.

    Returns:
        Difficulty scores in [0, 1), shape (num_queries,).
        ~0 = easy (one LLM is clearly best).
        ~1 = hard (all LLMs are essentially equivalent).
    """
    if logits.dim() != 2:
        raise ValueError(
            f"Expected logits with shape (num_queries, num_llms), got {logits.shape}"
        )

    probs = F.softmax(logits / temperature, dim=-1)
    max_prob = probs.max(dim=-1).values

    # Difficulty is inverse of confidence
    # max_prob in [1/num_llms, 1], so difficulty in [0, 1 - 1/num_llms]
    difficulty = 1.0 - max_prob

    return difficulty


def compute_difficulty_from_logit_gap(
    logits: torch.Tensor,
) -> torch.Tensor:
    """Extract difficulty based on gap between top-1 and top-2 logits.

    This provides a more robust difficulty estimate that's invariant to
    temperature scaling. A small gap means the model can't distinguish
    between the best two options, indicating a difficult query.

    Args:
        logits: Raw routing scores with shape (num_queries, num_llms).

    Returns:
        Difficulty scores in [0, 1], shape (num_queries,).
        Normalized assuming typical gap range is [0, 5] for well-trained models.
    """
    if logits.dim() != 2:
        raise ValueError(
            f"Expected logits with shape (num_queries, num_llms), got {logits.shape}"
        )

    if logits.size(-1) < 2:
        raise ValueError("Need at least 2 LLMs to compute logit gap")

    # Sort logits in descending order
    sorted_logits, _ = torch.sort(logits, dim=-1, descending=True)
    top1 = sorted_logits[:, 0]
    top2 = sorted_logits[:, 1]
    gap = top1 - top2

    # Normalize: empirically, gaps > 5 indicate clear winners
    # Clamp to [0, 1] and invert: large gap = easy, small gap = hard
    normalized_gap = torch.clamp(gap / 5.0, 0.0, 1.0)
    difficulty = 1.0 - normalized_gap

    return difficulty


def compute_routing_confidence(
    logits: torch.Tensor,
    temperature: float = 1.0,
) -> torch.Tensor:
    """Extract routing confidence as max probability after temperature scaling.

    Confidence is complementary to difficulty: confidence = 1 - difficulty
    (approximately, due to softmax normalization).

    Args:
        logits: Raw routing scores with shape (num_queries, num_llms).
        temperature: Temperature for softmax scaling. Default 1.0.

    Returns:
        Confidence scores in (0, 1], shape (num_queries,).
        ~1/num_llms = low confidence (uniform distribution).
        ~1 = high confidence (certain about best LLM).
    """
    if logits.dim() != 2:
        raise ValueError(
            f"Expected logits with shape (num_queries, num_llms), got {logits.shape}"
        )

    probs = F.softmax(logits / temperature, dim=-1)
    confidence = probs.max(dim=-1).values

    return confidence


def compute_routing_signals(
    logits: torch.Tensor,
    temperature: float = 1.0,
) -> dict[str, torch.Tensor]:
    """Compute all routing signals in a single pass.

    Efficient batch computation of all diagnostic signals.

    Args:
        logits: Raw routing scores with shape (num_queries, num_llms).
        temperature: Temperature for softmax scaling. Default 1.0.

    Returns:
        Dictionary containing:
        - 'uncertainty': Normalized entropy, shape (num_queries,)
        - 'difficulty': 1 - max(probs), shape (num_queries,)
        - 'difficulty_gap': Gap-based difficulty, shape (num_queries,)
        - 'confidence': max(probs), shape (num_queries,)
        - 'chosen_idx': Index of chosen LLM, shape (num_queries,)
    """
    if logits.dim() != 2:
        raise ValueError(
            f"Expected logits with shape (num_queries, num_llms), got {logits.shape}"
        )

    num_llms = logits.size(-1)

    # Single softmax computation
    probs = F.softmax(logits / temperature, dim=-1)

    # Confidence and difficulty from max prob
    max_prob, chosen_idx = probs.max(dim=-1)
    confidence = max_prob
    difficulty = 1.0 - max_prob

    # Entropy-based uncertainty
    entropy = -torch.sum(torch.xlogy(probs, probs), dim=-1)
    max_entropy = math.log(num_llms)
    uncertainty = entropy / max_entropy

    # Gap-based difficulty (temperature-invariant)
    sorted_logits, _ = torch.sort(logits, dim=-1, descending=True)
    gap = sorted_logits[:, 0] - sorted_logits[:, 1]
    difficulty_gap = 1.0 - torch.clamp(gap / 5.0, 0.0, 1.0)

    return {
        "uncertainty": uncertainty,
        "difficulty": difficulty,
        "difficulty_gap": difficulty_gap,
        "confidence": confidence,
        "chosen_idx": chosen_idx,
    }


def stratify_by_difficulty(
    signals: dict[str, torch.Tensor],
    num_quantiles: int = 4,
) -> dict[str, torch.Tensor]:
    """Compute quantile masks for stratifying analysis by difficulty.

    Args:
        signals: Dictionary from compute_routing_signals.
        num_quantiles: Number of difficulty buckets (default 4 = quartiles).

    Returns:
        Dictionary with boolean masks for each quantile:
        - 'q0': Easiest queries (difficulty < 25th percentile)
        - 'q1': Easy-medium
        - 'q2': Medium-hard
        - 'q3': Hardest queries (difficulty >= 75th percentile)
    """
    difficulty = signals["difficulty"]
    device = difficulty.device

    # Compute quantile boundaries
    boundaries = torch.quantile(
        difficulty,
        torch.linspace(0, 1, num_quantiles + 1, device=device),
    )

    masks = {}
    for i in range(num_quantiles):
        lower = boundaries[i]
        upper = boundaries[i + 1]
        if i == num_quantiles - 1:
            # Include upper boundary for last quantile
            mask = (difficulty >= lower) & (difficulty <= upper)
        else:
            mask = (difficulty >= lower) & (difficulty < upper)
        masks[f"q{i}"] = mask

    return masks
