#!/usr/bin/env python3
"""Evaluate ensemble with different aggregation methods.

Compares:
1. Naive probability averaging (baseline)
2. Shrinkage-weighted logit aggregation (gradient boosting)

Usage:
    python eval_ensemble.py --config checkpoints/ensemble/ensemble_config.json
"""

import argparse
import json
import logging
from pathlib import Path

import torch
import yaml

from embergraphrouter.data_processing.construct_router_data import load_data
from embergraphrouter.train_ensemble import EnsemblePredictor

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


def compute_metrics(predictions: torch.Tensor, rewards: torch.Tensor) -> dict:
    """Compute routing metrics from predictions and rewards.

    Args:
        predictions: Model predictions (N, num_llms) - probabilities
        rewards: Ground truth rewards (N, num_llms)

    Returns:
        Dict with top1_acc, regret, frac_of_oracle, etc.
    """
    pred_rank = torch.argsort(predictions, 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()

    # Handle ties correctly: 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_shifted),
        chosen_shifted / oracle_shifted.clamp_min(1e-6),
    )
    frac_of_oracle = per_query_frac.mean().item()

    # NDCG@k
    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()

    return {
        "top1_acc": top1_acc,
        "regret": regret,
        "frac_of_oracle": frac_of_oracle,
        "ndcg@1": ndcg_at_k(1),
        "ndcg@3": ndcg_at_k(3),
        "ndcg@All": ndcg_at_k(rewards.size(1)),
        "avg_chosen_reward": chosen_reward.mean().item(),
        "avg_oracle_reward": max_reward.mean().item(),
    }


def evaluate_ensemble(
    ensemble_config_path: str,
    data_config_path: str,
    device: str = "cuda",
) -> dict:
    """Evaluate ensemble with both aggregation methods.

    Args:
        ensemble_config_path: Path to ensemble_config.json
        data_config_path: Path to training config YAML (for data paths)
        device: Device for inference

    Returns:
        Dict with results for both methods
    """
    # Load data config
    with open(data_config_path, "r") as f:
        config = yaml.safe_load(f)

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

    # Load test data
    logger.info("Loading test data...")
    data = load_data(config, device=device)
    test_data = data["test"]

    # Load ensemble
    logger.info("Loading ensemble from %s", ensemble_config_path)
    ensemble = EnsemblePredictor(ensemble_config_path, device=device)

    logger.info("Ensemble type: %s", "boosted" if ensemble.is_boosted else "deep")
    if ensemble.is_boosted:
        logger.info("Shrinkage rate: %.2f", ensemble.shrinkage)

    # Compute rewards from test data
    # Use quality-only rewards for consistency (wp=1, wc=0, wl=0)
    combined_edge = test_data.combined_edge
    num_llms = 7  # Assuming 7 LLMs
    mask = test_data.edge_can_see.bool()

    # Extract quality scores (perf column = 1)
    quality = combined_edge[mask, 1]
    rewards = quality.reshape(-1, num_llms)

    results = {}

    # Method 1: Naive probability averaging
    logger.info("\n" + "=" * 60)
    logger.info("Method 1: Naive Probability Averaging")
    logger.info("=" * 60)

    result_naive = ensemble.predict(test_data, use_shrinkage=False)
    preds_naive = result_naive["predictions"]
    preds_naive = preds_naive.reshape(-1, num_llms)

    metrics_naive = compute_metrics(preds_naive, rewards)
    results["naive_averaging"] = metrics_naive

    logger.info("Results (Naive Averaging):")
    for k, v in metrics_naive.items():
        logger.info("  %s: %.4f", k, v)

    # Method 2: Shrinkage-weighted aggregation (only for boosted)
    if ensemble.is_boosted:
        logger.info("\n" + "=" * 60)
        logger.info("Method 2: Shrinkage-Weighted Logit Aggregation (η=%.2f)", ensemble.shrinkage)
        logger.info("=" * 60)

        result_shrinkage = ensemble.predict(test_data, use_shrinkage=True)
        preds_shrinkage = result_shrinkage["predictions"]
        preds_shrinkage = preds_shrinkage.reshape(-1, num_llms)

        metrics_shrinkage = compute_metrics(preds_shrinkage, rewards)
        results["shrinkage_weighted"] = metrics_shrinkage

        logger.info("Results (Shrinkage-Weighted):")
        for k, v in metrics_shrinkage.items():
            logger.info("  %s: %.4f", k, v)

        # Comparison
        logger.info("\n" + "=" * 60)
        logger.info("Comparison: Shrinkage vs Naive")
        logger.info("=" * 60)

        for metric in ["frac_of_oracle", "regret", "top1_acc", "ndcg@3"]:
            naive_val = metrics_naive[metric]
            shrinkage_val = metrics_shrinkage[metric]
            diff = shrinkage_val - naive_val
            diff_pct = 100 * diff / (abs(naive_val) + 1e-8)
            better = "↑" if diff > 0 else "↓" if diff < 0 else "="

            # For regret, lower is better
            if metric == "regret":
                better = "↓" if diff < 0 else "↑" if diff > 0 else "="

            logger.info(
                "  %s: %.4f → %.4f (%+.4f, %+.1f%%) %s",
                metric, naive_val, shrinkage_val, diff, diff_pct, better
            )

    return results


def main():
    parser = argparse.ArgumentParser(description="Evaluate ensemble aggregation methods")
    parser.add_argument(
        "--config",
        type=str,
        required=True,
        help="Path to ensemble_config.json",
    )
    parser.add_argument(
        "--data-config",
        type=str,
        default=None,
        help="Path to data config YAML (uses base_config from ensemble if not specified)",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device for inference",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Path to save results JSON",
    )

    args = parser.parse_args()

    # Get data config from ensemble config if not specified
    if args.data_config is None:
        with open(args.config, "r") as f:
            ensemble_config = json.load(f)
        args.data_config = ensemble_config.get("base_config")
        if args.data_config is None:
            raise ValueError(
                "No --data-config specified and ensemble_config.json has no base_config"
            )
        # Make relative to ensemble config dir
        if not Path(args.data_config).is_absolute():
            args.data_config = str(Path(args.config).parent.parent / args.data_config)

    logger.info("Ensemble config: %s", args.config)
    logger.info("Data config: %s", args.data_config)
    logger.info("Device: %s", args.device)

    results = evaluate_ensemble(args.config, args.data_config, args.device)

    if args.output:
        with open(args.output, "w") as f:
            json.dump(results, f, indent=2)
        logger.info("\nResults saved to %s", args.output)

    return results


if __name__ == "__main__":
    main()
