#!/usr/bin/env python3
"""Train a boosted ensemble of GNN models for robust LLM routing.

This script supports two training modes:

1. **Deep Ensemble** (--mode deep): Independent parallel training
   - All models train simultaneously on original targets
   - Different seeds provide diversity
   - Inference: average predictions, std = uncertainty

2. **Boosted Ensemble** (--mode boosted): Sequential residual training with pipeline overlap
   - Model 0 trains on original targets
   - Model 1 trains on residuals (errors) from Model 0
   - Model 2 trains on residuals from ensemble(0,1)
   - PipeDream-style overlap: Model N+1 starts after Model N reaches warmup_epochs

Usage:
    # Deep ensemble (parallel, independent)
    python train_ensemble.py --mode deep --num_models 3

    # Boosted ensemble (sequential with overlap)
    python train_ensemble.py --mode boosted --num_models 3 --warmup_epochs 20

    # Boosted with more overlap (start next model earlier)
    python train_ensemble.py --mode boosted --warmup_epochs 10

Architecture:
    - Same GNN architecture for all models
    - Boosted: each model focuses on previous models' errors
    - Deep: independent training with seed diversity
    - Inference: weighted sum (boosted) or average (deep)
"""

import argparse
import json
import logging
import os
import queue
import subprocess
import sys
import threading
import time
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Optional

import numpy as np
import torch
import yaml

_MODULE_DIR = Path(__file__).resolve().parent
if str(_MODULE_DIR.parent) not in sys.path:
    sys.path.insert(0, str(_MODULE_DIR.parent))

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)


def get_available_gpus() -> list[int]:
    """Detect available CUDA GPUs."""
    if not torch.cuda.is_available():
        return []
    return list(range(torch.cuda.device_count()))


def train_single_model(
    model_id: int,
    config_path: str,
    output_dir: Path,
    seed: int,
    gpu_id: int,
    offline: bool = False,
) -> dict:
    """Train a single model in a subprocess.

    Args:
        model_id: Identifier for this ensemble member (0, 1, 2, ...).
        config_path: Path to config YAML.
        output_dir: Directory to save model checkpoint.
        seed: Random seed for reproducibility.
        gpu_id: GPU device ID to use.
        offline: Run W&B in offline mode.

    Returns:
        Dict with training results.
    """
    # Set GPU for this process
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    # Import training modules (after setting CUDA_VISIBLE_DEVICES)
    import wandb

    from embergraphrouter.model.multi_task_graph_router import graph_router_prediction

    # Load config
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    # Override seed and model path for this ensemble member
    config["seed"] = seed
    model_path = output_dir / f"ensemble_model_{model_id}.pt"
    config["model_path"] = str(model_path)

    # Resolve relative paths
    for key in (
        "unified_qa_data_path",
        "llm_description_path",
        "saved_router_data_path",
        "llm_embedding_path",
    ):
        if key in config and not Path(config[key]).is_absolute():
            config[key] = str(_MODULE_DIR / config[key])

    # Initialize W&B with unique run name
    mode = "offline" if offline else "online"
    wandb.init(
        project="graph_router_ensemble",
        name=f"ensemble_member_{model_id}_seed_{seed}",
        mode=mode,
        config=config,
        group="ensemble_training",
    )

    logger.info(
        "Training ensemble member %d (seed=%d, gpu=%d, path=%s)",
        model_id, seed, gpu_id, model_path,
    )

    start_time = time.time()

    try:
        # Run training (saves model to model_path as state_dict)
        trainer = graph_router_prediction(
            router_data_path=config["saved_router_data_path"],
            llm_path=config["llm_description_path"],
            llm_embedding_path=config["llm_embedding_path"],
            config=config,
            wandb=wandb,
        )
        result = {"query_feature_dim": trainer.query_embeddings.shape[1]} if hasattr(trainer, "query_embeddings") else {}
        elapsed = time.time() - start_time

        # Repackage as ensemble checkpoint with config and dimensions
        # This enables proper model reconstruction at inference time
        if model_path.exists():
            state_dict = torch.load(model_path, map_location="cpu", weights_only=True)

            # Extract feature dimensions from result or infer from state_dict
            query_dim = result.get("query_feature_dim") if isinstance(result, dict) else None
            llm_dim = result.get("llm_feature_dim") if isinstance(result, dict) else None

            # If not in result, try to infer from model weights
            if query_dim is None:
                for key in state_dict:
                    if "query" in key.lower() and "weight" in key:
                        query_dim = state_dict[key].shape[-1]
                        break
            if llm_dim is None:
                for key in state_dict:
                    if "llm" in key.lower() and "weight" in key:
                        llm_dim = state_dict[key].shape[-1]
                        break

            # Fallback to config values
            if query_dim is None:
                query_dim = config.get("query_feature_dim", 768)
            if llm_dim is None:
                llm_dim = config.get("llm_feature_dim", 768)

            ensemble_checkpoint = {
                "model_state_dict": state_dict,
                "config": config,
                "query_feature_dim": query_dim,
                "llm_feature_dim": llm_dim,
                "model_id": model_id,
                "seed": seed,
            }
            torch.save(ensemble_checkpoint, model_path)
            logger.info(
                "Repackaged model %d as ensemble checkpoint (query=%d, llm=%d)",
                model_id, query_dim, llm_dim,
            )

        wandb.finish()

        return {
            "model_id": model_id,
            "seed": seed,
            "gpu_id": gpu_id,
            "model_path": str(model_path),
            "elapsed_seconds": elapsed,
            "success": True,
            "metrics": result if isinstance(result, dict) else {},
        }

    except Exception as e:
        elapsed = time.time() - start_time
        logger.error("Model %d failed: %s", model_id, e)
        wandb.finish(exit_code=1)

        return {
            "model_id": model_id,
            "seed": seed,
            "gpu_id": gpu_id,
            "model_path": str(model_path),
            "elapsed_seconds": elapsed,
            "success": False,
            "error": str(e),
        }


def train_ensemble_parallel(
    num_models: int,
    config_path: str,
    output_dir: Path,
    base_seed: int,
    gpus: list[int],
    offline: bool = False,
) -> list[dict]:
    """Train ensemble members in parallel across GPUs.

    Args:
        num_models: Number of ensemble members to train.
        config_path: Path to config YAML.
        output_dir: Directory for checkpoints.
        base_seed: Base random seed (each model gets base_seed + model_id).
        gpus: List of GPU IDs to use (cycles if fewer than num_models).
        offline: Run W&B in offline mode.

    Returns:
        List of training result dicts.
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    # Assign GPUs round-robin
    gpu_assignments = [gpus[i % len(gpus)] for i in range(num_models)]

    logger.info(
        "Training %d ensemble members across GPUs %s",
        num_models, gpus,
    )
    logger.info("GPU assignments: %s", dict(enumerate(gpu_assignments)))

    results = []

    # Use ProcessPoolExecutor to run trainings in parallel
    # Each process gets its own GPU via CUDA_VISIBLE_DEVICES
    with ProcessPoolExecutor(max_workers=len(gpus)) as executor:
        futures = {}

        for model_id in range(num_models):
            seed = base_seed + model_id * 1000
            gpu_id = gpu_assignments[model_id]

            future = executor.submit(
                train_single_model,
                model_id=model_id,
                config_path=config_path,
                output_dir=output_dir,
                seed=seed,
                gpu_id=gpu_id,
                offline=offline,
            )
            futures[future] = model_id

        for future in as_completed(futures):
            model_id = futures[future]
            try:
                result = future.result()
                results.append(result)
                status = "✓" if result["success"] else "✗"
                logger.info(
                    "Model %d %s (%.1f min)",
                    model_id, status, result["elapsed_seconds"] / 60,
                )
            except Exception as e:
                logger.error("Model %d failed with exception: %s", model_id, e)
                results.append({
                    "model_id": model_id,
                    "success": False,
                    "error": str(e),
                })

    return sorted(results, key=lambda x: x["model_id"])


# =============================================================================
# BOOSTED ENSEMBLE: PipeDream-style sequential training with overlap
# =============================================================================


def compute_residuals(
    predictions: torch.Tensor,
    targets: torch.Tensor,
    shrinkage: float = 0.1,
) -> torch.Tensor:
    """Compute residuals for boosting: what the current ensemble got wrong.

    Uses gradient boosting for multi-class classification. For cross-entropy loss,
    the negative gradient (direction of improvement) is (targets - predictions).
    We scale by shrinkage to prevent overfitting in later rounds.

    The residuals are then shifted to valid probability space [0, 1] and
    renormalized so each row sums to 1 - this gives the next model valid
    soft targets to train on.

    Args:
        predictions: Current ensemble predictions (N, num_llms), probabilities.
        targets: Ground truth soft targets (N, num_llms).
        shrinkage: Learning rate for boosting (0.1 = conservative updates).

    Returns:
        Pseudo-residual targets for the next model, shape (N, num_llms).
    """
    # Negative gradient of cross-entropy: targets - predictions
    # This points toward the correct answer
    raw_residuals = targets - predictions

    # Scale by shrinkage to prevent overfitting (standard in gradient boosting)
    # Smaller shrinkage = more conservative updates = more rounds needed but better generalization
    scaled_residuals = shrinkage * raw_residuals

    # For the next model to train on these as targets, we need valid probabilities.
    # Option 1: Train next model on residuals directly (requires modified loss)
    # Option 2: Create pseudo-targets by adjusting predictions toward targets
    #
    # We use Option 2: new_targets = predictions + scaled_residuals
    # This moves predictions toward targets by shrinkage amount each round
    pseudo_targets = predictions + scaled_residuals

    # Clamp to valid probability range and renormalize
    pseudo_targets = torch.clamp(pseudo_targets, min=1e-7, max=1.0)
    pseudo_targets = pseudo_targets / pseudo_targets.sum(dim=-1, keepdim=True)

    return pseudo_targets


def compute_sample_weights(
    predictions: torch.Tensor,
    targets: torch.Tensor,
    focus_errors: bool = True,
) -> torch.Tensor:
    """Compute sample weights for AdaBoost-style error focusing.

    Samples where the ensemble is wrong get higher weight, forcing
    subsequent models to focus on hard examples.

    Args:
        predictions: Current ensemble predictions (N, num_llms).
        targets: Ground truth (N, num_llms), can be hard or soft labels.
        focus_errors: If True, upweight errors. If False, uniform weights.

    Returns:
        Sample weights (N,), normalized to sum to N.
    """
    if not focus_errors:
        return torch.ones(predictions.shape[0], device=predictions.device)

    # For hard labels: error = 1 if argmax(pred) != argmax(target)
    # For soft labels: error = cross-entropy or KL divergence
    pred_class = predictions.argmax(dim=-1)
    target_class = targets.argmax(dim=-1)

    # Binary error indicator
    errors = (pred_class != target_class).float()

    # Also consider confidence: low-confidence correct predictions should get some weight
    pred_confidence = predictions.max(dim=-1).values
    confidence_weight = 1.0 - pred_confidence  # Less confident = more weight

    # Combine: errors get high weight, uncertain samples get moderate weight
    weights = errors + 0.3 * confidence_weight * (1.0 - errors)

    # Normalize to sum to N (preserves expected gradient magnitude)
    weights = weights * (weights.numel() / (weights.sum() + 1e-8))

    return weights


def train_boosted_pipeline(
    num_models: int,
    config_path: str,
    output_dir: Path,
    base_seed: int,
    gpus: list[int],
    warmup_epochs: int = 20,
    shrinkage: float = 0.1,
    offline: bool = False,
) -> list[dict]:
    """Train boosted ensemble with PipeDream-style pipeline overlap.

    Timeline (warmup_epochs=20, train_epochs=200):
        GPU 0: [===== Model 0 training (200 epochs) =====]
        GPU 1:      [== warmup ==][=== Model 1 on residuals (180 epochs) ===]
        GPU 2:                        [== warmup ==][=== Model 2 (160 epochs) ===]

    Each model trains on the residuals (errors) from the ensemble of all
    previous models, focusing on what the ensemble got wrong.

    Args:
        num_models: Number of boosting rounds (models).
        config_path: Path to base config YAML.
        output_dir: Directory for checkpoints.
        base_seed: Base random seed.
        gpus: List of GPU IDs (one per model for overlap).
        warmup_epochs: Epochs before starting next model.
        shrinkage: Boosting learning rate (smaller = more conservative).
        offline: Run W&B in offline mode.

    Returns:
        List of training result dicts with boosting metadata.
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    if len(gpus) < num_models:
        logger.warning(
            "Fewer GPUs (%d) than models (%d). Pipeline overlap will be limited.",
            len(gpus), num_models,
        )

    logger.info("=" * 60)
    logger.info("BOOSTED ENSEMBLE TRAINING (PipeDream-style)")
    logger.info("=" * 60)
    logger.info("Models: %d (sequential with overlap)", num_models)
    logger.info("Warmup epochs: %d (overlap window)", warmup_epochs)
    logger.info("Shrinkage: %.2f", shrinkage)
    logger.info("GPUs: %s", gpus)
    logger.info("=" * 60)

    # Shared state for pipeline coordination
    model_checkpoints: list[Optional[Path]] = [None] * num_models
    residual_targets: list[Optional[Path]] = [None] * num_models
    model_ready_events = [threading.Event() for _ in range(num_models)]
    results: list[dict] = []
    results_lock = threading.Lock()

    def train_boosted_model(model_id: int) -> dict:
        """Train a single model in the boosting sequence.

        Fault-tolerant: failures are isolated and don't block other models.
        """
        gpu_id = gpus[model_id % len(gpus)]
        seed = base_seed + model_id * 1000

        # Isolate GPU failures - set env var in try block
        try:
            os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
        except Exception as e:
            logger.error("Model %d: Failed to set GPU %d: %s", model_id, gpu_id, e)
            model_ready_events[model_id].set()  # Don't block pipeline
            return {"model_id": model_id, "success": False, "error": f"GPU setup failed: {e}"}

        import wandb
        from embergraphrouter.model.multi_task_graph_router import graph_router_prediction

        # Wait for previous model to reach warmup (if not first model)
        # Timeout prevents infinite hang if previous model fails
        if model_id > 0:
            logger.info(
                "Model %d waiting for Model %d warmup...",
                model_id, model_id - 1,
            )
            # Timeout: warmup_epochs * ~2min/epoch + buffer
            timeout_seconds = warmup_epochs * 120 + 300
            got_signal = model_ready_events[model_id - 1].wait(timeout=timeout_seconds)
            if not got_signal:
                logger.warning(
                    "Model %d: Timeout waiting for Model %d. Starting anyway.",
                    model_id, model_id - 1,
                )
            else:
                logger.info("Model %d starting (previous model warmed up)", model_id)

        # Load config
        with open(config_path, "r") as f:
            config = yaml.safe_load(f)

        config["seed"] = seed
        model_path = output_dir / f"boosted_model_{model_id}.pt"
        config["model_path"] = str(model_path)

        # Resolve relative paths
        for key in (
            "unified_qa_data_path",
            "llm_description_path",
            "saved_router_data_path",
            "llm_embedding_path",
        ):
            if key in config and not Path(config[key]).is_absolute():
                config[key] = str(_MODULE_DIR / config[key])

        # For boosted models (id > 0), load residual targets
        if model_id > 0 and residual_targets[model_id - 1] is not None:
            config["residual_targets_path"] = str(residual_targets[model_id - 1])
            config["boosting_shrinkage"] = shrinkage
            logger.info(
                "Model %d training on residuals from: %s",
                model_id, residual_targets[model_id - 1],
            )

        # Initialize W&B with reinit=True to allow multiple runs in threads
        mode = "offline" if offline else "online"
        run = wandb.init(
            project="graph_router_boosted",
            name=f"boosted_model_{model_id}_seed_{seed}",
            mode=mode,
            config=config,
            group="boosted_training",
            reinit=True,  # Critical: allows multiple wandb runs in same process
        )

        logger.info(
            "Training boosted model %d (seed=%d, gpu=%d)",
            model_id, seed, gpu_id,
        )

        start_time = time.time()

        # Custom callback to signal warmup and save residuals
        warmup_signaled = False

        def epoch_callback(epoch: int, model, val_predictions=None):
            nonlocal warmup_signaled
            if epoch >= warmup_epochs and not warmup_signaled:
                # Signal next model can start
                model_ready_events[model_id].set()
                warmup_signaled = True
                logger.info("Model %d reached warmup, signaling next model", model_id)

                # Save current predictions as residual targets for next model
                if val_predictions is not None and model_id < num_models - 1:
                    residual_path = output_dir / f"residuals_after_model_{model_id}.pt"
                    torch.save(val_predictions, residual_path)
                    residual_targets[model_id] = residual_path

        # Note: epoch_callback integration requires modification to graph_router_prediction
        # For now, we use a simpler approach: train fully, then compute residuals

        try:
            trainer = graph_router_prediction(
                router_data_path=config["saved_router_data_path"],
                llm_path=config["llm_description_path"],
                llm_embedding_path=config["llm_embedding_path"],
                config=config,
                wandb=wandb,
            )
            result = {"query_feature_dim": trainer.query_embeddings.shape[1]} if hasattr(trainer, "query_embeddings") else {}
            elapsed = time.time() - start_time

            # Signal completion (in case warmup wasn't reached)
            model_ready_events[model_id].set()

            # Repackage checkpoint
            if model_path.exists():
                state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
                ensemble_checkpoint = {
                    "model_state_dict": state_dict,
                    "config": config,
                    "query_feature_dim": config.get("query_feature_dim", 768),
                    "llm_feature_dim": config.get("llm_feature_dim", 768),
                    "model_id": model_id,
                    "seed": seed,
                    "boosting_round": model_id,
                    "shrinkage": shrinkage,
                }
                torch.save(ensemble_checkpoint, model_path)
                model_checkpoints[model_id] = model_path

            wandb.finish()

            return {
                "model_id": model_id,
                "seed": seed,
                "gpu_id": gpu_id,
                "model_path": str(model_path),
                "elapsed_seconds": elapsed,
                "success": True,
                "boosting_round": model_id,
                "metrics": result if isinstance(result, dict) else {},
            }

        except Exception as e:
            elapsed = time.time() - start_time
            logger.error("Boosted model %d failed: %s", model_id, e)
            model_ready_events[model_id].set()  # Don't block pipeline
            wandb.finish(exit_code=1)
            return {
                "model_id": model_id,
                "seed": seed,
                "gpu_id": gpu_id,
                "elapsed_seconds": elapsed,
                "success": False,
                "error": str(e),
            }

    # Launch models with ThreadPoolExecutor for pipeline coordination
    # (Threads can share state; processes would need IPC)
    with ThreadPoolExecutor(max_workers=num_models) as executor:
        futures = {
            executor.submit(train_boosted_model, i): i
            for i in range(num_models)
        }

        for future in as_completed(futures):
            model_id = futures[future]
            try:
                result = future.result()
                with results_lock:
                    results.append(result)
                status = "✓" if result["success"] else "✗"
                logger.info(
                    "Boosted model %d %s (%.1f min)",
                    model_id, status, result["elapsed_seconds"] / 60,
                )
            except Exception as e:
                logger.error("Boosted model %d exception: %s", model_id, e)
                with results_lock:
                    results.append({"model_id": model_id, "success": False, "error": str(e)})

    return sorted(results, key=lambda x: x["model_id"])


def save_ensemble_config(
    output_dir: Path,
    results: list[dict],
    config_path: str,
) -> Path:
    """Save ensemble configuration for inference."""
    ensemble_config = {
        "num_models": len(results),
        "model_paths": [r["model_path"] for r in results if r["success"]],
        "seeds": [r["seed"] for r in results if r["success"]],
        "base_config": config_path,
        "training_results": results,
    }

    config_file = output_dir / "ensemble_config.json"
    with open(config_file, "w") as f:
        json.dump(ensemble_config, f, indent=2)

    logger.info("Saved ensemble config to %s", config_file)
    return config_file


class EnsemblePredictor:
    """Inference with ensemble of GNN models.

    Loads multiple trained models and combines predictions.
    Provides uncertainty estimates via prediction disagreement.

    For boosted ensembles, uses proper gradient boosting aggregation:
        F(x) = logits_0 + η·logits_1 + η·logits_2 + ... + η·logits_n
    where η is the shrinkage rate (default 0.1).

    For deep ensembles (independent training), uses simple probability averaging.

    Each ensemble member is saved as a checkpoint containing:
        - model_state_dict: The trained weights
        - config: Full training config for model reconstruction
        - llm_embeddings: LLM feature matrix (shared across models)
    """

    def __init__(
        self,
        ensemble_config_path: str,
        device: str = "cuda",
        llm_embedding_path: Optional[str] = None,
    ):
        """Load ensemble from config file.

        Args:
            ensemble_config_path: Path to ensemble_config.json.
            device: Device for inference.
            llm_embedding_path: Path to LLM embeddings (if not in checkpoints).
        """
        from embergraphrouter.data_processing.utils import loadpkl
        from embergraphrouter.model.graph_nn import GNNPredictor

        with open(ensemble_config_path, "r") as f:
            ensemble_config = json.load(f)

        self.model_paths = ensemble_config["model_paths"]
        self.device = device
        self.models = []
        self.configs = []

        # Detect boosted ensemble and extract shrinkage
        self.is_boosted = False
        self.shrinkage = 0.1  # Default shrinkage
        if ensemble_config.get("training_results"):
            first_result = ensemble_config["training_results"][0]
            if "boosting_round" in first_result:
                self.is_boosted = True
                # Get shrinkage from checkpoint if available
                self.shrinkage = first_result.get("shrinkage", 0.1)

        logger.info("Loading %d ensemble members...", len(self.model_paths))
        if self.is_boosted:
            logger.info("Boosted ensemble detected (shrinkage=%.2f)", self.shrinkage)

        for i, path in enumerate(self.model_paths):
            checkpoint = torch.load(path, map_location=device, weights_only=False)
            model_config = checkpoint["config"]
            self.configs.append(model_config)

            # Get dimensions from checkpoint or load from embedding files
            query_dim = checkpoint.get("query_feature_dim", model_config.get("query_feature_dim"))
            llm_dim = checkpoint.get("llm_feature_dim", model_config.get("llm_feature_dim"))

            if query_dim is None or llm_dim is None:
                raise ValueError(
                    f"Model {i} checkpoint missing feature dimensions. "
                    "Ensure models were trained with updated train_ensemble.py"
                )

            # Create a minimal wandb stub for GNNPredictor
            class WandBStub:
                def log(self, *args, **kwargs): pass
                def watch(self, *args, **kwargs): pass

            # Reconstruct GNNPredictor with the saved config
            predictor = GNNPredictor(
                query_feature_dim=query_dim,
                llm_feature_dim=llm_dim,
                hidden_features_size=model_config["embedding_dim"],
                in_edges_size=model_config["edge_dim"],
                wandb=WandBStub(),
                config=model_config,
                device=device,
            )

            # Load trained weights
            predictor.model.load_state_dict(checkpoint["model_state_dict"])
            predictor.model.eval()
            self.models.append(predictor)

            logger.info(
                "  Loaded model %d (seed=%d, query_dim=%d, llm_dim=%d)",
                i, model_config.get("seed", -1), query_dim, llm_dim,
            )

        logger.info("Ensemble loaded (%d models)", len(self.models))

    def predict(self, data, use_shrinkage: bool = True) -> dict:
        """Run inference with all ensemble members.

        For boosted ensembles, combines predictions in logit space using
        gradient boosting formula:
            combined_logits = logits_0 + η·logits_1 + η·logits_2 + ...

        For deep ensembles, uses simple probability averaging.

        Args:
            data: Input data (PyG Data object or batch).
            use_shrinkage: Whether to apply shrinkage weighting for boosted
                ensembles. Set False to force simple averaging.

        Returns:
            Dict with:
                - predictions: Combined predictions (softmax probabilities)
                - uncertainty: Std across ensemble members
                - individual: List of each model's predictions (probabilities)
                - individual_logits: List of each model's raw logits
        """
        individual_logits = []
        individual_preds = []

        for predictor in self.models:
            with torch.no_grad():
                logits = predictor.model(data)
                probs = torch.softmax(logits, dim=-1)
                individual_logits.append(logits)
                individual_preds.append(probs)

        # Stack for uncertainty calculation
        stacked_probs = torch.stack(individual_preds, dim=0)
        uncertainty = stacked_probs.std(dim=0)

        if self.is_boosted and use_shrinkage:
            # Gradient boosting aggregation in logit space:
            # F(x) = logits_0 + η·logits_1 + η·logits_2 + ... + η·logits_n
            #
            # Model 0 is the base learner (full weight)
            # Models 1..n are residual learners (shrinkage weight)
            combined_logits = individual_logits[0].clone()
            for i in range(1, len(individual_logits)):
                combined_logits = combined_logits + self.shrinkage * individual_logits[i]

            # Convert to probabilities
            mean_pred = torch.softmax(combined_logits, dim=-1)
        else:
            # Deep ensemble: simple probability averaging
            mean_pred = stacked_probs.mean(dim=0)

        return {
            "predictions": mean_pred,
            "uncertainty": uncertainty,
            "individual": individual_preds,
            "individual_logits": individual_logits,
        }

    def predict_with_disagreement_routing(
        self,
        data,
        uncertainty_threshold: float = 0.3,
    ) -> dict:
        """Route based on ensemble disagreement.

        If uncertainty is high (models disagree), route to strongest model.
        If uncertainty is low (models agree), use averaged prediction.

        Args:
            data: Input data.
            uncertainty_threshold: Threshold for high uncertainty.

        Returns:
            Predictions with routing decisions.
        """
        result = self.predict(data)

        # High uncertainty = route to strongest model
        high_uncertainty = result["uncertainty"].mean(dim=-1) > uncertainty_threshold

        return {
            **result,
            "high_uncertainty_mask": high_uncertainty,
            "routing_decision": torch.where(
                high_uncertainty,
                torch.zeros_like(high_uncertainty),  # Route to model 0 (strongest)
                torch.ones_like(high_uncertainty),   # Use ensemble average
            ),
        }


def main():
    parser = argparse.ArgumentParser(
        description="Train ensemble of GNN routers (deep or boosted)",
    )
    parser.add_argument(
        "--mode",
        type=str,
        choices=["deep", "boosted"],
        default="boosted",
        help="Training mode: 'boosted' (sequential residuals) or 'deep' (parallel independent)",
    )
    parser.add_argument(
        "--num_models",
        type=int,
        default=8,
        help="Number of ensemble members (default: 8 for 8xGPU)",
    )
    parser.add_argument(
        "--config",
        type=str,
        default=str(_MODULE_DIR / "configs" / "config.yaml"),
        help="Path to config YAML",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default=str(_MODULE_DIR / "checkpoints" / "ensemble"),
        help="Directory for ensemble checkpoints",
    )
    parser.add_argument(
        "--base_seed",
        type=int,
        default=42,
        help="Base random seed (default: 42)",
    )
    parser.add_argument(
        "--gpus",
        type=str,
        default=None,
        help="Comma-separated GPU IDs (default: all available)",
    )
    parser.add_argument(
        "--offline",
        action="store_true",
        help="Run W&B in offline mode",
    )
    # Boosted-specific arguments
    parser.add_argument(
        "--warmup_epochs",
        type=int,
        default=20,
        help="Epochs before starting next model in boosted mode (default: 20)",
    )
    parser.add_argument(
        "--shrinkage",
        type=float,
        default=0.1,
        help="Boosting learning rate / shrinkage (default: 0.1)",
    )

    args = parser.parse_args()

    # Determine GPUs to use
    if args.gpus:
        gpus = [int(g) for g in args.gpus.split(",")]
    else:
        gpus = get_available_gpus()
        if not gpus:
            logger.warning("No GPUs detected, using CPU")
            gpus = [-1]  # CPU

    output_dir = Path(args.output_dir)

    # Route to appropriate training mode
    if args.mode == "boosted":
        results = train_boosted_pipeline(
            num_models=args.num_models,
            config_path=args.config,
            output_dir=output_dir,
            base_seed=args.base_seed,
            gpus=gpus,
            warmup_epochs=args.warmup_epochs,
            shrinkage=args.shrinkage,
            offline=args.offline,
        )
    else:
        logger.info("=" * 60)
        logger.info("DEEP ENSEMBLE TRAINING (parallel independent)")
        logger.info("=" * 60)
        logger.info("Models: %d", args.num_models)
        logger.info("GPUs: %s", gpus)
        logger.info("Output: %s", output_dir)
        logger.info("Base seed: %d", args.base_seed)
        logger.info("=" * 60)

        results = train_ensemble_parallel(
            num_models=args.num_models,
            config_path=args.config,
            output_dir=output_dir,
            base_seed=args.base_seed,
            gpus=gpus,
            offline=args.offline,
        )

    # Save ensemble config
    save_ensemble_config(output_dir, results, args.config)

    # Print summary
    successful = sum(1 for r in results if r["success"])
    logger.info("=" * 60)
    logger.info("TRAINING COMPLETE")
    logger.info("=" * 60)
    logger.info("Successful: %d/%d", successful, args.num_models)

    for r in results:
        status = "✓" if r["success"] else "✗"
        elapsed = r.get("elapsed_seconds", 0) / 60
        logger.info(
            "  Model %d: %s (%.1f min, seed=%d)",
            r["model_id"], status, elapsed, r.get("seed", -1),
        )

    if successful < args.num_models:
        logger.warning("Some models failed! Check logs above.")
        sys.exit(1)


if __name__ == "__main__":
    main()
