"""GNN training orchestrator for LLM routing prediction.

Contains GNNPredictor, the main class that trains and evaluates the GNN router
model. It composes layers, architectures, losses, and data conversion utilities
from sibling modules.
"""

from __future__ import annotations

import logging
import math
import random
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
import torch
from sklearn.metrics import f1_score, roc_auc_score
from torch import nn
from torch.nn import functional as F
from torch.nn.utils import clip_grad_norm_
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR

from embergraphrouter.model.architectures import (
    EncoderDecoderNet,
    GATEncoderDecoderNet,
)
from embergraphrouter.model.layers import ModelEMA
from embergraphrouter.model.losses import (
    BrierLoss,
    FocalLoss,
    LabelSmoothingBCELoss,
    compute_calibration_metrics,
    compute_uncertainty_metrics,
    heteroscedastic_nll_loss,
    overconfidence_penalty,
)
from embergraphrouter.model.optimizers import SAM, get_optimizer
from embergraphrouter.model.pairwise_losses import (
    lambda_rank_loss,
    pairwise_logistic_loss,
    pairwise_margin_loss,
)

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    import wandb as wandb_module


class GNNPredictor:
    """Trains and evaluates the GNN router model.

    The model learns to predict which LLM will perform best for a given query,
    conditioned on user preferences for performance, cost, and latency.
    """

    def __init__(
        self,
        query_feature_dim: int,
        llm_feature_dim: int,
        hidden_features_size: int,
        in_edges_size: int,
        wandb: wandb_module,
        config: dict,
        device: torch.device | str,
    ) -> None:
        dropout = float(config.get("dropout", 0.0))
        num_gnn_layers = int(config.get("num_gnn_layers", 2))
        self.use_difficulty_head = config.get("use_difficulty_head", False)
        self.difficulty_head_mode = config.get("difficulty_head_mode", "regression")
        self.difficulty_loss_weight = float(config.get("difficulty_loss_weight", 0.1))
        self.difficulty_entropy_weight = float(config.get("difficulty_entropy_weight", 0.0))
        self.use_focal_loss = config.get("use_focal_loss", False)

        # Task-based sample weighting for oversampling hard tasks (e.g., BBEH)
        # Format: {"bbeh_": 5.0} means tasks starting with "bbeh_" get 5x weight
        self.task_weights: dict[str, float] = config.get("task_weights", {})

        # Architecture selection: GCN (GeneralConv) or GAT (GATv2Conv)
        gnn_architecture = config.get("gnn_architecture", "gcn").lower()

        # Hidden channel multiplier (channels = embedding_dim * multiplier)
        hidden_multiplier = int(config.get("hidden_multiplier", 2))

        # Query MLP depth and input dropout (data augmentation)
        query_mlp_layers = int(config.get("query_mlp_layers", 1))
        llm_mlp_layers = int(config.get("llm_mlp_layers", 1))
        input_dropout = float(config.get("input_dropout", 0.0))

        # Advanced MLP options
        use_residual = bool(config.get("use_residual", False))
        use_swiglu = bool(config.get("use_swiglu", False))
        use_variational = bool(config.get("use_variational", False))
        vib_beta = float(config.get("vib_beta", 0.001))

        use_contrastive = bool(config.get("use_contrastive", False))
        contrastive_weight = float(config.get("contrastive_weight", 0.1))
        use_score_head = bool(config.get("use_score_head", False))
        use_gated_skip = bool(config.get("use_gated_skip", False))
        use_query_moe = bool(config.get("use_query_moe", False))
        moe_num_experts = int(config.get("moe_num_experts", 4))
        moe_top_k = int(config.get("moe_top_k", 2))
        use_uniformity_loss = bool(config.get("use_uniformity_loss", False))
        uniformity_weight = float(config.get("uniformity_weight", 0.05))
        self.use_llm_balance_loss = bool(config.get("use_llm_balance_loss", False))
        self.llm_balance_weight = float(config.get("llm_balance_weight", 0.1))
        self.quality_entropy_weight = float(config.get("quality_entropy_weight", 0.0))

        # Dynamic bias adjustment: loss-free anti-collapse from DeepSeek research
        # Instead of auxiliary loss, apply per-LLM bias based on selection frequency
        # LLMs selected too often get negative bias, forcing model to explore others
        self.use_dynamic_bias = bool(config.get("use_dynamic_bias", False))
        self.dynamic_bias_lr = float(config.get("dynamic_bias_lr", 0.1))
        self.dynamic_bias_ema = float(config.get("dynamic_bias_ema", 0.99))
        self.llm_num = int(config.get("llm_num", 7))
        # Initialize bias tensor (will be moved to device later)
        self._llm_selection_bias: torch.Tensor | None = None
        self._llm_selection_counts: torch.Tensor | None = None

        # Uncertainty estimation for heteroscedastic predictions
        # Staged training: delay uncertainty loss until quality head stabilizes
        self.use_uncertainty_head = bool(config.get("use_uncertainty_head", False))
        self.uncertainty_weight = float(config.get("uncertainty_weight", 0.1))
        self.overconfidence_penalty_weight = float(
            config.get("overconfidence_penalty_weight", 0.1)
        )
        # Warmup epochs before uncertainty loss kicks in (0 = immediate)
        self.uncertainty_warmup_epochs = int(config.get("uncertainty_warmup_epochs", 0))
        # Whether to freeze quality backbone when training uncertainty
        self.freeze_backbone_for_uncertainty = bool(
            config.get("freeze_backbone_for_uncertainty", False)
        )

        anti_collapse_enabled = []
        if use_contrastive:
            anti_collapse_enabled.append(f"contrastive(w={contrastive_weight})")
        if use_uniformity_loss:
            anti_collapse_enabled.append(f"uniformity(w={uniformity_weight})")
        if self.use_llm_balance_loss:
            anti_collapse_enabled.append(f"llm_balance(w={self.llm_balance_weight})")
        if self.quality_entropy_weight > 0:
            anti_collapse_enabled.append(f"quality_entropy(w={self.quality_entropy_weight})")
        if self.use_uncertainty_head:
            anti_collapse_enabled.append(
                f"uncertainty(w={self.uncertainty_weight}, "
                f"overconf={self.overconfidence_penalty_weight})"
            )
        if self.use_dynamic_bias:
            anti_collapse_enabled.append(
                f"dynamic_bias(lr={self.dynamic_bias_lr}, ema={self.dynamic_bias_ema})"
            )
        if anti_collapse_enabled:
            logger.info("Anti-collapse mechanisms: %s", ", ".join(anti_collapse_enabled))
        else:
            logger.info("Anti-collapse mechanisms: NONE (consider enabling for better routing)")

        if gnn_architecture == "gat":
            num_attention_heads = int(config.get("num_attention_heads", 4))
            attention_dropout = float(config.get("attention_dropout", 0.1))
            self.model = GATEncoderDecoderNet(
                query_feature_dim=query_feature_dim,
                llm_feature_dim=llm_feature_dim,
                hidden_features=hidden_features_size,
                in_edges=in_edges_size,
                dropout=dropout,
                num_gnn_layers=num_gnn_layers,
                use_difficulty_head=self.use_difficulty_head,
                difficulty_head_mode=self.difficulty_head_mode,
                num_attention_heads=num_attention_heads,
                attention_dropout=attention_dropout,
                hidden_multiplier=hidden_multiplier,
                query_mlp_layers=query_mlp_layers,
                llm_mlp_layers=llm_mlp_layers,
                input_dropout=input_dropout,
                use_residual=use_residual,
                use_swiglu=use_swiglu,
                use_variational=use_variational,
                vib_beta=vib_beta,
                use_gated_skip=use_gated_skip,
                use_query_moe=use_query_moe,
                moe_num_experts=moe_num_experts,
                moe_top_k=moe_top_k,
                use_contrastive=use_contrastive,
                contrastive_weight=contrastive_weight,
                use_uniformity_loss=use_uniformity_loss,
                uniformity_weight=uniformity_weight,
                use_score_head=use_score_head,
                use_uncertainty_head=self.use_uncertainty_head,
            ).to(device)
            logger.info(
                "Using GAT architecture: %d layers, %d attention heads, "
                "hidden_mult=%d, dropout=%.2f, query_mlp=%d, llm_mlp=%d, "
                "input_drop=%.2f, residual=%s, swiglu=%s, vib=%s, score_head=%s",
                num_gnn_layers,
                num_attention_heads,
                hidden_multiplier,
                attention_dropout,
                query_mlp_layers,
                llm_mlp_layers,
                input_dropout,
                use_residual,
                use_swiglu,
                use_variational,
                use_score_head,
            )
        else:
            self.model = EncoderDecoderNet(
                query_feature_dim=query_feature_dim,
                llm_feature_dim=llm_feature_dim,
                hidden_features=hidden_features_size,
                in_edges=in_edges_size,
                dropout=dropout,
                num_gnn_layers=num_gnn_layers,
                use_difficulty_head=self.use_difficulty_head,
                difficulty_head_mode=self.difficulty_head_mode,
                hidden_multiplier=hidden_multiplier,
                query_mlp_layers=query_mlp_layers,
                input_dropout=input_dropout,
            ).to(device)
            logger.info(
                "Using GCN architecture: %d layers, hidden_mult=%d, dropout=%.2f, "
                "query_mlp=%d, input_drop=%.2f",
                num_gnn_layers,
                hidden_multiplier,
                dropout,
                query_mlp_layers,
                input_dropout,
            )

        # Multi-GPU support via DataParallel
        self.use_multi_gpu = config.get("use_multi_gpu", False)
        if self.use_multi_gpu and torch.cuda.device_count() > 1:
            num_gpus = torch.cuda.device_count()
            self.model = nn.DataParallel(self.model)
            logger.info(
                "Multi-GPU enabled: using DataParallel across %d GPUs", num_gpus
            )

        # Difficulty loss function
        if self.use_difficulty_head:
            # Get loss configuration
            label_smoothing = float(config.get("label_smoothing", 0.0))
            difficulty_loss_type = config.get("difficulty_loss_type", "auto")

            # Auto-select based on legacy use_focal_loss flag
            if difficulty_loss_type == "auto":
                if self.use_focal_loss and self.difficulty_head_mode == "classification":
                    difficulty_loss_type = "focal"
                elif self.difficulty_head_mode == "classification":
                    difficulty_loss_type = "bce"
                else:
                    difficulty_loss_type = "mse"

            # Set up criterion based on loss type
            if difficulty_loss_type == "focal":
                focal_alpha = float(config.get("focal_alpha", 0.25))
                focal_gamma = float(config.get("focal_gamma", 2.0))
                self.difficulty_criterion = FocalLoss(
                    alpha=focal_alpha,
                    gamma=focal_gamma,
                )
                logger.info(
                    "Difficulty head: classification with focal loss "
                    "(alpha=%.2f, gamma=%.2f)",
                    focal_alpha,
                    focal_gamma,
                )
            elif difficulty_loss_type == "brier":
                # Brier score - optimal for calibration
                self.difficulty_criterion = BrierLoss(label_smoothing=label_smoothing)
                logger.info(
                    "Difficulty head: classification with Brier loss "
                    "(label_smoothing=%.2f)",
                    label_smoothing,
                )
            elif difficulty_loss_type == "label_smoothing_bce":
                self.difficulty_criterion = LabelSmoothingBCELoss(smoothing=label_smoothing)
                logger.info(
                    "Difficulty head: classification with label-smoothed BCE "
                    "(smoothing=%.2f)",
                    label_smoothing,
                )
            elif difficulty_loss_type == "bce":
                self.difficulty_criterion = nn.BCELoss()
                logger.info("Difficulty head: classification with BCE loss")
            else:  # mse for regression
                self.difficulty_criterion = nn.MSELoss()
                logger.info("Difficulty head: regression with MSE loss")

            logger.info(
                "Difficulty head enabled: %d GNN layers, loss_weight=%.2f",
                num_gnn_layers,
                self.difficulty_loss_weight,
            )

        self.device = device
        self.wandb = wandb
        self.config = config
        self.llm_num = int(config["llm_num"])

        # Quality-only training mode: learn P(correct | query, LLM) directly
        # Cost and latency are applied at inference time via user preferences
        self.quality_only_training = config.get("quality_only_training", False)
        if self.quality_only_training:
            logger.info(
                "Quality-only training enabled: learning calibrated P(correct) per LLM. "
                "Cost/latency will be composed at inference time."
            )

        # Needle query weighting: upweight queries where few LLMs succeed
        # These "needle in haystack" queries benefit most from focused learning
        self.use_needle_weighting = config.get("use_needle_weighting", False)
        if self.use_needle_weighting:
            logger.info(
                "Needle query weighting enabled: queries with 1 correct LLM get %dx weight",
                self.llm_num,
            )

        # Importance weighting for listwise loss: upweight informative queries
        # Strategies:
        #   - "uniform": All queries weighted equally (baseline)
        #   - "needle": Inverse frequency (7x for 1-correct, 1x for 7-correct)
        #   - "variance": Weight by performance variance across LLMs
        self.importance_weighting = config.get("importance_weighting", "uniform").lower()
        self.zero_oracle_importance = float(config.get("zero_oracle_importance", 1.0))
        if self.importance_weighting != "uniform":
            logger.info(
                "Importance weighting: %s (zero_oracle_weight=%.1f)",
                self.importance_weighting,
                self.zero_oracle_importance,
            )

        # Optimizer selection with specialized options for routing
        self.optimizer_type = config.get("optimizer", "adamw").lower()
        optimizer_kwargs = {}

        # Optimizer-specific configuration
        if self.optimizer_type == "muon":
            optimizer_kwargs["momentum"] = float(config.get("muon_momentum", 0.95))
            optimizer_kwargs["ns_iterations"] = int(config.get("muon_ns_iterations", 5))
            optimizer_kwargs["min_dim_for_orth"] = int(
                config.get("muon_min_dim_for_orth", 128)
            )
            logger.info(
                "Using Muon optimizer (momentum=%.2f, ns_iter=%d) for collapse prevention",
                optimizer_kwargs["momentum"],
                optimizer_kwargs["ns_iterations"],
            )
        elif self.optimizer_type == "sam":
            optimizer_kwargs["rho"] = float(config.get("sam_rho", 0.05))
            optimizer_kwargs["adaptive"] = bool(config.get("sam_adaptive", False))
            optimizer_kwargs["base_optimizer"] = config.get("sam_base_optimizer", "adamw")
            logger.info(
                "Using SAM optimizer (rho=%.3f, adaptive=%s) for flat minima",
                optimizer_kwargs["rho"],
                optimizer_kwargs["adaptive"],
            )
        elif self.optimizer_type == "lion":
            optimizer_kwargs["betas"] = tuple(config.get("lion_betas", (0.9, 0.99)))
            logger.info("Using Lion optimizer for noise-robust updates")
        elif self.optimizer_type == "prodigy":
            optimizer_kwargs["d_coef"] = float(config.get("prodigy_d_coef", 1.0))
            logger.info("Using Prodigy optimizer with auto-tuning LR")

        self.optimizer = get_optimizer(
            name=self.optimizer_type,
            params=self.model.parameters(),
            lr=config["learning_rate"],
            weight_decay=config["weight_decay"],
            **optimizer_kwargs,
        )

        # LR scheduler with optional warmup
        epochs = int(config.get("train_epoch", 200))
        warmup_epochs = int(config.get("warmup_epochs", 0))

        if warmup_epochs > 0:
            warmup_scheduler = LinearLR(
                self.optimizer,
                start_factor=0.1,
                end_factor=1.0,
                total_iters=warmup_epochs,
            )
            cosine_scheduler = CosineAnnealingLR(
                self.optimizer,
                T_max=epochs - warmup_epochs,
                eta_min=1e-6,
            )
            self.scheduler = SequentialLR(
                self.optimizer,
                schedulers=[warmup_scheduler, cosine_scheduler],
                milestones=[warmup_epochs],
            )
            logger.info(
                "LR scheduler: %d epochs warmup + cosine decay (total %d epochs)",
                warmup_epochs,
                epochs,
            )
        else:
            self.scheduler = CosineAnnealingLR(
                self.optimizer,
                T_max=epochs,
                eta_min=1e-6,
            )

        # Loss function
        self.loss_type = config.get("loss_type", "bce").lower()
        self.tau = float(config.get("soft_target_temperature", 0.5))

        # Quality-only mode uses pointwise BCE (optionally with Brier for calibration)
        if self.quality_only_training:
            quality_loss_type = config.get("quality_loss_type", "bce").lower()
            quality_label_smoothing = float(config.get("quality_label_smoothing", 0.0))

            if quality_loss_type == "brier":
                self.quality_criterion = BrierLoss(label_smoothing=quality_label_smoothing)
                logger.info(
                    "Quality-only training: using Brier loss (label_smoothing=%.2f)",
                    quality_label_smoothing,
                )
            elif quality_loss_type == "focal":
                focal_alpha = float(config.get("quality_focal_alpha", 0.25))
                focal_gamma = float(config.get("quality_focal_gamma", 2.0))
                self.quality_criterion = FocalLoss(alpha=focal_alpha, gamma=focal_gamma)
                logger.info(
                    "Quality-only training: using Focal loss (alpha=%.2f, gamma=%.2f)",
                    focal_alpha,
                    focal_gamma,
                )
            elif quality_loss_type == "label_smoothing_bce":
                self.quality_criterion = LabelSmoothingBCELoss(
                    smoothing=quality_label_smoothing
                )
                logger.info(
                    "Quality-only training: using label-smoothed BCE (smoothing=%.2f)",
                    quality_label_smoothing,
                )
            else:
                self.quality_criterion = nn.BCELoss()
                logger.info("Quality-only training: using BCE loss")

        if self.loss_type == "listwise":
            # Use reduction="none" to enable per-query importance weighting
            # Loss is computed per-query and weighted before averaging
            self.criterion = nn.KLDivLoss(reduction="none")
        else:
            self.criterion = nn.BCELoss()

        # Calibration settings
        self.calibration_weight = float(config.get("calibration_weight", 0.0))
        self.learned_temperature = nn.Parameter(
            torch.tensor(1.0), requires_grad=False
        )  # For post-hoc scaling

        # Hybrid loss: BCE (calibration) + Pairwise ranking (discrimination)
        # BCE learns P(success | query, LLM) - distinguishes trivial/zero-oracle
        # Pairwise learns ranking among LLMs - sharp discrimination
        self.use_hybrid_loss = config.get("use_hybrid_loss", False)
        self.bce_weight = float(config.get("bce_weight", 0.3))
        self.pairwise_weight = float(config.get("pairwise_weight", 0.3))
        self.pairwise_margin = float(config.get("pairwise_margin", 0.5))
        # Pairwise loss variant: "margin" (hinge), "logistic" (Bradley-Terry),
        # "lambda_rank" (NDCG-weighted logistic)
        self.pairwise_loss_type = config.get("pairwise_loss_type", "margin").lower()
        self.pairwise_temperature = float(config.get("pairwise_temperature", 1.0))
        if self.use_hybrid_loss:
            logger.info(
                "Hybrid loss enabled: BCE=%.2f, Pairwise(%s)=%.2f, margin=%.2f, temp=%.2f",
                self.bce_weight,
                self.pairwise_loss_type,
                self.pairwise_weight,
                self.pairwise_margin,
                self.pairwise_temperature,
            )

        # Training settings
        self.grad_clip = float(config.get("grad_clip_norm", 0.0))
        self.early_stop_patience = int(config.get("early_stop_patience", 0))

        # Model EMA for better generalization
        self.use_ema = config.get("use_ema", False)
        self.ema_decay = float(config.get("ema_decay", 0.999))
        self.ema: ModelEMA | None = None
        if self.use_ema:
            self.ema = ModelEMA(self.model, decay=self.ema_decay, device=self.device)
            logger.info("Model EMA enabled with decay=%.4f", self.ema_decay)

        # Adaptive multi-task weighting: gradually increase difficulty loss weight
        self.adaptive_difficulty_weight = config.get("adaptive_difficulty_weight", False)
        if self.adaptive_difficulty_weight and self.use_difficulty_head:
            self.difficulty_weight_start = float(
                config.get("difficulty_weight_start", self.difficulty_loss_weight)
            )
            self.difficulty_weight_end = float(
                config.get("difficulty_weight_end", self.difficulty_loss_weight * 5)
            )
            self.difficulty_weight_schedule = config.get(
                "difficulty_weight_schedule", "linear"
            ).lower()
            logger.info(
                "Adaptive difficulty weighting: %.2f -> %.2f (%s schedule)",
                self.difficulty_weight_start,
                self.difficulty_weight_end,
                self.difficulty_weight_schedule,
            )

        # Preference sampling
        self.pref_mode = config.get("preference_mode", "grid")
        self.pref_grid = config.get(
            "preference_grid3",
            [[1.0, 0.0, 0.0], [0.6, 0.2, 0.2], [0.4, 0.2, 0.4]],
        )

        # Query-adaptive weighting: adjust reward weights based on per-query perf variance
        self.adaptive_weights = config.get("adaptive_weights", False)
        self.adaptive_easy_threshold = config.get("adaptive_easy_threshold", 0.1)
        self.adaptive_hard_threshold = config.get("adaptive_hard_threshold", 0.3)
        if self.adaptive_weights:
            logger.info(
                "Adaptive weights enabled: easy_t=%.2f, hard_t=%.2f",
                self.adaptive_easy_threshold,
                self.adaptive_hard_threshold,
            )

        # Inference preferences (fixed)
        inference_weights = config.get(
            "inference_weights",
            {"perf": 0.5, "cost": 0.5, "latency": 0.0},
        )
        self.infer_wp = float(inference_weights.get("perf", 0.5))
        self.infer_wc = float(inference_weights.get("cost", 0.5))
        self.infer_wl = float(inference_weights.get("latency", 0.0))

        # Adaptive inference: use difficulty predictions to adjust weights at inference
        self.use_adaptive_inference = config.get("use_adaptive_inference", False)
        if self.use_adaptive_inference and not self.use_difficulty_head:
            logger.warning(
                "use_adaptive_inference requires use_difficulty_head=True; disabling"
            )
            self.use_adaptive_inference = False
        if self.use_adaptive_inference:
            logger.info("Adaptive inference enabled: difficulty predictions will "
                        "influence routing weights")

        # Masks (set during training)
        self.train_mask: torch.Tensor | None = None
        self.valid_mask: torch.Tensor | None = None
        self.test_mask: torch.Tensor | None = None
        self.save_path: str = ""

    def _checkpoint_path(self) -> Path:
        """Return the checkpoint file path derived from save_path."""
        if not self.save_path:
            return Path("checkpoint.pt")
        base = Path(self.save_path)
        return base.parent / f"{base.stem}_checkpoint.pt"

    def _get_difficulty_weight(self, epoch: int, total_epochs: int) -> float:
        """Compute adaptive difficulty loss weight for current epoch.

        Gradually increases weight from difficulty_weight_start to difficulty_weight_end
        over training. This helps the model learn routing first (when difficulty head
        is under-trained) and focus on difficulty later (when routing is stable).

        Schedules:
        - linear: weight = start + (end - start) * (epoch / total_epochs)
        - cosine: smooth S-curve using cosine annealing
        - exp: exponential growth (fast start, slow end)

        Args:
            epoch: Current epoch (0-indexed).
            total_epochs: Total number of training epochs.

        Returns:
            Difficulty loss weight for this epoch.
        """
        if not self.adaptive_difficulty_weight or not self.use_difficulty_head:
            return self.difficulty_loss_weight

        progress = epoch / max(1, total_epochs - 1)  # 0 to 1
        start = self.difficulty_weight_start
        end = self.difficulty_weight_end

        if self.difficulty_weight_schedule == "cosine":
            # Cosine annealing: smooth S-curve
            factor = 0.5 * (1 - math.cos(math.pi * progress))
        elif self.difficulty_weight_schedule == "exp":
            # Exponential: fast start, slow end
            factor = progress ** 2
        else:  # linear (default)
            factor = progress

        weight = start + (end - start) * factor
        return float(weight)

    def _sam_backward_step(self, loss: torch.Tensor) -> None:
        """Perform SAM's two-step backward pass (approximation).

        SAM requires two forward-backward passes:
        1. First pass: Compute gradient, move to adversarial point
        2. Second pass: Compute gradient at adversarial point, update weights

        This implementation uses a common approximation where we:
        1. Retain the computation graph for both backward passes
        2. Use the same graph for both gradients (slight approximation)

        The approximation is effective because:
        - The perturbation is small (rho=0.05 default)
        - The second gradient at perturbed params still captures curvature
        - This approach is used in many SAM implementations for efficiency

        For the routing problem, SAM helps find flat minima that generalize
        better across easy and hard queries.

        Args:
            loss: The loss tensor from the first forward pass.
        """
        # First backward pass - retain graph for second pass
        self.optimizer.zero_grad()
        loss.backward(retain_graph=True)
        if self.grad_clip > 0.0:
            clip_grad_norm_(self.model.parameters(), self.grad_clip)

        # First step: move to adversarial point (stores original params)
        self.optimizer.first_step(zero_grad=True)

        # Second backward pass at adversarial point
        loss.backward()
        if self.grad_clip > 0.0:
            clip_grad_norm_(self.model.parameters(), self.grad_clip)

        # Second step: restore params and apply adversarial gradient update
        self.optimizer.second_step(zero_grad=True)

    def _architecture_signature(self) -> dict:
        """Return a dict describing the model architecture for compatibility checks."""
        # Extract layer shapes that define the architecture
        state_dict = self.model.state_dict()
        return {
            key: list(tensor.shape)
            for key, tensor in state_dict.items()
            if "weight" in key or "bias" in key
        }

    def _save_checkpoint(
        self,
        epoch: int,
        best_score: float,
        no_improve: int,
    ) -> None:
        """Save training checkpoint for crash recovery."""
        checkpoint = {
            "epoch": epoch,
            "model_state_dict": self.model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "scheduler_state_dict": self.scheduler.state_dict(),
            "best_score": best_score,
            "no_improve": no_improve,
            "architecture": self._architecture_signature(),
        }
        # Save EMA state if enabled
        if self.ema is not None:
            checkpoint["ema_state_dict"] = self.ema.state_dict()
        ckpt_path = self._checkpoint_path()
        torch.save(checkpoint, ckpt_path)
        logger.debug("Checkpoint saved at epoch %d to %s", epoch + 1, ckpt_path)

    def _load_checkpoint(self) -> dict | None:
        """Load checkpoint if it exists and is compatible. Returns None otherwise.

        If the checkpoint exists but has an incompatible architecture (e.g., different
        hidden_dim), logs a warning and returns None to start fresh training. The
        incompatible checkpoint is preserved for potential future use.
        """
        ckpt_path = self._checkpoint_path()
        if not ckpt_path.exists():
            return None

        device = next(self.model.parameters()).device
        checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)

        # Validate architecture compatibility
        saved_arch = checkpoint.get("architecture")
        current_arch = self._architecture_signature()

        if saved_arch is not None and saved_arch != current_arch:
            # Find specific mismatches for helpful error message
            mismatches = []
            for key in current_arch:
                if key in saved_arch and saved_arch[key] != current_arch[key]:
                    mismatches.append(
                        f"  {key}: checkpoint {saved_arch[key]} vs current {current_arch[key]}"
                    )
            mismatch_details = "\n".join(mismatches[:5])  # Show first 5
            if len(mismatches) > 5:
                mismatch_details += f"\n  ... and {len(mismatches) - 5} more"

            logger.warning(
                "Checkpoint architecture mismatch - starting fresh training.\n"
                "Checkpoint preserved at: %s\n"
                "Mismatched layers:\n%s",
                ckpt_path,
                mismatch_details,
            )
            return None

        # For legacy checkpoints without architecture info, try to load and catch errors
        if saved_arch is None:
            try:
                self.model.load_state_dict(checkpoint["model_state_dict"])
            except RuntimeError as e:
                if "size mismatch" in str(e):
                    logger.warning(
                        "Legacy checkpoint has incompatible architecture - starting fresh.\n"
                        "Checkpoint preserved at: %s\n"
                        "Error: %s",
                        ckpt_path,
                        str(e).split("\n")[0],
                    )
                    return None
                raise
        else:
            self.model.load_state_dict(checkpoint["model_state_dict"])

        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])

        # Load EMA state if available and EMA is enabled
        if self.ema is not None and "ema_state_dict" in checkpoint:
            self.ema.load_state_dict(checkpoint["ema_state_dict"])

        logger.info(
            "Resumed from checkpoint: epoch %d, best_score=%.4f",
            checkpoint["epoch"] + 1,
            checkpoint["best_score"],
        )
        return checkpoint

    def _clear_checkpoint(self) -> None:
        """Remove checkpoint file after successful training completion."""
        ckpt_path = self._checkpoint_path()
        if ckpt_path.exists():
            ckpt_path.unlink()
            logger.debug("Checkpoint cleared: %s", ckpt_path)

    def _sample_pref(self) -> tuple[float, float, float]:
        """Sample preference weights for training."""
        mode = (self.pref_mode or "grid").lower()

        if mode == "grid" and self.pref_grid:
            weights = random.choice(self.pref_grid)
            return float(weights[0]), float(weights[1]), float(weights[2])

        if mode == "uniform":
            weights = np.random.random(3)
            weights = weights / (weights.sum() or 1.0)
            return float(weights[0]), float(weights[1]), float(weights[2])

        return self.infer_wp, self.infer_wc, self.infer_wl

    def _build_query_mask(
        self,
        selected_queries: torch.Tensor,
        total_edges: int,
    ) -> torch.Tensor:
        """Build a boolean mask for edges belonging to selected queries.

        Vectorized implementation that avoids Python loops over queries.
        Each query owns `llm_num` consecutive edges.

        Args:
            selected_queries: Query indices to include in mask.
            total_edges: Total number of edges in the graph.

        Returns:
            Boolean mask of shape (total_edges,) with True for selected query edges.
        """
        device = selected_queries.device
        mask = torch.zeros(total_edges, dtype=torch.bool, device=device)

        if selected_queries.numel() == 0:
            return mask

        # Compute all edge indices for selected queries in one vectorized operation
        # starts: (num_selected,) containing first edge index for each query
        starts = selected_queries * self.llm_num
        # offsets: (llm_num,) containing [0, 1, 2, ..., llm_num-1]
        offsets = torch.arange(self.llm_num, device=device)
        # indices: (num_selected, llm_num) -> flatten to (num_selected * llm_num,)
        indices = (starts.unsqueeze(1) + offsets).flatten()
        mask[indices] = True

        return mask

    @staticmethod
    def _edge_features_with_pref(
        combined_edge: torch.Tensor,
        wp: float,
        wc: float,
        wl: float,
    ) -> torch.Tensor:
        """Prepend preference weights to edge features."""
        weights = torch.tensor(
            [wp, wc, wl],
            dtype=combined_edge.dtype,
            device=combined_edge.device,
        )
        expanded = weights.unsqueeze(0).expand(combined_edge.size(0), -1)
        return torch.cat((expanded, combined_edge), dim=1)

    @staticmethod
    def _rewards_from_pref(
        combined_edge: torch.Tensor,
        wp: float,
        wc: float,
        wl: float,
    ) -> torch.Tensor:
        """Compute reward = wp*perf - wc*cost - wl*latency."""
        cost = combined_edge[:, 0]
        perf = combined_edge[:, 1]
        latency = combined_edge[:, 2] if combined_edge.size(1) > 2 else torch.zeros_like(cost)
        return wp * perf - wc * cost - wl * latency

    def _adaptive_rewards_from_pref(
        self,
        combined_edge: torch.Tensor,
        wp: float,
        wc: float,
        wl: float,
    ) -> torch.Tensor:
        """Compute query-adaptive rewards based on per-query performance variance.

        For easy queries (all LLMs perform similarly): emphasize cost/latency
        For hard queries (LLMs differ): emphasize performance

        Args:
            combined_edge: (num_edges, 3) tensor with [cost, perf, latency]
            wp, wc, wl: Base preference weights

        Returns:
            Per-edge rewards with query-adaptive weighting
        """
        cost = combined_edge[:, 0]
        perf = combined_edge[:, 1]
        latency = combined_edge[:, 2] if combined_edge.size(1) > 2 else torch.zeros_like(cost)

        # Reshape to per-query: (num_queries, num_llms)
        num_edges = combined_edge.size(0)
        num_queries = num_edges // self.llm_num
        perf_matrix = perf.view(num_queries, self.llm_num)

        # Compute per-query perf range (max - min)
        perf_range = perf_matrix.max(dim=1).values - perf_matrix.min(dim=1).values

        # Compute adaptive weights per query
        # Easy queries (range < easy_threshold): cost/latency dominate
        # Hard queries (range > hard_threshold): perf dominates
        # In between: linear interpolation
        easy_t = self.adaptive_easy_threshold
        hard_t = self.adaptive_hard_threshold

        # Interpolation factor: 0 = easy (cost-focused), 1 = hard (perf-focused)
        alpha = ((perf_range - easy_t) / (hard_t - easy_t + 1e-8)).clamp(0, 1)

        # Easy weights: minimize cost/latency (perf=0, cost=0.6, latency=0.4)
        # Hard weights: maximize perf (perf=0.95, cost=0.025, latency=0.025)
        wp_adaptive = alpha * 0.95 + (1 - alpha) * 0.0
        wc_adaptive = alpha * 0.025 + (1 - alpha) * 0.6
        wl_adaptive = alpha * 0.025 + (1 - alpha) * 0.4

        # Expand to per-edge: (num_queries,) -> (num_queries, num_llms) -> (num_edges,)
        wp_per_edge = wp_adaptive.unsqueeze(1).expand(-1, self.llm_num).reshape(-1)
        wc_per_edge = wc_adaptive.unsqueeze(1).expand(-1, self.llm_num).reshape(-1)
        wl_per_edge = wl_adaptive.unsqueeze(1).expand(-1, self.llm_num).reshape(-1)

        return wp_per_edge * perf - wc_per_edge * cost - wl_per_edge * latency

    def _init_dynamic_bias(self, device: torch.device | str) -> None:
        """Initialize dynamic bias tensors for loss-free load balancing.

        Based on DeepSeek's auxiliary-loss-free balancing: instead of adding
        gradient-interfering auxiliary losses, we apply a learned bias to
        routing scores based on selection frequency. LLMs selected too often
        get negative bias, forcing the model to explore other options.
        """
        if self._llm_selection_bias is None:
            self._llm_selection_bias = torch.zeros(
                self.llm_num, device=device, dtype=torch.float32
            )
            self._llm_selection_counts = torch.zeros(
                self.llm_num, device=device, dtype=torch.float32
            )
            self._total_selections = 0

    def _apply_dynamic_bias(self, logits_grouped: torch.Tensor) -> torch.Tensor:
        """Apply dynamic bias to logits before softmax/loss computation.

        Args:
            logits_grouped: Routing logits (batch_size, num_llms)

        Returns:
            Bias-adjusted logits for more balanced selection
        """
        if not self.use_dynamic_bias or self._llm_selection_bias is None:
            return logits_grouped
        return logits_grouped + self._llm_selection_bias.unsqueeze(0)

    def _update_dynamic_bias(self, logits_grouped: torch.Tensor) -> None:
        """Update dynamic bias based on batch selection frequencies.

        After each batch, track which LLMs were selected and adjust bias
        to discourage over-selected LLMs and encourage under-selected ones.

        Args:
            logits_grouped: Routing logits (batch_size, num_llms)
        """
        if not self.use_dynamic_bias or self._llm_selection_bias is None:
            return

        batch_size = logits_grouped.size(0)
        with torch.no_grad():
            # Get selections from adjusted logits
            adjusted = logits_grouped + self._llm_selection_bias.unsqueeze(0)
            selections = adjusted.argmax(dim=1)

            # Count selections per LLM in this batch
            batch_counts = torch.zeros(
                self.llm_num, device=logits_grouped.device, dtype=torch.float32
            )
            for i in range(self.llm_num):
                batch_counts[i] = (selections == i).sum().float()

            # EMA update of selection counts
            self._llm_selection_counts = (
                self.dynamic_bias_ema * self._llm_selection_counts
                + (1 - self.dynamic_bias_ema) * batch_counts
            )
            self._total_selections += batch_size

            # Compute selection frequency deviation from uniform
            # target: each LLM should be selected 1/num_llms of the time
            target_freq = 1.0 / self.llm_num
            current_freq = self._llm_selection_counts / (
                self._llm_selection_counts.sum() + 1e-8
            )
            freq_deviation = current_freq - target_freq

            # Update bias: negative for over-selected, positive for under-selected
            self._llm_selection_bias = (
                self._llm_selection_bias - self.dynamic_bias_lr * freq_deviation
            )

    def _weights_from_difficulty_pred(
        self,
        difficulty_pred: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Compute per-query preference weights from predicted difficulty scores.

        This enables adaptive inference: the difficulty head predicts how "hard"
        each query is, and we use that to interpolate between cost-focused weights
        (for easy queries) and performance-focused weights (for hard queries).

        The difficulty head outputs:
        - Classification mode: P(easy) in [0,1] where 1=easy, 0=hard
        - Regression mode: predicted perf_range in [0,1] where 0=easy, 1=hard

        Args:
            difficulty_pred: Predicted difficulty scores (num_queries,) in [0, 1].
                             Interpretation depends on difficulty_head_mode.

        Returns:
            Tuple of (wp, wc, wl) tensors, each shape (num_queries,).
            - wp: Performance weight per query
            - wc: Cost weight per query
            - wl: Latency weight per query
        """
        # Convert difficulty_pred to "hardness" alpha where 0 = easy, 1 = hard
        # Classification mode: high pred = easy, so invert
        # Regression mode: high pred = hard, use directly
        if self.difficulty_head_mode == "classification":
            # P(easy) -> hardness: invert so hard queries have alpha ≈ 1
            alpha = 1.0 - difficulty_pred
        else:
            # Regression predicts perf_range directly: high = hard
            alpha = difficulty_pred

        # Clamp to [0, 1] for numerical safety
        alpha = alpha.clamp(0.0, 1.0)

        # Weight profiles (matching _adaptive_rewards_from_pref)
        # Easy: prioritize cost/latency (wp=0.0, wc=0.6, wl=0.4)
        # Hard: prioritize performance (wp=0.95, wc=0.025, wl=0.025)
        wp = alpha * 0.95 + (1.0 - alpha) * 0.0
        wc = alpha * 0.025 + (1.0 - alpha) * 0.6
        wl = alpha * 0.025 + (1.0 - alpha) * 0.4

        return wp, wc, wl

    def _adaptive_rewards_from_difficulty(
        self,
        combined_edge: torch.Tensor,
        difficulty_pred: torch.Tensor,
    ) -> torch.Tensor:
        """Compute rewards using predicted difficulty for weight interpolation.

        This is the inference-time counterpart to _adaptive_rewards_from_pref.
        Instead of using ground truth perf_range, we use the difficulty head's
        predictions to determine per-query weights.

        Args:
            combined_edge: (num_edges, 3+) tensor with [cost, perf, latency, ...]
            difficulty_pred: Predicted difficulty (num_queries,) from difficulty head.

        Returns:
            Per-edge rewards with difficulty-adaptive weighting.
        """
        cost = combined_edge[:, 0]
        perf = combined_edge[:, 1]
        latency = (
            combined_edge[:, 2]
            if combined_edge.size(1) > 2
            else torch.zeros_like(cost)
        )

        # Get per-query weights from difficulty predictions
        wp_per_query, wc_per_query, wl_per_query = self._weights_from_difficulty_pred(
            difficulty_pred
        )

        # Expand to per-edge: (num_queries,) -> (num_queries, num_llms) -> (num_edges,)
        wp_per_edge = wp_per_query.unsqueeze(1).expand(-1, self.llm_num).reshape(-1)
        wc_per_edge = wc_per_query.unsqueeze(1).expand(-1, self.llm_num).reshape(-1)
        wl_per_edge = wl_per_query.unsqueeze(1).expand(-1, self.llm_num).reshape(-1)

        return wp_per_edge * perf - wc_per_edge * cost - wl_per_edge * latency

    def _compute_difficulty_targets(
        self,
        combined_edge: torch.Tensor,
        mask: torch.Tensor,
        selected_queries: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Compute ground truth difficulty for each query.

        Supports two modes:
        1. Regression: Returns perf_range (continuous in [0, 1])
        2. Classification: Returns binary labels (1=easy, 0=hard)

        Difficulty is based on max(perf) - min(perf) across LLMs:
        - High perf_range = hard query (LLM choice matters)
        - Low perf_range = easy query (any LLM works)

        For classification, uses adaptive threshold based on distribution:
        - Threshold ≈ median(perf_range) to balance classes
        - Can be configured via difficulty_classification_threshold

        Args:
            combined_edge: (num_edges, 3) tensor with [cost, perf, latency]
            mask: Boolean mask for edges being predicted
            selected_queries: Optional tensor of query indices (for ensuring shape match)

        Returns:
            Difficulty targets (num_queries,) in:
            - Regression mode: [0, 1] continuous
            - Classification mode: {0, 1} binary (1=easy, 0=hard)
        """
        perf = combined_edge[:, 1]

        if selected_queries is not None:
            # Use selected_queries to compute targets for exactly those queries
            # This ensures shape matches with difficulty_pred[selected_queries]
            perf_ranges = []
            for q_idx in selected_queries.tolist():
                q_start = q_idx * self.llm_num
                q_end = q_start + self.llm_num
                if q_end <= perf.size(0):
                    query_perf = perf[q_start:q_end]
                    perf_range = (query_perf.max() - query_perf.min()).clamp(0, 1)
                    perf_ranges.append(perf_range)
            perf_range = torch.stack(perf_ranges)
        else:
            # Fallback: use mask (old behavior)
            perf_masked = perf[mask].view(-1, self.llm_num)
            perf_range = perf_masked.max(dim=1).values - perf_masked.min(dim=1).values
            perf_range = perf_range.clamp(0, 1)

        # For classification mode, convert to binary labels
        if self.difficulty_head_mode == "classification":
            # Threshold: low perf_range = easy (1), high = hard (0)
            threshold = self.config.get("difficulty_classification_threshold", 0.15)
            # Invert: easy queries have LOW perf_range, so we want perf_range < threshold
            binary_labels = (perf_range < threshold).float()
            return binary_labels

        # For regression mode, return continuous perf_range
        return perf_range

    def _compute_hard_negative_weights(
        self,
        combined_edge: torch.Tensor,
        query_indices: torch.Tensor,
    ) -> torch.Tensor:
        """Compute sampling weights for hard negative mining.

        Focuses training on queries where model selection matters most
        (high perf_range = model choice has big impact on performance).

        Strategy:
        1. Compute perf_range for each query
        2. Identify "hard" queries (top 30% by perf_range)
        3. Give hard queries 2-3x sampling weight

        This ensures the model learns to make good decisions on high-value
        queries while still seeing easy queries for calibration.

        Args:
            combined_edge: (num_edges, 3) tensor with [cost, perf, latency]
            query_indices: Tensor of query indices to compute weights for

        Returns:
            Sampling weights (num_queries,) summing to 1.0
        """
        # Get performance values
        perf = combined_edge[:, 1]

        # Reshape to per-query matrix (num_queries, num_llms)
        perf_matrix = perf.view(-1, self.llm_num)

        # Compute per-query perf range (difficulty)
        perf_range = perf_matrix.max(dim=1).values - perf_matrix.min(dim=1).values

        # Get perf_range for specified queries
        perf_range_subset = perf_range[query_indices]

        # Define "hard" queries as top 30% by perf_range
        hard_negative_ratio = self.config.get("hard_negative_ratio", 0.3)
        threshold = torch.quantile(perf_range_subset, 1 - hard_negative_ratio)

        # Create sampling weights
        weights = torch.ones_like(perf_range_subset)

        # Hard queries get 2-3x weight
        hard_mask = perf_range_subset > threshold
        hard_boost = self.config.get("hard_negative_boost", 2.5)
        weights[hard_mask] = hard_boost

        # Normalize to sum to 1
        weights = weights / weights.sum()

        return weights

    def _compute_needle_query_weights(
        self,
        targets: torch.Tensor,
        num_queries: int,
    ) -> torch.Tensor:
        """Compute per-edge weights based on number of correct LLMs per query.

        "Needle queries" are those where only 1-2 LLMs succeed - these are
        the hardest to route and benefit most from focused learning.

        Weighting strategy:
        - 1 correct LLM (needle): weight = num_llms (7x for 7 LLMs)
        - 2 correct LLMs: weight = num_llms / 2 (3.5x)
        - All correct LLMs: weight = 1 (baseline)

        This makes needle queries contribute proportionally more to the loss,
        forcing the model to learn the distinguishing features.

        Args:
            targets: Binary targets (num_edges,) where 1 = correct LLM.
            num_queries: Hint for number of queries (overridden by actual tensor size).

        Returns:
            Per-edge weights (num_edges,) for weighted loss computation.
        """
        # Derive actual query count from tensor size (batch boundary may truncate)
        num_edges = targets.numel()
        actual_queries = num_edges // self.llm_num
        if actual_queries * self.llm_num != num_edges:
            # Edge count not divisible by llm_num — fall back to uniform weights
            return torch.ones_like(targets)

        # Reshape to (num_queries, num_llms)
        targets_matrix = targets.view(actual_queries, self.llm_num)

        # Count correct LLMs per query
        num_correct = targets_matrix.sum(dim=1).clamp(min=1)  # (num_queries,)

        # Weight = num_llms / num_correct (inverse weighting)
        # Needle queries (1 correct) get 7x weight vs all-correct queries
        query_weights = self.llm_num / num_correct  # (num_queries,)

        # Expand to edge-level weights (repeat for each LLM)
        edge_weights = query_weights.unsqueeze(1).expand(-1, self.llm_num).reshape(-1)

        return edge_weights

    def _compute_query_importance(
        self,
        combined_edge: torch.Tensor,
        mask: torch.Tensor,
    ) -> torch.Tensor:
        """Compute per-query importance weights for loss weighting.

        Different strategies:
        - "uniform": All queries weighted equally (baseline)
        - "needle": Inverse frequency (7x for 1-correct, 1x for 7-correct)
        - "variance": Weight by performance variance across LLMs

        Zero-oracle queries (where no LLM succeeds) are weighted by
        zero_oracle_importance - set > 0 to learn cost preferences for these.

        Args:
            combined_edge: Edge features (num_edges, features) with perf at [:, 1].
            mask: Boolean mask selecting edges for current batch.

        Returns:
            Per-query importance weights (num_queries,), normalized to mean ~1.
        """
        perf = combined_edge[mask, 1].view(-1, self.llm_num)
        num_correct = perf.sum(dim=1)

        if self.importance_weighting == "needle":
            # Inverse frequency: upweight queries where routing matters most
            # - Needle (1 correct): weight = num_llms (max signal)
            # - Trivial (all correct): weight = 1 (min signal)
            # - Zero-oracle (0 correct): configurable importance
            importance = torch.where(
                num_correct > 0,
                self.llm_num / num_correct.clamp(min=1),
                torch.full_like(num_correct, self.zero_oracle_importance),
            )
        elif self.importance_weighting == "variance":
            # Weight by performance variance (queries where LLMs differ)
            # Higher variance = more discriminative signal for routing
            importance = perf.var(dim=1) * self.llm_num + 1.0
        else:  # uniform
            importance = torch.ones(num_correct.size(0), device=perf.device)

        # Normalize to mean=1 for interpretable loss scale
        mean_importance = importance.mean().clamp(min=1e-8)
        return importance / mean_importance

    def _labels_from_rewards(self, rewards: torch.Tensor) -> torch.Tensor:
        """Convert rewards to one-hot labels (best LLM = 1)."""
        rewards_matrix = rewards.view(-1, self.llm_num)
        labels = torch.zeros_like(rewards_matrix)
        if rewards_matrix.size(0) == 0:
            return labels.view(-1)

        best_idx = torch.argmax(rewards_matrix, dim=1)
        row_idx = torch.arange(rewards_matrix.size(0), device=rewards_matrix.device)
        labels[row_idx, best_idx] = 1.0
        return labels.view(-1)

    @staticmethod
    def _ece(conf: np.ndarray, correct: np.ndarray, n_bins: int = 10) -> float:
        """Compute Expected Calibration Error."""
        if conf.size == 0:
            return 0.0

        bins = np.linspace(0.0, 1.0, n_bins + 1)
        total = float(conf.size)
        ece = 0.0

        for idx in range(n_bins):
            if idx < n_bins - 1:
                mask = (conf >= bins[idx]) & (conf < bins[idx + 1])
            else:
                mask = (conf >= bins[idx]) & (conf <= bins[idx + 1])

            if not mask.any():
                continue

            acc = float(correct[mask].mean())
            conf_mean = float(conf[mask].mean())
            ece += (mask.sum() / total) * abs(acc - conf_mean)

        return float(ece)

    @staticmethod
    def _brier_score(
        probs: torch.Tensor,
        rewards: torch.Tensor,
    ) -> torch.Tensor:
        """Compute multi-class Brier score for calibration.

        Brier score = mean((probs - one_hot_best)^2), a proper scoring rule
        that penalizes miscalibrated probabilities. Lower is better.

        Args:
            probs: Predicted probabilities [batch, num_llms], should sum to 1.
            rewards: Reward values [batch, num_llms] used to determine best LLM.

        Returns:
            Scalar Brier score loss (differentiable).
        """
        best_idx = rewards.argmax(dim=1)
        one_hot = torch.zeros_like(probs)
        one_hot.scatter_(1, best_idx.unsqueeze(1), 1.0)
        return ((probs - one_hot) ** 2).sum(dim=1).mean()

    def _pairwise_ranking_loss(
        self,
        scores: torch.Tensor,
        successes: torch.Tensor,
        margin: float = 0.5,
        threshold: float = 0.5,
    ) -> torch.Tensor:
        """Compute pairwise margin ranking loss from binary outcomes.

        For each (winner, loser) pair where winner succeeds and loser fails,
        enforces: score(winner) > score(loser) + margin.

        This directly optimizes the ranking constraint using all pairwise
        comparisons implied by binary success/failure labels.

        Args:
            scores: Model routing scores [batch, num_llms].
            successes: Binary success indicators [batch, num_llms], 1 = success.
            margin: Required score gap between winners and losers. Default 0.5.
            threshold: Threshold for binarizing continuous targets. Default 0.5.

        Returns:
            Scalar pairwise ranking loss. Returns 0 if no valid pairs.
        """
        device = scores.device

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

        # All pairwise score differences: diff[b,i,j] = score[b,i] - score[b,j]
        score_diff = scores.unsqueeze(2) - scores.unsqueeze(1)  # (B, L, L)

        # Valid pair mask: pref[b,i,j] = 1 if LLM i succeeds and LLM j fails
        valid_pair_mask = winners.unsqueeze(2) * losers.unsqueeze(1)  # (B, L, L)

        # Hinge loss: max(0, margin - (score_winner - score_loser))
        pair_loss = F.relu(margin - score_diff)  # (B, 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
        return torch.tensor(0.0, device=device, requires_grad=True)

    def _pointwise_bce_loss(
        self,
        scores: torch.Tensor,
        successes: torch.Tensor,
    ) -> torch.Tensor:
        """Compute pointwise BCE loss for calibrated success probabilities.

        Learns P(success | query, LLM) for each LLM independently.
        This enables distinguishing:
        - Trivial queries: all high scores -> pick cheapest
        - Zero-oracle queries: all low scores -> escalate
        - Discriminative queries: mixed scores -> pick argmax

        Args:
            scores: Model routing scores (logits) [batch, num_llms].
            successes: Binary success indicators [batch, num_llms].

        Returns:
            Scalar BCE loss.
        """
        return F.binary_cross_entropy_with_logits(
            scores, successes, reduction="mean"
        )

    def _to_probs(self, scores: torch.Tensor) -> torch.Tensor:
        """Convert logits to probabilities."""
        if scores.numel() == 0:
            return scores

        if self.loss_type == "listwise":
            return torch.softmax(scores, dim=1)

        probs = torch.sigmoid(scores)
        denom = probs.sum(dim=1, keepdim=True) + 1e-8
        return probs / denom

    def _compute_eval_metrics(
        self,
        pred_scores: torch.Tensor,
        rewards: torch.Tensor,
    ) -> dict[str, float]:
        """Compute evaluation metrics for routing predictions."""
        empty_metrics = {
            "top1_acc": 0.0,
            "regret": 0.0,
            "frac_of_oracle": 0.0,
            "uplift_norm": 0.0,
            "ndcg@1": 0.0,
            "ndcg@3": 0.0,
            "ndcg@All": 0.0,
            "ece": 0.0,
            "avg_chosen_reward": 0.0,
            "avg_oracle_reward": 0.0,
        }

        if pred_scores.numel() == 0:
            return empty_metrics

        probs = self._to_probs(pred_scores)
        pred_rank = torch.argsort(pred_scores, dim=1, descending=True)
        true_rank = torch.argsort(rewards, dim=1, descending=True)

        max_reward = rewards.max(dim=1).values
        min_reward = rewards.min(dim=1).values
        span = max_reward - min_reward
        shifted_rewards = rewards - min_reward.unsqueeze(1)

        row_indices = torch.arange(rewards.size(0), device=rewards.device)
        chosen_reward = rewards[row_indices, pred_rank[:, 0]]
        chosen_shifted = shifted_rewards[row_indices, pred_rank[:, 0]]
        oracle_shifted = shifted_rewards[row_indices, true_rank[:, 0]]

        top1_acc = (pred_rank[:, 0] == true_rank[:, 0]).float().mean().item()
        regret = (max_reward - chosen_reward).mean().item()
        uplift = (chosen_shifted / span.clamp_min(1e-6)).mean().item()

        # frac_of_oracle: Handle ties correctly (all-correct or all-wrong).
        # For ties (span < epsilon), any choice matches oracle -> frac = 1.0.
        # For non-ties, compute chosen/oracle ratio as usual.
        is_tie = span < 1e-6
        per_query_frac = torch.where(
            is_tie,
            torch.ones_like(chosen_shifted),
            chosen_shifted / oracle_shifted.clamp_min(1e-6),
        )
        frac_of_oracle = per_query_frac.mean().item()

        def ndcg_at_k(k: int) -> float:
            k = min(k, rewards.size(1))
            denom = torch.log2(
                torch.arange(2, k + 2, device=rewards.device, dtype=rewards.dtype)
            )
            gains_pred = shifted_rewards[row_indices.unsqueeze(1), pred_rank[:, :k]]
            gains_true = shifted_rewards[row_indices.unsqueeze(1), true_rank[:, :k]]
            dcg = (gains_pred / denom).sum(dim=1)
            idcg = (gains_true / denom).sum(dim=1) + 1e-8
            return (dcg / idcg).mean().item()

        probs_max = probs.max(dim=1).values.detach().cpu().numpy()
        correct = (pred_rank[:, 0] == true_rank[:, 0]).float().cpu().numpy()

        return {
            "top1_acc": top1_acc,
            "regret": regret,
            "frac_of_oracle": frac_of_oracle,
            "uplift_norm": uplift,
            "ndcg@1": ndcg_at_k(1),
            "ndcg@3": ndcg_at_k(3),
            "ndcg@All": ndcg_at_k(rewards.size(1)),
            "ece": self._ece(probs_max, correct),
            "avg_chosen_reward": chosen_reward.mean().item(),
            "avg_oracle_reward": max_reward.mean().item(),
        }

    def _train_epoch_listwise(
        self,
        data: Data,
        train_query_indices: torch.Tensor,
        batch_size: int,
        current_epoch: int = 0,
        total_epochs: int = 1,
    ) -> tuple[float, float]:
        """Run one training epoch with listwise loss.

        Returns:
            Tuple of (total_loss, difficulty_loss). If difficulty head is disabled,
            difficulty_loss will be 0.0.

        Uses standard mini-batch SGD: backward() is called per batch to avoid
        accumulating a massive computation graph across all batches.
        """
        if train_query_indices.numel() == 0:
            return 0.0, 0.0

        device = data.task_id.device
        # Initialize dynamic bias tensors if enabled
        if self.use_dynamic_bias:
            self._init_dynamic_bias(device)

        effective_batch = max(1, batch_size)
        num_batches = max(1, math.ceil(train_query_indices.numel() / effective_batch))

        # Hard negative mining: weight sampling by query difficulty
        use_hard_mining = self.config.get("use_hard_negative_mining", False)
        if use_hard_mining:
            # Compute sampling weights based on perf_range
            sample_weights = self._compute_hard_negative_weights(
                data.combined_edge,
                train_query_indices,
            )
            # Sample with replacement using weights
            num_samples = train_query_indices.numel()
            perm = torch.multinomial(
                sample_weights,
                num_samples=num_samples,
                replacement=True,
            )
        else:
            # Uniform random sampling (original behavior)
            perm = torch.randperm(train_query_indices.numel(), device=device)

        total_loss = 0.0
        total_diff_loss = 0.0
        log_interval = max(1, num_batches // 5)  # Log ~5 times per epoch

        model = self.model.module if hasattr(self.model, "module") else self.model
        best_llm_idx = None
        if hasattr(model, "use_contrastive") and model.use_contrastive:
            quality_scores = data.combined_edge[:, 1]  # effect column
            best_llm_idx = quality_scores.view(-1, self.llm_num).argmax(dim=1)

        for batch_idx in range(num_batches):
            if batch_idx > 0 and batch_idx % log_interval == 0:
                logger.info(
                    "  Batch %d/%d (%.0f%%), loss: %.4f",
                    batch_idx,
                    num_batches,
                    100 * batch_idx / num_batches,
                    total_loss / batch_idx,
                )

            start = batch_idx * effective_batch
            end = min((batch_idx + 1) * effective_batch, train_query_indices.numel())
            selected_queries = train_query_indices[perm[start:end]]

            if selected_queries.numel() == 0:
                continue

            # Build mask for selected queries (vectorized)
            mask = self._build_query_mask(selected_queries, self.train_mask.numel())

            edge_can_see = self.train_mask & ~mask
            wp, wc, wl = self._sample_pref()

            augmented_edge = self._edge_features_with_pref(data.combined_edge, wp, wc, wl)

            # Forward pass - optionally get difficulty and/or uncertainty predictions
            difficulty_pred = None
            log_var = None
            forward_outputs = self.model(
                task_features=data.task_id,
                query_features=data.query_features,
                llm_features=data.llm_features,
                edge_index=data.edge_index,
                edge_mask=mask,
                edge_can_see=edge_can_see,
                edge_weight=augmented_edge,
                return_difficulty=self.use_difficulty_head,
                return_uncertainty=self.use_uncertainty_head,
                best_llm_idx=best_llm_idx,
            )

            # Unpack based on what was requested
            if self.use_difficulty_head and self.use_uncertainty_head:
                logits, difficulty_pred, log_var = forward_outputs
            elif self.use_difficulty_head:
                logits, difficulty_pred = forward_outputs
            elif self.use_uncertainty_head:
                logits, log_var = forward_outputs
            else:
                logits = forward_outputs

            # Get difficulty targets if needed
            if self.use_difficulty_head:
                difficulty_target = self._compute_difficulty_targets(
                    data.combined_edge, mask, selected_queries
                )

            if self.adaptive_weights:
                rewards = self._adaptive_rewards_from_pref(data.combined_edge, wp, wc, wl)
            else:
                rewards = self._rewards_from_pref(data.combined_edge, wp, wc, wl)
            target = rewards[mask].reshape(-1, self.llm_num)
            soft_target = torch.softmax(target / self.tau, dim=1).clamp_min(1e-8)
            logits_grouped = logits.reshape(-1, self.llm_num)

            # Apply dynamic bias for loss-free load balancing (anti-collapse)
            logits_for_loss = self._apply_dynamic_bias(logits_grouped)

            # Compute per-query KL divergence with importance weighting
            # criterion returns (num_queries, num_llms), sum over LLMs for per-query loss
            # Use bias-adjusted logits for loss computation
            per_query_kl = self.criterion(
                torch.log_softmax(logits_for_loss, dim=1), soft_target
            ).sum(dim=1)

            # Apply importance weighting (needle queries get higher weight)
            importance = self._compute_query_importance(data.combined_edge, mask)
            ranking_loss = (importance * per_query_kl).mean()

            # Hybrid loss: combines listwise KL + pointwise BCE + pairwise ranking
            # - KL: distributional matching to soft targets (existing)
            # - BCE: calibrated P(success|query,LLM) for regime detection
            # - Pairwise: sharp ranking among success/failure pairs
            if self.use_hybrid_loss:
                # Get binary success indicators for hybrid losses
                perf = data.combined_edge[mask, 1].view(-1, self.llm_num)

                # BCE loss for calibration (learns absolute success probability)
                bce_loss = self._pointwise_bce_loss(logits_for_loss, perf)

                # Pairwise ranking loss (learns relative ranking)
                # Dispatch to appropriate loss based on pairwise_loss_type
                if self.pairwise_loss_type == "logistic":
                    pairwise_loss = pairwise_logistic_loss(
                        logits_for_loss, perf,
                        temperature=self.pairwise_temperature,
                    )
                elif self.pairwise_loss_type == "lambda_rank":
                    pairwise_loss = lambda_rank_loss(logits_for_loss, perf)
                else:  # Default: margin (hinge) loss
                    pairwise_loss = pairwise_margin_loss(
                        logits_for_loss, perf, margin=self.pairwise_margin
                    )

                # Combine: KL (implicit from ranking_loss weight) + BCE + Pairwise
                # The KL weight is implicitly 1 - bce_weight - pairwise_weight
                kl_weight = 1.0 - self.bce_weight - self.pairwise_weight
                loss = (
                    kl_weight * ranking_loss
                    + self.bce_weight * bce_loss
                    + self.pairwise_weight * pairwise_loss
                )

                # Track individual loss components for gradient diagnosis
                if not hasattr(self, '_loss_components'):
                    self._loss_components = {
                        'kl': [], 'bce': [], 'pairwise': []
                    }
                self._loss_components['kl'].append(ranking_loss.item())
                self._loss_components['bce'].append(bce_loss.item())
                self._loss_components['pairwise'].append(pairwise_loss.item())
            # Add calibration loss (Brier score) if weight > 0
            elif self.calibration_weight > 0:
                probs = torch.softmax(logits_for_loss, dim=1)
                calibration_loss = self._brier_score(probs, target)
                loss = ranking_loss + self.calibration_weight * calibration_loss
            else:
                loss = ranking_loss

            # Add difficulty prediction loss if enabled
            if self.use_difficulty_head:
                # Select predictions for queries in this batch
                difficulty_pred_batch = difficulty_pred[selected_queries]
                difficulty_loss = self.difficulty_criterion(
                    difficulty_pred_batch,
                    difficulty_target,
                )
                total_diff_loss += difficulty_loss.item()
                # Use adaptive weighting if enabled
                current_difficulty_weight = self._get_difficulty_weight(
                    current_epoch, total_epochs
                )
                loss = loss + current_difficulty_weight * difficulty_loss

                # Entropy regularization to prevent collapse (constant predictions)
                # Penalize when predictions cluster near 0 or 1 with low variance
                if self.difficulty_entropy_weight > 0:
                    pred_mean = difficulty_pred_batch.mean()
                    # Binary entropy: -p*log(p) - (1-p)*log(1-p), max at p=0.5
                    eps = 1e-8
                    entropy = -(
                        pred_mean * torch.log(pred_mean + eps)
                        + (1 - pred_mean) * torch.log(1 - pred_mean + eps)
                    )
                    # Max entropy is log(2) ≈ 0.693; penalize low entropy
                    max_entropy = 0.693
                    entropy_loss = max_entropy - entropy
                    loss = loss + self.difficulty_entropy_weight * entropy_loss

            if hasattr(model, "get_auxiliary_losses"):
                aux_losses = model.get_auxiliary_losses()
                for aux_name, aux_loss in aux_losses.items():
                    if aux_loss is not None:
                        loss = loss + aux_loss

            if self.use_llm_balance_loss:
                probs = torch.softmax(logits_for_loss, dim=1)
                avg_selection = probs.mean(dim=0)
                target_uniform = torch.ones_like(avg_selection) / self.llm_num
                llm_balance_loss = F.kl_div(
                    (avg_selection + 1e-8).log(), target_uniform, reduction='batchmean'
                )
                loss = loss + self.llm_balance_weight * llm_balance_loss

            if self.quality_entropy_weight > 0:
                probs = torch.softmax(logits_for_loss, dim=1)
                selection_entropy = -(probs * (probs + 1e-8).log()).sum(dim=1).mean()
                loss = loss - self.quality_entropy_weight * selection_entropy  # maximize

            # Uncertainty loss: heteroscedastic NLL + overconfidence penalty
            # This teaches the model to predict higher uncertainty when likely wrong
            # Staged training: only apply after warmup epochs (quality head stabilizes first)
            uncertainty_active = (
                self.use_uncertainty_head
                and log_var is not None
                and current_epoch >= self.uncertainty_warmup_epochs
            )
            if uncertainty_active:
                # Get binary quality targets (correct/incorrect)
                quality_targets = data.combined_edge[mask, 1]  # perf column

                # Predicted quality probabilities
                quality_preds = torch.sigmoid(logits)

                # Heteroscedastic NLL: penalizes wrong predictions, allows high
                # variance to reduce penalty (but log(var) prevents infinite var)
                uncertainty_nll = heteroscedastic_nll_loss(
                    quality_preds, quality_targets, log_var
                )

                # Overconfidence penalty: explicitly penalize low variance + wrong
                overconf_loss = overconfidence_penalty(
                    quality_preds, quality_targets, log_var, threshold=0.5
                )

                # Optional: ramp up weight linearly over post-warmup epochs
                warmup_progress = min(
                    1.0,
                    (current_epoch - self.uncertainty_warmup_epochs + 1)
                    / max(1, total_epochs - self.uncertainty_warmup_epochs),
                )
                effective_weight = self.uncertainty_weight * warmup_progress

                uncertainty_loss = (
                    effective_weight * uncertainty_nll
                    + self.overconfidence_penalty_weight * overconf_loss
                )
                loss = loss + uncertainty_loss

            # Standard mini-batch SGD: zero_grad -> backward -> step
            self.optimizer.zero_grad()
            loss.backward()

            if self.grad_clip > 0.0:
                clip_grad_norm_(self.model.parameters(), self.grad_clip)

            self.optimizer.step()

            # Update EMA weights after optimizer step
            if self.ema is not None:
                self.ema.update(self.model)

            # Update dynamic bias based on selection frequencies (loss-free balancing)
            self._update_dynamic_bias(logits_grouped)

            total_loss += loss.item()

        avg_loss = total_loss / max(1, num_batches)
        avg_diff_loss = total_diff_loss / max(1, num_batches) if self.use_difficulty_head else 0.0
        return avg_loss, avg_diff_loss

    def _train_epoch_bce(self, data: Data, batch_size: int) -> tuple[float, float]:
        """Run one training epoch with BCE loss.

        Returns:
            Tuple of (total_loss, difficulty_loss). BCE mode does not support
            difficulty head, so difficulty_loss is always 0.0.

        Uses standard mini-batch SGD with per-batch backward passes.
        """
        mask_train = data.edge_mask
        total_loss = 0.0

        for _ in range(batch_size):
            mask = mask_train.clone()
            random_mask = torch.rand(mask.shape, device=mask.device) < self.config["train_mask_rate"]
            mask = mask & ~random_mask
            edge_can_see = ~mask & self.train_mask

            wp, wc, wl = self._sample_pref()
            augmented_edge = self._edge_features_with_pref(data.combined_edge, wp, wc, wl)

            logits = self.model(
                task_features=data.task_id,
                query_features=data.query_features,
                llm_features=data.llm_features,
                edge_index=data.edge_index,
                edge_mask=mask,
                edge_can_see=edge_can_see,
                edge_weight=augmented_edge,
            )

            rewards = self._rewards_from_pref(data.combined_edge, wp, wc, wl)
            labels = self._labels_from_rewards(rewards)
            preds = torch.sigmoid(logits).reshape(-1)
            loss = self.criterion(preds, labels[mask].reshape(-1))

            # Standard mini-batch SGD: zero_grad -> backward -> step
            self.optimizer.zero_grad()
            loss.backward()

            if self.grad_clip > 0.0:
                clip_grad_norm_(self.model.parameters(), self.grad_clip)

            self.optimizer.step()

            # Update EMA weights after optimizer step
            if self.ema is not None:
                self.ema.update(self.model)

            total_loss += loss.item()

        return total_loss / batch_size, 0.0  # BCE mode doesn't support difficulty head

    def _train_epoch_quality_only(
        self,
        data: Data,
        train_query_indices: torch.Tensor,
        batch_size: int,
        current_epoch: int,
        total_epochs: int,
    ) -> tuple[float, float]:
        """Train one epoch using quality-only mode with pointwise BCE.

        Learns P(correct | query, LLM) directly without preference sampling.
        Cost and latency are NOT used during training - they're composed at inference.

        Key differences from reward-based training:
        1. No preference sampling (wp, wc, wl)
        2. Target is binary correctness (perf column), not composite reward
        3. Edge features don't include preferences (simpler)
        4. Loss is pointwise BCE or Brier, not listwise KL

        Args:
            data: PyG Data object with training data.
            train_query_indices: Indices of queries in training set.
            batch_size: Number of queries per mini-batch.
            current_epoch: Current epoch number (for difficulty weight scheduling).
            total_epochs: Total epochs (for difficulty weight scheduling).

        Returns:
            Tuple of (avg_loss, avg_difficulty_loss).
        """
        total_loss = 0.0
        total_diff_loss = 0.0

        # Compute sample weights combining hard negative mining and task weighting
        use_hard_mining = self.config.get("use_hard_negative_mining", False)
        use_task_weights = bool(self.task_weights) and data.query_task_labels is not None

        if use_hard_mining or use_task_weights:
            # Start with uniform weights
            num_train = train_query_indices.numel()
            sample_weights = torch.ones(num_train, device=train_query_indices.device)

            # Apply hard negative mining weights if enabled
            if use_hard_mining:
                hard_weights = self._compute_hard_negative_weights(
                    data.combined_edge,
                    train_query_indices,
                )
                sample_weights = sample_weights * hard_weights

            # Apply task-based weights if enabled
            if use_task_weights:
                task_multipliers = torch.ones(num_train, device=train_query_indices.device)
                for i, q_idx in enumerate(train_query_indices.tolist()):
                    task_label = data.query_task_labels[q_idx]
                    for prefix, weight in self.task_weights.items():
                        if task_label.startswith(prefix):
                            task_multipliers[i] = weight
                            break
                sample_weights = sample_weights * task_multipliers

            # Normalize and sample with replacement
            sample_weights = sample_weights / sample_weights.sum()
            perm = torch.multinomial(
                sample_weights,
                num_samples=num_train,
                replacement=True,
            )
        else:
            # Uniform random sampling (original behavior)
            perm = torch.randperm(len(train_query_indices), device=train_query_indices.device)

        shuffled_indices = train_query_indices[perm]
        num_batches = max(1, len(shuffled_indices) // batch_size)

        # Get quality targets (perf column from combined_edge)
        quality_targets = data.combined_edge[:, 1]  # perf is column 1

        # In quality-only mode, we still need to pad edge features to match model's
        # expected input dimension. Use zeros for preference weights since they're
        # not used during training - only at inference for score composition.
        # This maintains model architecture compatibility.
        edge_weight_padded = self._edge_features_with_pref(
            data.combined_edge, wp=0.0, wc=0.0, wl=0.0
        )

        model = self.model.module if hasattr(self.model, "module") else self.model
        best_llm_idx = None
        if hasattr(model, "use_contrastive") and model.use_contrastive:
            perf_matrix = quality_targets.view(-1, self.llm_num)
            best_llm_idx = perf_matrix.argmax(dim=1)

        preferences = None
        if hasattr(model, "use_hypernetwork") and model.use_hypernetwork:
            inf_weights = self.config.get("inference_weights", {})
            preferences = torch.tensor(
                [inf_weights.get("perf", 0.7), inf_weights.get("cost", 0.2), inf_weights.get("latency", 0.1)],
                dtype=torch.float32,
                device=data.combined_edge.device,
            )

        # Compute difficulty predictions once per epoch if using difficulty head
        difficulty_pred = None
        if self.use_difficulty_head:
            with torch.no_grad():
                _, difficulty_pred = self.model(
                    task_features=data.task_id,
                    query_features=data.query_features,
                    llm_features=data.llm_features,
                    edge_index=data.edge_index,
                    edge_mask=data.edge_mask,
                    edge_can_see=self.train_mask,
                    edge_weight=edge_weight_padded,
                    return_difficulty=True,
                    best_llm_idx=best_llm_idx,
                    preferences=preferences,
                )

        for batch_idx in range(num_batches):
            start_idx = batch_idx * batch_size
            end_idx = min(start_idx + batch_size, len(shuffled_indices))
            selected_queries = shuffled_indices[start_idx:end_idx]

            # Build mask for selected queries (vectorized)
            mask = self._build_query_mask(selected_queries, data.combined_edge.size(0))

            edge_can_see = ~mask & self.train_mask

            # Forward pass
            if self.use_difficulty_head:
                logits, difficulty_pred = self.model(
                    task_features=data.task_id,
                    query_features=data.query_features,
                    llm_features=data.llm_features,
                    edge_index=data.edge_index,
                    edge_mask=mask,
                    edge_can_see=edge_can_see,
                    edge_weight=edge_weight_padded,
                    return_difficulty=True,
                    best_llm_idx=best_llm_idx,
                    preferences=preferences,
                )
            else:
                logits = self.model(
                    task_features=data.task_id,
                    query_features=data.query_features,
                    llm_features=data.llm_features,
                    edge_index=data.edge_index,
                    edge_mask=mask,
                    edge_can_see=edge_can_see,
                    edge_weight=edge_weight_padded,
                    best_llm_idx=best_llm_idx,
                    preferences=preferences,
                )

            # Pointwise quality prediction
            quality_pred = torch.sigmoid(logits).reshape(-1)
            targets = quality_targets[mask]

            # Compute sample weights for needle query upweighting
            sample_weight = None
            if self.use_needle_weighting:
                num_queries_batch = len(selected_queries)
                sample_weight = self._compute_needle_query_weights(
                    targets, num_queries_batch
                )

            # Compute quality loss (with optional needle weighting).
            # FocalLoss/BrierLoss accept sample_weight; BCELoss does not.
            _supports_weight = isinstance(
                self.quality_criterion, (FocalLoss, BrierLoss)
            )
            if sample_weight is not None and _supports_weight:
                loss = self.quality_criterion(quality_pred, targets, sample_weight)
            elif sample_weight is not None:
                # Manual weighting for losses that don't accept sample_weight
                per_sample = F.binary_cross_entropy(
                    quality_pred.float(), targets.float(), reduction="none"
                )
                loss = (per_sample * sample_weight).sum() / sample_weight.sum().clamp(min=1e-8)
            else:
                loss = self.quality_criterion(quality_pred, targets)

            model = self.model.module if hasattr(self.model, "module") else self.model

            # Add calibration loss if enabled (Brier score)
            if self.calibration_weight > 0:
                calibration_loss = self._brier_score(
                    quality_pred.unsqueeze(0), targets.unsqueeze(0)
                )
                loss = loss + self.calibration_weight * calibration_loss

            # Add difficulty prediction loss if enabled
            if self.use_difficulty_head and difficulty_pred is not None:
                difficulty_target = self._compute_difficulty_targets(
                    data.combined_edge, mask, selected_queries
                )
                difficulty_pred_batch = difficulty_pred[selected_queries]
                difficulty_loss = self.difficulty_criterion(
                    difficulty_pred_batch, difficulty_target
                )
                total_diff_loss += difficulty_loss.item()

                current_difficulty_weight = self._get_difficulty_weight(
                    current_epoch, total_epochs
                )
                loss = loss + current_difficulty_weight * difficulty_loss

            if hasattr(model, "get_auxiliary_losses"):
                aux_losses = model.get_auxiliary_losses()
                for aux_name, aux_loss in aux_losses.items():
                    if aux_loss is not None:
                        loss = loss + aux_loss

            if self.use_llm_balance_loss:
                # Compute actual query count from predictions (handles missing data)
                actual_edges = quality_pred.size(0)
                if actual_edges % self.llm_num == 0:
                    num_queries_actual = actual_edges // self.llm_num
                    probs_grouped = quality_pred.view(num_queries_actual, self.llm_num)
                    selection_probs = F.softmax(probs_grouped, dim=1)
                    avg_selection = selection_probs.mean(dim=0)
                    target = torch.ones_like(avg_selection) / self.llm_num
                    llm_balance_loss = F.kl_div(
                        (avg_selection + 1e-8).log(), target, reduction='batchmean'
                    )
                    loss = loss + self.llm_balance_weight * llm_balance_loss

            if self.quality_entropy_weight > 0:
                p = quality_pred.clamp(1e-7, 1 - 1e-7)
                entropy = -p * p.log() - (1 - p) * (1 - p).log()
                loss = loss - self.quality_entropy_weight * entropy.mean()  # maximize

            # Backward pass (SAM requires special two-step handling)
            if isinstance(self.optimizer, SAM):
                self._sam_backward_step(loss)
            else:
                self.optimizer.zero_grad()
                loss.backward()
                if self.grad_clip > 0.0:
                    clip_grad_norm_(self.model.parameters(), self.grad_clip)
                self.optimizer.step()

            # Update EMA weights after optimizer step
            if self.ema is not None:
                self.ema.update(self.model)

            total_loss += loss.item()

        avg_loss = total_loss / max(1, num_batches)
        avg_diff_loss = total_diff_loss / max(1, num_batches) if self.use_difficulty_head else 0.0
        return avg_loss, avg_diff_loss

    def _evaluate_difficulty_head(
        self,
        data: Data,
        mask: torch.Tensor,
    ) -> dict[str, float]:
        """Evaluate difficulty head predictions on given data.

        Computes accuracy and AUC-ROC for the difficulty classification/regression head.

        Args:
            data: Data object containing query features and combined_edge.
            mask: Boolean mask indicating which edges to evaluate.

        Returns:
            Dict with difficulty head metrics:
            - diff_acc: Classification accuracy (for classification mode)
            - diff_auc: AUC-ROC score
            - diff_loss: Average loss on this data
        """
        if not self.use_difficulty_head:
            return {}

        self.model.eval()

        # Get query indices from mask
        query_mask = mask.view(-1, self.llm_num).all(dim=1)
        query_indices = torch.nonzero(query_mask, as_tuple=False).squeeze(-1)

        if query_indices.numel() == 0:
            return {"diff_acc": 0.0, "diff_auc": 0.5, "diff_loss": 0.0}

        wp, wc, wl = self.infer_wp, self.infer_wc, self.infer_wl

        with torch.no_grad():
            augmented_edge = self._edge_features_with_pref(data.combined_edge, wp, wc, wl)
            _, difficulty_pred = self.model(
                task_features=data.task_id,
                query_features=data.query_features,
                llm_features=data.llm_features,
                edge_index=data.edge_index,
                edge_mask=mask,
                edge_can_see=self.train_mask,
                edge_weight=augmented_edge,
                return_difficulty=True,
            )

            # Get predictions for evaluated queries
            difficulty_pred_eval = difficulty_pred[query_indices]

            # Compute targets
            difficulty_target = self._compute_difficulty_targets(
                data.combined_edge, mask, query_indices
            )

            # Compute loss
            diff_loss = self.difficulty_criterion(difficulty_pred_eval, difficulty_target)

        # Convert to numpy for sklearn metrics
        pred_np = difficulty_pred_eval.cpu().numpy().squeeze()
        target_np = difficulty_target.cpu().numpy()

        # For classification mode, compute accuracy
        if self.difficulty_head_mode == "classification":
            pred_binary = (pred_np > 0.5).astype(float)
            diff_acc = (pred_binary == target_np).mean()
        else:
            # For regression, use threshold-based accuracy
            threshold = self.config.get("difficulty_classification_threshold", 0.15)
            pred_binary = (pred_np < threshold).astype(float)
            target_binary = (target_np < threshold).astype(float)
            diff_acc = (pred_binary == target_binary).mean()

        # Compute AUC-ROC (requires both classes present)
        try:
            if len(np.unique(target_np)) > 1:
                diff_auc = roc_auc_score(target_np, pred_np)
            else:
                diff_auc = 0.5  # Default when only one class present
        except ValueError:
            diff_auc = 0.5

        # Compute comprehensive calibration metrics
        calib_metrics = compute_calibration_metrics(pred_np, target_np, n_bins=10)

        return {
            "diff_acc": float(diff_acc),
            "diff_auc": float(diff_auc),
            "diff_loss": float(diff_loss.item()),
            "diff_ece": calib_metrics["ece"],
            "diff_mce": calib_metrics["mce"],
            "diff_brier": calib_metrics["brier"],
            "diff_reliability_bins": calib_metrics["reliability_bins"],
        }

    def _validate(self, data: Data) -> tuple[dict[str, float], float, float, float, float, float]:
        """Run validation with fixed inference preferences."""
        self.model.eval()

        # Apply EMA weights for validation if enabled
        if self.ema is not None:
            self.ema.apply_shadow(self.model)

        try:
            mask = data.edge_mask.clone()

            # Use FIXED inference weights for validation (not random!)
            wp, wc, wl = self.infer_wp, self.infer_wc, self.infer_wl

            if self.quality_only_training:
                return self._validate_quality_only(data, mask, wp, wc, wl)

            with torch.no_grad():
                augmented_edge = self._edge_features_with_pref(data.combined_edge, wp, wc, wl)

                # Forward pass with optional uncertainty
                log_var = None
                if self.use_uncertainty_head:
                    forward_outputs = self.model(
                        task_features=data.task_id,
                        query_features=data.query_features,
                        llm_features=data.llm_features,
                        edge_index=data.edge_index,
                        edge_mask=mask,
                        edge_can_see=self.train_mask,
                        edge_weight=augmented_edge,
                        return_uncertainty=True,
                    )
                    logits, log_var = forward_outputs
                else:
                    logits = self.model(
                        task_features=data.task_id,
                        query_features=data.query_features,
                        llm_features=data.llm_features,
                        edge_index=data.edge_index,
                        edge_mask=mask,
                        edge_can_see=self.train_mask,
                        edge_weight=augmented_edge,
                    )

            if self.adaptive_weights:
                rewards = self._adaptive_rewards_from_pref(data.combined_edge, wp, wc, wl)
            else:
                rewards = self._rewards_from_pref(data.combined_edge, wp, wc, wl)
            logits_grouped = logits.reshape(-1, self.llm_num)
            rewards_grouped = rewards[mask].reshape(-1, self.llm_num)

            # Compute accuracy and F1
            observe_idx = torch.argmax(logits_grouped, dim=1)
            label_idx = torch.argmax(rewards_grouped, dim=1)
            total = label_idx.size(0)

            if total > 0:
                accuracy = (observe_idx == label_idx).float().mean().item()
                f1 = f1_score(label_idx.cpu().numpy(), observe_idx.cpu().numpy(), average="macro")
            else:
                accuracy = 0.0
                f1 = 0.0

            # Compute loss
            if self.loss_type == "listwise":
                soft_target = torch.softmax(rewards_grouped / self.tau, dim=1).clamp_min(1e-8)
                # criterion uses reduction="none", manually compute mean
                per_query_kl = self.criterion(
                    torch.log_softmax(logits_grouped, dim=1), soft_target
                ).sum(dim=1)
                loss = per_query_kl.mean()
            else:
                labels = self._labels_from_rewards(rewards)
                loss = self.criterion(torch.sigmoid(logits).reshape(-1), labels[mask].reshape(-1))

            # Compute detailed metrics
            metrics = self._compute_eval_metrics(logits_grouped, rewards_grouped)

            # Compute avg cost/perf/latency for chosen LLMs
            feature_dim = data.combined_edge.size(1)
            combined = data.combined_edge[mask].reshape(-1, self.llm_num, feature_dim)
            row_indices = torch.arange(combined.size(0), device=combined.device)

            avg_cost = combined[row_indices, observe_idx, 0].mean().item() if combined.size(0) else 0.0
            avg_perf = combined[row_indices, observe_idx, 1].mean().item() if combined.size(0) else 0.0
            avg_lat = (
                combined[row_indices, observe_idx, 2].mean().item()
                if (feature_dim > 2 and combined.size(0))
                else 0.0
            )

            # Compute uncertainty metrics if enabled
            if self.use_uncertainty_head and log_var is not None:
                quality_targets = data.combined_edge[mask, 1]  # perf column
                quality_preds = torch.sigmoid(logits)
                uncertainty_metrics = compute_uncertainty_metrics(
                    quality_preds, quality_targets, log_var
                )
                metrics.update({f"unc_{k}": v for k, v in uncertainty_metrics.items()})

            return metrics, loss.item(), accuracy, f1, avg_cost, avg_perf, avg_lat
        finally:
            # Restore original weights after validation
            if self.ema is not None:
                self.ema.restore(self.model)

    def _validate_quality_only(
        self,
        data: Data,
        mask: torch.Tensor,
        wp: float,
        wc: float,
        wl: float,
    ) -> tuple[dict[str, float], float, float, float, float, float]:
        """Validate in quality-only mode with inference-time score composition.

        In quality-only mode:
        1. Model outputs P(correct | query, LLM) for each edge
        2. At inference, we compose: score = wp * quality - wc * cost - wl * latency
        3. Selection is based on composed scores

        This decouples quality learning from cost/latency preferences.
        """
        # Pad edge features with zeros to match model's expected input dimension
        edge_weight_padded = self._edge_features_with_pref(
            data.combined_edge, wp=0.0, wc=0.0, wl=0.0
        )

        with torch.no_grad():
            # Forward pass with zero-padded edge weights for model compatibility
            logits = self.model(
                task_features=data.task_id,
                query_features=data.query_features,
                llm_features=data.llm_features,
                edge_index=data.edge_index,
                edge_mask=mask,
                edge_can_see=self.train_mask,
                edge_weight=edge_weight_padded,
            )

        # Get quality predictions (calibrated probabilities)
        # Note: model returns predictions only for masked edges
        quality_pred = torch.sigmoid(logits).reshape(-1)  # (num_masked_edges,)

        # Extract cost, quality (ground truth), latency for MASKED edges only
        # to match quality_pred shape
        cost = data.combined_edge[mask, 0]
        quality_gt = data.combined_edge[mask, 1]  # Ground truth perf
        latency = (
            data.combined_edge[mask, 2]
            if data.combined_edge.size(1) > 2
            else torch.zeros_like(cost)
        )

        # Compose inference-time scores: score = wp * quality_pred - wc * cost - wl * latency
        # All tensors now have shape (num_masked_edges,)
        composed_scores = wp * quality_pred - wc * cost - wl * latency

        # Ground truth rewards (for comparison) use actual quality
        gt_rewards = wp * quality_gt - wc * cost - wl * latency

        # Group by query - reshape directly since we already have masked edges
        scores_grouped = composed_scores.reshape(-1, self.llm_num)
        gt_rewards_grouped = gt_rewards.reshape(-1, self.llm_num)
        quality_pred_grouped = quality_pred.reshape(-1, self.llm_num)
        quality_gt_grouped = quality_gt.reshape(-1, self.llm_num)

        # Selection based on composed scores
        observe_idx = torch.argmax(scores_grouped, dim=1)
        label_idx = torch.argmax(gt_rewards_grouped, dim=1)
        total = label_idx.size(0)

        if total > 0:
            accuracy = (observe_idx == label_idx).float().mean().item()
            f1 = f1_score(label_idx.cpu().numpy(), observe_idx.cpu().numpy(), average="macro")
        else:
            accuracy = 0.0
            f1 = 0.0

        # Compute quality loss (BCE on quality predictions vs ground truth)
        # quality_pred and quality_gt are already masked
        quality_loss = self.quality_criterion(quality_pred, quality_gt)

        # Compute detailed metrics using composed scores
        metrics = self._compute_eval_metrics(scores_grouped, gt_rewards_grouped)

        # Add quality-specific metrics
        metrics["quality_brier"] = F.mse_loss(quality_pred, quality_gt).item()
        metrics["quality_pred_mean"] = quality_pred.mean().item()
        metrics["quality_pred_std"] = quality_pred.std().item()

        # Compute difficulty from quality predictions (variance across LLMs per query)
        quality_variance = quality_pred_grouped.std(dim=1)
        metrics["avg_quality_variance"] = quality_variance.mean().item()

        # Per-query mean quality: distinguishes trivial (all high) from impossible (all low)
        quality_mean_per_query = quality_pred_grouped.mean(dim=1)  # [num_queries]
        metrics["quality_mean_per_query_avg"] = quality_mean_per_query.mean().item()

        # Stratify frac_oracle by GROUND-TRUTH difficulty quartiles
        # Key insight (Dean): Use ground-truth difficulty (num LLMs succeed) instead of
        # model-estimated difficulty (entropy). Ground-truth is always trustworthy,
        # while model estimates require calibration which isn't available early in training.
        #
        # Difficulty = num_winners (how many LLMs succeed on each query)
        # - q0 (trivial): many LLMs succeed → LLM selection doesn't matter
        # - q3 (hard/needle): few LLMs succeed → LLM selection is critical
        num_winners = gt_rewards_grouped.sum(dim=1)  # [num_queries]

        # Use quartile boundaries from the data itself (robust to class imbalance)
        # Lower num_winners = harder query
        quartile_boundaries = torch.quantile(
            num_winners.float(),
            torch.tensor([0.25, 0.5, 0.75], device=num_winners.device)
        )

        # Create masks: q3 = hardest (fewest winners), q0 = easiest (most winners)
        q3_mask = num_winners <= quartile_boundaries[0]  # Hardest 25%
        q2_mask = (num_winners > quartile_boundaries[0]) & (num_winners <= quartile_boundaries[1])
        q1_mask = (num_winners > quartile_boundaries[1]) & (num_winners <= quartile_boundaries[2])
        q0_mask = num_winners > quartile_boundaries[2]  # Easiest 25%

        difficulty_masks = {"q0": q0_mask, "q1": q1_mask, "q2": q2_mask, "q3": q3_mask}

        for quartile, qmask in difficulty_masks.items():
            if qmask.sum() > 0:
                quartile_scores = scores_grouped[qmask]
                quartile_rewards = gt_rewards_grouped[qmask]
                chosen_vals = quartile_rewards[
                    torch.arange(quartile_rewards.size(0), device=quartile_rewards.device),
                    quartile_scores.argmax(dim=1)
                ]
                oracle_vals = quartile_rewards.max(dim=1).values
                min_vals = quartile_rewards.min(dim=1).values
                denom = (oracle_vals - min_vals).clamp(min=1e-8)
                frac = ((chosen_vals - min_vals) / denom).mean().item()
                metrics[f"frac_oracle_{quartile}"] = frac

        # Log ground-truth difficulty distribution for debugging
        metrics["gt_difficulty_mean"] = num_winners.float().mean().item()
        metrics["gt_difficulty_std"] = num_winners.float().std().item()

        # Compute correlation between model signals and ground-truth difficulty
        # This tells us if the router "knows what it doesn't know"
        from .logit_signals import (
            compute_uncertainty_from_logits,
            compute_difficulty_from_logits,
            compute_difficulty_from_logit_gap,
        )
        model_uncertainty = compute_uncertainty_from_logits(scores_grouped)
        model_difficulty = compute_difficulty_from_logits(scores_grouped)
        model_gap_difficulty = compute_difficulty_from_logit_gap(scores_grouped)

        # num_winners is INVERTED difficulty (more winners = easier)
        # So correlations should be NEGATIVE if well-calibrated
        # (high uncertainty/difficulty when few winners)
        if len(model_uncertainty) > 10:
            gt = num_winners.float()
            metrics["uncertainty_gt_corr"] = torch.corrcoef(
                torch.stack([model_uncertainty, gt])
            )[0, 1].item()
            metrics["difficulty_gt_corr"] = torch.corrcoef(
                torch.stack([model_difficulty, gt])
            )[0, 1].item()
            metrics["gap_difficulty_gt_corr"] = torch.corrcoef(
                torch.stack([model_gap_difficulty, gt])
            )[0, 1].item()

            # Quality mean correlation: should be POSITIVE if calibrated
            # (higher mean quality when more LLMs succeed)
            metrics["quality_mean_gt_corr"] = torch.corrcoef(
                torch.stack([quality_mean_per_query, gt])
            )[0, 1].item()

            # Distinguish trivial (all succeed) from impossible (all fail)
            trivial_mask = num_winners == self.llm_num  # All 7 succeed
            impossible_mask = num_winners == 0  # All 7 fail
            if trivial_mask.sum() > 0:
                metrics["trivial_quality_mean"] = quality_mean_per_query[trivial_mask].mean().item()
                metrics["trivial_entropy"] = model_uncertainty[trivial_mask].mean().item()
                metrics["trivial_count"] = trivial_mask.sum().item()
            if impossible_mask.sum() > 0:
                metrics["impossible_quality_mean"] = quality_mean_per_query[impossible_mask].mean().item()
                metrics["impossible_entropy"] = model_uncertainty[impossible_mask].mean().item()
                metrics["impossible_count"] = impossible_mask.sum().item()

        # Compute avg cost/perf/latency for chosen LLMs
        feature_dim = data.combined_edge.size(1)
        combined = data.combined_edge[mask].reshape(-1, self.llm_num, feature_dim)
        row_indices = torch.arange(combined.size(0), device=combined.device)

        avg_cost = combined[row_indices, observe_idx, 0].mean().item() if combined.size(0) else 0.0
        avg_perf = combined[row_indices, observe_idx, 1].mean().item() if combined.size(0) else 0.0
        avg_lat = (
            combined[row_indices, observe_idx, 2].mean().item()
            if (feature_dim > 2 and combined.size(0))
            else 0.0
        )

        return metrics, quality_loss.item(), accuracy, f1, avg_cost, avg_perf, avg_lat

    def train_validate(
        self,
        data: Data,
        data_validate: Data,
        data_for_test: Data,
    ) -> None:
        """Train the model with validation-based early stopping."""
        logger.info("train_validate: Setting up...")

        self.save_path = self.config["model_path"]
        self.train_mask = data.train_mask.clone()
        self.valid_mask = data.valid_mask.clone()
        self.test_mask = data.test_mask.clone()

        epochs = int(self.config["train_epoch"])
        batch_size = int(self.config["batch_size"])

        # Eval-only mode: load weights and return without training.
        if self.config.get("eval_only", False):
            best_path = Path(self.save_path)
            if best_path.is_file():
                device = next(self.model.parameters()).device
                state_dict = torch.load(
                    best_path, map_location=device, weights_only=True,
                )
                self.model.load_state_dict(state_dict)
                logger.info("eval_only: loaded weights from %s", best_path)
            else:
                logger.warning(
                    "eval_only: no weights at %s, using random init",
                    best_path,
                )
            return

        best_score = float("-inf")
        no_improve = 0
        start_epoch = 0

        # Try to resume from checkpoint
        checkpoint = self._load_checkpoint()
        if checkpoint is not None:
            start_epoch = checkpoint["epoch"] + 1
            best_score = checkpoint["best_score"]
            no_improve = checkpoint["no_improve"]
            logger.info(
                "Resuming from epoch %d (best_score=%.4f, no_improve=%d)",
                start_epoch + 1,
                best_score,
                no_improve,
            )

        # Get training query indices
        train_queries = self.train_mask.view(-1, self.llm_num).all(dim=1)
        train_query_indices = torch.nonzero(train_queries, as_tuple=False).squeeze(-1)

        logger.info(
            "Starting training loop: epochs %d-%d, %d train queries",
            start_epoch + 1,
            epochs,
            len(train_query_indices),
        )

        for epoch in range(start_epoch, epochs):
            logger.info("Epoch %d/%d starting...", epoch + 1, epochs)

            self.model.train()

            # Training step (backward is called inside each method)
            if self.quality_only_training:
                # Quality-only mode: learn P(correct | query, LLM) directly
                mean_loss, train_diff_loss = self._train_epoch_quality_only(
                    data, train_query_indices, batch_size, epoch, epochs
                )
            elif self.loss_type == "listwise":
                mean_loss, train_diff_loss = self._train_epoch_listwise(
                    data, train_query_indices, batch_size, epoch, epochs
                )
            else:
                mean_loss, train_diff_loss = self._train_epoch_bce(data, batch_size)

            self.scheduler.step()

            # Validation
            metrics_val, loss_val, acc_val, f1_val, avg_cost, avg_perf, avg_lat = self._validate(
                data_validate
            )

            # Evaluate difficulty head on validation data
            diff_metrics_val = self._evaluate_difficulty_head(data_validate, data_validate.edge_mask)

            # Log epoch summary
            logger.info(
                "Epoch %d complete: train_loss=%.4f, val_acc=%.3f, val_f1=%.3f, frac_oracle=%.3f",
                epoch + 1,
                mean_loss,
                acc_val,
                f1_val,
                metrics_val.get("frac_of_oracle", 0.0),
            )
            # Log difficulty head metrics to console for debugging collapse
            if diff_metrics_val:
                logger.info(
                    "  Difficulty: acc=%.3f, auc=%.3f, ece=%.3f, loss=%.4f",
                    diff_metrics_val.get("diff_acc", 0.0),
                    diff_metrics_val.get("diff_auc", 0.0),
                    diff_metrics_val.get("diff_ece", 0.0),
                    diff_metrics_val.get("diff_loss", 0.0),
                )

            # Log uncertainty-difficulty correlations (negative = well-calibrated)
            if "uncertainty_gt_corr" in metrics_val:
                logger.info(
                    "  Uncertainty-GT correlations: uncertainty=%.3f, difficulty=%.3f, gap=%.3f",
                    metrics_val.get("uncertainty_gt_corr", 0.0),
                    metrics_val.get("difficulty_gt_corr", 0.0),
                    metrics_val.get("gap_difficulty_gt_corr", 0.0),
                )
                # Log quality mean correlation and trivial/impossible discrimination
                # Two orthogonal signals: entropy (routing) vs quality_mean (success)
                logger.info(
                    "  Quality-mean corr=%.3f | Trivial(n=%d): qmean=%.3f, entropy=%.3f | Impossible(n=%d): qmean=%.3f, entropy=%.3f",
                    metrics_val.get("quality_mean_gt_corr", 0.0),
                    int(metrics_val.get("trivial_count", 0)),
                    metrics_val.get("trivial_quality_mean", 0.0),
                    metrics_val.get("trivial_entropy", 0.0),
                    int(metrics_val.get("impossible_count", 0)),
                    metrics_val.get("impossible_quality_mean", 0.0),
                    metrics_val.get("impossible_entropy", 0.0),
                )

            # Log hybrid loss components for gradient diagnosis
            if hasattr(self, '_loss_components') and self._loss_components.get('kl'):
                import numpy as np
                kl_mean = np.mean(self._loss_components['kl'])
                bce_mean = np.mean(self._loss_components['bce'])
                pw_mean = np.mean(self._loss_components['pairwise'])
                logger.info(
                    "  Loss components: KL=%.4f, BCE=%.4f, Pairwise=%.4f (ratio KL:BCE:PW = 1:%.2f:%.2f)",
                    kl_mean, bce_mean, pw_mean,
                    bce_mean / (kl_mean + 1e-8),
                    pw_mean / (kl_mean + 1e-8),
                )
                # Reset for next epoch
                self._loss_components = {'kl': [], 'bce': [], 'pairwise': []}

            # Early stopping check - support weighted metrics that prioritize hard queries
            # If early_stop_weights is configured, use weighted combination of quartile metrics
            # Default: uniform frac_of_oracle (backwards compatible)
            #
            # NOTE: With ground-truth difficulty (num_winners), warmup is OPTIONAL.
            # Ground-truth difficulty is always valid, unlike model-estimated difficulty
            # which requires calibration. Keep small warmup (e.g., 5 epochs) if desired
            # to avoid selecting checkpoints before meaningful learning occurs.
            early_stop_warmup = self.config.get("early_stop_warmup", 0)
            early_stop_weights = self.config.get("early_stop_weights", None)
            if early_stop_weights and all(
                f"frac_oracle_{q}" in metrics_val for q in ["q0", "q1", "q2", "q3"]
            ):
                # Compute weighted frac_oracle: sum(weight_i * frac_oracle_qi)
                w0 = early_stop_weights.get("q0", 0.25)
                w1 = early_stop_weights.get("q1", 0.25)
                w2 = early_stop_weights.get("q2", 0.25)
                w3 = early_stop_weights.get("q3", 0.25)
                score_monitor = (
                    w0 * metrics_val["frac_oracle_q0"] +
                    w1 * metrics_val["frac_oracle_q1"] +
                    w2 * metrics_val["frac_oracle_q2"] +
                    w3 * metrics_val["frac_oracle_q3"]
                )
                logger.info(
                    "  Weighted score: %.4f (q0=%.3f×%.1f, q1=%.3f×%.1f, q2=%.3f×%.1f, q3=%.3f×%.1f)",
                    score_monitor,
                    metrics_val["frac_oracle_q0"], w0,
                    metrics_val["frac_oracle_q1"], w1,
                    metrics_val["frac_oracle_q2"], w2,
                    metrics_val["frac_oracle_q3"], w3,
                )
            else:
                score_monitor = metrics_val.get("frac_of_oracle", acc_val)

            # Skip model saving during warmup period (uncertainty signals unreliable)
            if epoch < early_stop_warmup:
                logger.info("  Warmup epoch %d/%d - skipping model selection", epoch + 1, early_stop_warmup)
                no_improve = 0  # Don't count warmup epochs toward early stopping
                continue

            if score_monitor > best_score:
                best_score = score_monitor
                torch.save(self.model.state_dict(), self.save_path)
                logger.info("  New best model saved (score=%.4f)", score_monitor)
                no_improve = 0
            else:
                no_improve += 1
                if self.early_stop_patience > 0 and no_improve >= self.early_stop_patience:
                    logger.info("Early stopping at epoch %d (no improvement for %d epochs)", epoch + 1, no_improve)
                    self._clear_checkpoint()
                    break

            # Test evaluation
            test_metrics, test_loss = self.test(data_for_test, self.config["model_path"])

            # Evaluate difficulty head on test data
            diff_metrics_test = self._evaluate_difficulty_head(data_for_test, data_for_test.edge_mask)

            # Log to W&B
            log_payload = {
                "epoch": epoch + 1,
                "lr": self.optimizer.param_groups[0]["lr"],
                "train_loss": mean_loss,
                "train_diff_loss": train_diff_loss,
                "validate_loss": loss_val,
                "test_loss": test_loss,
                "val/top1_acc": acc_val,
                "val/f1_macro": f1_val,
                "val/avg_cost_norm": avg_cost,
                "val/avg_perf": avg_perf,
                "val/avg_latency_norm": avg_lat,
            }
            log_payload.update({f"val/{k}": v for k, v in metrics_val.items()})
            # Filter out non-scalar metrics (reliability_bins) for WandB logging
            log_payload.update({
                f"val/{k}": v for k, v in diff_metrics_val.items()
                if not isinstance(v, (list, dict))
            })
            log_payload.update({f"test/{k}": v for k, v in test_metrics.items()})
            log_payload.update({
                f"test/{k}": v for k, v in diff_metrics_test.items()
                if not isinstance(v, (list, dict))
            })
            self.wandb.log(log_payload)

            # Save checkpoint for crash recovery
            self._save_checkpoint(epoch, best_score, no_improve)

        # Training complete - clear checkpoint and load best model
        self._clear_checkpoint()
        best_path = Path(self.save_path)
        if best_path.is_file():
            device = next(self.model.parameters()).device
            state_dict = torch.load(best_path, map_location=device, weights_only=True)
            self.model.load_state_dict(state_dict)
            logger.info("Loaded best model from %s", best_path)

    def test(self, data: Data, model_path: str) -> tuple[dict[str, float], float]:
        """Evaluate the model on test data.

        When use_adaptive_inference is enabled, the difficulty head predictions
        are used to compute per-query weights, enabling dynamic cost-vs-accuracy
        tradeoffs based on predicted query difficulty.
        """
        self.model.eval()
        mask = data.edge_mask.clone()
        edge_can_see = self.valid_mask | self.train_mask

        wp, wc, wl = self.infer_wp, self.infer_wc, self.infer_wl

        with torch.no_grad():
            augmented_edge = self._edge_features_with_pref(data.combined_edge, wp, wc, wl)

            # Get logits and optionally difficulty predictions
            if self.use_adaptive_inference:
                logits, difficulty_pred = self.model(
                    task_features=data.task_id,
                    query_features=data.query_features,
                    llm_features=data.llm_features,
                    edge_index=data.edge_index,
                    edge_mask=mask,
                    edge_can_see=edge_can_see,
                    edge_weight=augmented_edge,
                    return_difficulty=True,
                )
            else:
                logits = self.model(
                    task_features=data.task_id,
                    query_features=data.query_features,
                    llm_features=data.llm_features,
                    edge_index=data.edge_index,
                    edge_mask=mask,
                    edge_can_see=edge_can_see,
                    edge_weight=augmented_edge,
                )
                difficulty_pred = None

        # Compute rewards: use adaptive weights if enabled, else fixed weights
        if self.use_adaptive_inference and difficulty_pred is not None:
            rewards = self._adaptive_rewards_from_difficulty(
                data.combined_edge, difficulty_pred
            )
        else:
            rewards = self._rewards_from_pref(data.combined_edge, wp, wc, wl)

        # Compute loss
        if self.loss_type == "listwise":
            grouped_rewards = rewards[mask].reshape(-1, self.llm_num)
            logits_grouped = logits.reshape(-1, self.llm_num)
            soft_target = torch.softmax(grouped_rewards / self.tau, dim=1).clamp_min(1e-8)
            # criterion uses reduction="none", manually compute mean
            per_query_kl = self.criterion(
                torch.log_softmax(logits_grouped, dim=1), soft_target
            ).sum(dim=1)
            loss = per_query_kl.mean()
        else:
            labels = self._labels_from_rewards(rewards)
            loss = self.criterion(torch.sigmoid(logits).reshape(-1), labels[mask].reshape(-1))

        scores = logits.reshape(-1, self.llm_num)
        rewards_grouped = rewards[mask].reshape(-1, self.llm_num)

        if rewards_grouped.size(0) == 0:
            empty_metrics = {
                "chosen_reward": 0.0,
                "oracle_reward": 0.0,
                "top1_acc": 0.0,
                "regret": 0.0,
                "frac_of_oracle": 0.0,
                "uplift_norm": 0.0,
                "ndcg@1": 0.0,
                "ndcg@3": 0.0,
                "ndcg@All": 0.0,
                "ece": 0.0,
                "avg_cost": 0.0,
                "avg_perf": 0.0,
                "avg_latency": 0.0,
                "avg_chosen_reward": 0.0,
                "avg_oracle_reward": 0.0,
            }
            return empty_metrics, float(loss.item())

        metrics_core = self._compute_eval_metrics(scores, rewards_grouped)

        row_indices = torch.arange(rewards_grouped.size(0), device=rewards_grouped.device)
        chosen_idx = torch.argmax(scores, dim=1)
        chosen_reward = rewards_grouped[row_indices, chosen_idx].mean().item()
        oracle_reward = rewards_grouped.max(dim=1).values.mean().item()

        feature_dim = data.combined_edge.size(1)
        combined = data.combined_edge[mask].reshape(-1, self.llm_num, feature_dim)
        avg_cost = combined[row_indices, chosen_idx, 0].mean().item()
        avg_perf = combined[row_indices, chosen_idx, 1].mean().item()
        avg_lat = combined[row_indices, chosen_idx, 2].mean().item() if feature_dim > 2 else 0.0

        # Zero-cost signal extraction (no trainable parameters)
        from .logit_signals import compute_routing_signals, stratify_by_difficulty

        temperature = self.learned_temperature.item() if hasattr(self, 'learned_temperature') else 1.0
        signals = compute_routing_signals(scores, temperature=temperature)

        avg_uncertainty = signals["uncertainty"].mean().item()
        avg_difficulty = signals["difficulty"].mean().item()
        avg_confidence = signals["confidence"].mean().item()

        # Stratify frac_oracle by difficulty quartiles
        difficulty_masks = stratify_by_difficulty(signals, num_quantiles=4)

        stratified_metrics = {}
        for quartile, mask in difficulty_masks.items():
            if mask.sum() > 0:
                quartile_scores = scores[mask]
                quartile_rewards = rewards_grouped[mask]
                # frac_of_oracle for this quartile
                chosen_vals = quartile_rewards[
                    torch.arange(quartile_rewards.size(0), device=quartile_rewards.device),
                    quartile_scores.argmax(dim=1)
                ]
                oracle_vals = quartile_rewards.max(dim=1).values
                min_vals = quartile_rewards.min(dim=1).values
                span = oracle_vals - min_vals
                # Handle ties: for all-correct or all-wrong, any choice = oracle -> frac = 1.0
                is_tie = span < 1e-6
                per_query_frac = torch.where(
                    is_tie,
                    torch.ones_like(chosen_vals),
                    (chosen_vals - min_vals) / span.clamp_min(1e-6),
                )
                frac = per_query_frac.mean().item()
                stratified_metrics[f"frac_oracle_{quartile}"] = frac

        metrics = {
            "chosen_reward": chosen_reward,
            "oracle_reward": oracle_reward,
            "avg_cost": avg_cost,
            "avg_perf": avg_perf,
            "avg_latency": avg_lat,
            "avg_uncertainty": avg_uncertainty,
            "avg_difficulty": avg_difficulty,
            "avg_confidence": avg_confidence,
            **stratified_metrics,
            **metrics_core,
        }

        logger.info("Test metrics: %s", metrics)
        return metrics, float(loss.item())

    def calibrate_temperature(
        self,
        data: Data,
        lr: float = 0.01,
        max_iters: int = 50,
    ) -> float:
        """Learn optimal temperature for calibration using validation data.

        Post-hoc temperature scaling (Platt scaling) finds a scalar T that
        minimizes NLL on the validation set: calibrated_probs = softmax(logits/T).
        This preserves ranking order while improving probability calibration.

        Args:
            data: Graph data with validation mask.
            lr: Learning rate for temperature optimization.
            max_iters: Maximum optimization iterations.

        Returns:
            Optimal temperature value (also stored in self.learned_temperature).
        """
        self.model.eval()
        mask = self.valid_mask
        edge_can_see = self.train_mask

        wp, wc, wl = self.infer_wp, self.infer_wc, self.infer_wl

        with torch.no_grad():
            augmented_edge = self._edge_features_with_pref(data.combined_edge, wp, wc, wl)
            logits = self.model(
                task_features=data.task_id,
                query_features=data.query_features,
                llm_features=data.llm_features,
                edge_index=data.edge_index,
                edge_mask=mask,
                edge_can_see=edge_can_see,
                edge_weight=augmented_edge,
            )
            rewards = self._rewards_from_pref(data.combined_edge, wp, wc, wl)

        logits_grouped = logits.reshape(-1, self.llm_num)
        rewards_grouped = rewards[mask].reshape(-1, self.llm_num)
        best_idx = rewards_grouped.argmax(dim=1)

        # Optimize temperature using gradient descent on NLL
        temperature = torch.tensor(1.0, requires_grad=True, device=logits.device)
        optimizer = torch.optim.LBFGS([temperature], lr=lr, max_iter=max_iters)

        def closure() -> torch.Tensor:
            optimizer.zero_grad()
            scaled_logits = logits_grouped / temperature.clamp(min=0.01)
            nll = F.cross_entropy(scaled_logits, best_idx)
            nll.backward()
            return nll

        optimizer.step(closure)

        optimal_temp = float(temperature.detach().clamp(min=0.01).item())
        self.learned_temperature.data.fill_(optimal_temp)

        # Compute ECE before and after calibration
        probs_before = torch.softmax(logits_grouped, dim=1)
        probs_after = torch.softmax(logits_grouped / optimal_temp, dim=1)

        correct = (probs_before.argmax(dim=1) == best_idx).cpu().numpy()
        ece_before = self._ece(probs_before.max(dim=1).values.cpu().numpy(), correct)
        ece_after = self._ece(probs_after.max(dim=1).values.cpu().numpy(), correct)

        logger.info(
            "Temperature calibration: T=%.3f, ECE %.3f -> %.3f",
            optimal_temp,
            ece_before,
            ece_after,
        )

        return optimal_temp


# Backwards compatibility alias
GNN_prediction = GNNPredictor
