"""Loss functions and calibration metrics for the GNN router.

Includes heteroscedastic NLL, focal loss, Brier loss, label-smoothed BCE,
overconfidence penalty, uncertainty metrics, and calibration measurement.
"""

from __future__ import annotations

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F


def heteroscedastic_nll_loss(
    predictions: torch.Tensor,
    targets: torch.Tensor,
    log_var: torch.Tensor,
    reduction: str = "mean",
) -> torch.Tensor:
    """Compute heteroscedastic negative log-likelihood loss.

    This loss naturally learns calibrated uncertainty: the model is rewarded
    for predicting high variance on difficult examples where it makes errors.

    Formula: L = 0.5 * (log(sigma^2) + (y - mu)^2 / sigma^2)

    Args:
        predictions: Predicted values mu (num_samples,).
        targets: Ground truth values y (num_samples,).
        log_var: Predicted log variance log(sigma^2) (num_samples,).
        reduction: "mean", "sum", or "none".

    Returns:
        Heteroscedastic NLL loss.
    """
    variance = torch.exp(log_var)
    squared_error = (predictions - targets) ** 2

    # NLL = 0.5 * (log(sigma^2) + (y - mu)^2 / sigma^2)
    nll = 0.5 * (log_var + squared_error / (variance + 1e-8))

    if reduction == "mean":
        return nll.mean()
    elif reduction == "sum":
        return nll.sum()
    return nll


def overconfidence_penalty(
    predictions: torch.Tensor,
    targets: torch.Tensor,
    log_var: torch.Tensor,
    threshold: float = 0.5,
) -> torch.Tensor:
    """Explicit penalty for low uncertainty on incorrect predictions.

    This implements the user's "smart policy": penalize the model when it's
    confident (low variance) but wrong. Correct predictions are not penalized
    regardless of confidence level.

    Formula: penalty = mean((1 - sigma^2) * I(|error| > threshold))

    Args:
        predictions: Predicted values (num_samples,).
        targets: Ground truth values (num_samples,).
        log_var: Predicted log variance (num_samples,).
        threshold: Error threshold for "incorrect" (default 0.5 for binary).

    Returns:
        Mean overconfidence penalty.
    """
    variance = torch.exp(log_var).clamp(max=1.0)  # Clamp for stability
    error = (predictions - targets).abs()

    # Only penalize incorrect predictions (error > threshold)
    incorrect_mask = (error > threshold).float()

    # Penalty is high when variance is low (confident) and prediction is wrong
    # (1 - variance) ranges from 0 (uncertain) to 1 (confident)
    penalty = (1.0 - variance) * incorrect_mask

    return penalty.mean()


def compute_uncertainty_metrics(
    predictions: torch.Tensor,
    targets: torch.Tensor,
    log_var: torch.Tensor,
) -> dict[str, float]:
    """Compute metrics for evaluating uncertainty calibration.

    Good uncertainty should:
    1. Be high when predictions are wrong
    2. Be low when predictions are correct
    3. Correlate with actual error magnitude

    Args:
        predictions: Predicted values (num_samples,).
        targets: Ground truth values (num_samples,).
        log_var: Predicted log variance (num_samples,).

    Returns:
        Dictionary with:
        - mean_variance: Average predicted variance
        - mean_log_var: Average log variance
        - error_var_correlation: Spearman correlation between |error| and variance
        - overconfidence_rate: Fraction of wrong predictions with low variance
        - underconfidence_rate: Fraction of correct predictions with high variance
    """
    with torch.no_grad():
        variance = torch.exp(log_var)
        error = (predictions - targets).abs()

        # Basic stats
        mean_var = variance.mean().item()
        mean_log_var = log_var.mean().item()

        # Error-variance correlation (should be positive for good calibration)
        # Move to CPU for numpy correlation
        error_np = error.cpu().numpy()
        var_np = variance.cpu().numpy()

        # Spearman correlation (rank-based, robust to outliers)
        from scipy.stats import spearmanr
        if len(error_np) > 10:
            corr, _ = spearmanr(error_np, var_np)
            corr = float(corr) if not np.isnan(corr) else 0.0
        else:
            corr = 0.0

        # Overconfidence: wrong (error > 0.5) but confident (variance < median)
        median_var = variance.median()
        wrong_mask = error > 0.5
        confident_mask = variance < median_var
        overconf_rate = (wrong_mask & confident_mask).float().mean().item()

        # Underconfidence: correct (error < 0.5) but uncertain (variance > median)
        correct_mask = error <= 0.5
        uncertain_mask = variance >= median_var
        underconf_rate = (correct_mask & uncertain_mask).float().mean().item()

        return {
            "mean_variance": mean_var,
            "mean_log_var": mean_log_var,
            "error_var_correlation": corr,
            "overconfidence_rate": overconf_rate,
            "underconfidence_rate": underconf_rate,
        }


class FocalLoss(nn.Module):
    """Focal loss for handling class imbalance in difficulty prediction.

    Focal loss down-weights easy examples and focuses learning on hard
    examples. This is critical for difficulty prediction where:
    - 26.5% of queries are "easy" (minority class)
    - 73.5% of queries are "non_easy" (majority class)

    Formula: FL(p_t) = -alpha * (1 - p_t)^gamma * log(p_t)

    Where:
    - p_t is the model's confidence for the correct class
    - gamma controls focusing (gamma=2 is standard)
    - alpha balances class weights (alpha=0.25 for minority)

    When gamma=0, focal loss = BCE loss.
    As gamma increases, easy examples (high p_t) contribute less to loss.

    Note: This implementation is AMP-safe by casting to FP32 for BCE
    computation, as binary_cross_entropy is numerically unstable in FP16.
    """

    def __init__(self, alpha: float = 0.25, gamma: float = 2.0) -> None:
        """Initialize focal loss.

        Args:
            alpha: Weighting factor for positive class (easy queries).
                   Should be set to minority_class_ratio / total.
                   Default 0.25 ~ 26.5% easy queries.
            gamma: Focusing parameter. Higher = more focus on hard examples.
                   Standard value is 2.0.
        """
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(
        self,
        pred: torch.Tensor,
        target: torch.Tensor,
        sample_weight: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Compute focal loss with optional per-sample weighting.

        Args:
            pred: Predicted probabilities (num_samples,) in [0, 1].
            target: Ground truth labels (num_samples,) in {0, 1}.
            sample_weight: Optional per-sample weights (num_samples,).
                Useful for query-level weighting (e.g., upweight needle queries
                where only 1 LLM succeeds).

        Returns:
            Scalar focal loss.
        """
        # Cast to FP32 for numerically stable BCE (required for AMP compatibility)
        pred_f32 = pred.float()
        target_f32 = target.float()
        bce_loss = F.binary_cross_entropy(pred_f32, target_f32, reduction="none")

        # Confidence: p_t = p if y=1, else (1-p)
        p_t = pred * target + (1 - pred) * (1 - target)

        # Focal weight: (1 - p_t)^gamma
        focal_weight = (1 - p_t) ** self.gamma

        # Class balance weight: alpha if y=1, else (1-alpha)
        alpha_weight = self.alpha * target + (1 - self.alpha) * (1 - target)

        # Focal loss = focal_weight * alpha_weight * bce_loss
        focal_loss = focal_weight * alpha_weight * bce_loss

        # Apply sample-level weighting if provided
        if sample_weight is not None:
            focal_loss = focal_loss * sample_weight
            # Weighted mean: sum(loss * weight) / sum(weight)
            return focal_loss.sum() / sample_weight.sum().clamp(min=1e-8)

        return focal_loss.mean()


class BrierLoss(nn.Module):
    """Brier score loss for calibrated probability predictions.

    Brier score = mean((p - y)^2) directly measures calibration.
    Unlike focal loss, it doesn't down-weight confident predictions,
    making it ideal for learning well-calibrated probabilities.

    Properties:
    - Strictly proper scoring rule (optimal predictions = true probabilities)
    - Smooth gradients everywhere
    - Equally penalizes over-confidence and under-confidence
    """

    def __init__(self, label_smoothing: float = 0.0) -> None:
        """Initialize Brier loss.

        Args:
            label_smoothing: Smoothing factor in [0, 0.5). If > 0, converts
                hard labels {0, 1} to soft labels {eps, 1-eps} where
                eps = label_smoothing. This prevents overconfidence.
        """
        super().__init__()
        self.label_smoothing = label_smoothing

    def forward(
        self,
        pred: torch.Tensor,
        target: torch.Tensor,
    ) -> torch.Tensor:
        """Compute Brier score.

        Args:
            pred: Predicted probabilities (num_queries,) in [0, 1].
            target: Ground truth labels (num_queries,) in {0, 1}.

        Returns:
            Scalar Brier score (lower = better calibration).
        """
        # Apply label smoothing if enabled
        if self.label_smoothing > 0:
            target = target * (1 - self.label_smoothing) + 0.5 * self.label_smoothing

        # Brier score = mean squared error between predictions and targets
        return F.mse_loss(pred, target)


class LabelSmoothingBCELoss(nn.Module):
    """BCE loss with label smoothing for improved calibration.

    Label smoothing converts hard labels {0, 1} to soft labels
    {eps, 1-eps}, preventing the model from becoming overconfident.
    """

    def __init__(self, smoothing: float = 0.1) -> None:
        """Initialize label smoothing BCE.

        Args:
            smoothing: Smoothing factor in [0, 0.5). Hard labels y are
                transformed to y' = y * (1 - smoothing) + 0.5 * smoothing.
        """
        super().__init__()
        self.smoothing = smoothing

    def forward(
        self,
        pred: torch.Tensor,
        target: torch.Tensor,
    ) -> torch.Tensor:
        """Compute label-smoothed BCE loss."""
        if self.smoothing > 0:
            target = target * (1 - self.smoothing) + 0.5 * self.smoothing
        return F.binary_cross_entropy(pred, target)


def compute_calibration_metrics(
    pred: np.ndarray,
    target: np.ndarray,
    n_bins: int = 10,
) -> dict:
    """Compute comprehensive calibration metrics.

    Args:
        pred: Predicted probabilities (num_samples,) in [0, 1].
        target: Ground truth binary labels (num_samples,) in {0, 1}.
        n_bins: Number of bins for ECE/MCE computation.

    Returns:
        Dictionary with:
        - ece: Expected Calibration Error
        - mce: Maximum Calibration Error
        - brier: Brier score
        - reliability_bins: List of (bin_center, accuracy, confidence, count)
    """
    # Ensure numpy arrays
    pred = np.asarray(pred).flatten()
    target = np.asarray(target).flatten()

    # Brier score
    brier = np.mean((pred - target) ** 2)

    # Bin predictions for ECE/MCE
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    bin_lowers = bin_boundaries[:-1]
    bin_uppers = bin_boundaries[1:]

    ece = 0.0
    mce = 0.0
    reliability_bins = []

    for bin_lower, bin_upper in zip(bin_lowers, bin_uppers, strict=False):
        # Find samples in this bin
        in_bin = (pred > bin_lower) & (pred <= bin_upper)
        prop_in_bin = in_bin.mean()

        if prop_in_bin > 0:
            # Accuracy = fraction of positives in this bin
            accuracy = target[in_bin].mean()
            # Confidence = average predicted probability in this bin
            confidence = pred[in_bin].mean()
            # Calibration gap
            gap = np.abs(accuracy - confidence)

            ece += prop_in_bin * gap
            mce = max(mce, gap)

            reliability_bins.append({
                "bin_center": (bin_lower + bin_upper) / 2,
                "accuracy": float(accuracy),
                "confidence": float(confidence),
                "count": int(in_bin.sum()),
                "gap": float(gap),
            })

    return {
        "ece": float(ece),
        "mce": float(mce),
        "brier": float(brier),
        "reliability_bins": reliability_bins,
    }
