"""Evaluation harness for router baseline comparison.

This module provides comprehensive evaluation utilities for comparing router
architectures across multiple dimensions:

1. **Routing Quality**: frac_oracle, top-1 accuracy, NDCG
2. **Calibration**: ECE, reliability diagrams
3. **Efficiency**: Inference latency, parameter count, memory usage
4. **Robustness**: Performance across query types, OOD generalization

Usage:
    from embergraphrouter.model.baselines.evaluation import BaselineEvaluator

    evaluator = BaselineEvaluator(test_dataset)
    results = evaluator.compare([
        ("Two Tower", two_tower_router),
        ("Multi-Head MLP", mlp_router),
        ("GNN", gnn_router),
    ])
    evaluator.print_comparison_table(results)
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Optional

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader

from .base import BaseRouter, DifficultySignals
from .trainer import RouterDataset, collate_router_batch


@dataclass
class EvaluationResult:
    """Complete evaluation results for a single router.

    Attributes:
        name: Router name for display.
        architecture: Architecture type (gnn, two_tower, multi_head_mlp).

        # Routing Quality
        frac_oracle: Fraction of oracle performance (0-1).
        top1_accuracy: Accuracy of top model selection.
        top3_accuracy: Accuracy of best model in top-3.
        ndcg_at_1: NDCG@1 (same as top-1 accuracy for binary).
        ndcg_at_3: NDCG@3 ranking quality.
        ndcg_at_5: NDCG@5 ranking quality.
        mean_reciprocal_rank: MRR for ranking.
        kendall_tau: Kendall's Tau rank correlation.

        # Calibration
        expected_calibration_error: ECE (lower is better).
        max_calibration_error: MCE (worst bin error).
        brier_score: Brier score for probability accuracy.

        # Efficiency
        parameter_count: Total trainable parameters.
        inference_latency_ms: Mean inference time per batch.
        memory_usage_mb: Peak GPU memory during inference.

        # Per-Difficulty Performance
        frac_oracle_easy: Frac oracle on easy queries (entropy < 0.3).
        frac_oracle_medium: Frac oracle on medium queries.
        frac_oracle_hard: Frac oracle on hard queries (entropy > 0.7).
    """

    name: str
    architecture: str

    # Routing Quality
    frac_oracle: float = 0.0
    top1_accuracy: float = 0.0
    top3_accuracy: float = 0.0
    ndcg_at_1: float = 0.0
    ndcg_at_3: float = 0.0
    ndcg_at_5: float = 0.0
    mean_reciprocal_rank: float = 0.0
    kendall_tau: float = 0.0

    # Calibration
    expected_calibration_error: float = 0.0
    max_calibration_error: float = 0.0
    brier_score: float = 0.0

    # Efficiency
    parameter_count: int = 0
    inference_latency_ms: float = 0.0
    memory_usage_mb: float = 0.0

    # Per-Difficulty Performance
    frac_oracle_easy: float = 0.0
    frac_oracle_medium: float = 0.0
    frac_oracle_hard: float = 0.0

    def to_dict(self) -> dict:
        """Convert to dictionary for serialization."""
        return {
            "name": self.name,
            "architecture": self.architecture,
            "frac_oracle": self.frac_oracle,
            "top1_accuracy": self.top1_accuracy,
            "top3_accuracy": self.top3_accuracy,
            "ndcg_at_3": self.ndcg_at_3,
            "mrr": self.mean_reciprocal_rank,
            "ece": self.expected_calibration_error,
            "parameter_count": self.parameter_count,
            "inference_latency_ms": self.inference_latency_ms,
        }


class BaselineEvaluator:
    """Comprehensive evaluator for router baselines."""

    def __init__(
        self,
        test_dataset: RouterDataset,
        batch_size: int = 32,
        device: Optional[torch.device] = None,
        num_latency_runs: int = 100,
    ) -> None:
        """Initialize evaluator.

        Args:
            test_dataset: Test dataset for evaluation.
            batch_size: Batch size for evaluation.
            device: Device for evaluation (default: auto-detect).
            num_latency_runs: Number of runs for latency measurement.
        """
        self.test_dataset = test_dataset
        self.batch_size = batch_size
        self.device = device or torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        self.num_latency_runs = num_latency_runs

        self.dataloader = DataLoader(
            test_dataset,
            batch_size=batch_size,
            shuffle=False,
            collate_fn=collate_router_batch,
        )

    def evaluate(
        self,
        router: BaseRouter,
        name: str = "Router",
    ) -> EvaluationResult:
        """Evaluate a single router.

        Args:
            router: Router to evaluate.
            name: Display name for results.

        Returns:
            EvaluationResult with comprehensive metrics.
        """
        router.to(self.device)
        router.eval()

        # Collect predictions
        all_scores = []
        all_targets = []
        all_signals = []

        with torch.no_grad():
            for query_embs, model_embs, targets in self.dataloader:
                query_embs = query_embs.to(self.device)
                model_embs = model_embs.to(self.device)
                targets = targets.to(self.device)

                output = router(query_embs, model_embs, return_signals=True)
                all_scores.append(output.scores.cpu())
                all_targets.append(targets.cpu())
                if output.signals:
                    all_signals.append(output.signals)

        scores = torch.cat(all_scores, dim=0)
        targets = torch.cat(all_targets, dim=0)

        # Compute metrics
        result = EvaluationResult(
            name=name,
            architecture=router.get_architecture_name(),
        )

        # Routing quality
        result.frac_oracle = self._compute_frac_oracle(scores, targets)
        result.top1_accuracy = self._compute_topk_accuracy(scores, targets, k=1)
        result.top3_accuracy = self._compute_topk_accuracy(scores, targets, k=3)
        result.ndcg_at_1 = self._compute_ndcg(scores, targets, k=1)
        result.ndcg_at_3 = self._compute_ndcg(scores, targets, k=3)
        result.ndcg_at_5 = self._compute_ndcg(scores, targets, k=5)
        result.mean_reciprocal_rank = self._compute_mrr(scores, targets)
        result.kendall_tau = self._compute_kendall_tau(scores, targets)

        # Calibration
        probs = F.softmax(scores, dim=-1)
        result.expected_calibration_error = self._compute_ece(probs, targets)
        result.max_calibration_error = self._compute_mce(probs, targets)
        result.brier_score = self._compute_brier(probs, targets)

        # Efficiency
        result.parameter_count = router.count_parameters()
        result.inference_latency_ms = self._measure_latency(router)
        result.memory_usage_mb = self._measure_memory(router)

        # Per-difficulty performance
        difficulty_results = self._compute_per_difficulty(scores, targets)
        result.frac_oracle_easy = difficulty_results["easy"]
        result.frac_oracle_medium = difficulty_results["medium"]
        result.frac_oracle_hard = difficulty_results["hard"]

        return result

    def compare(
        self,
        routers: list[tuple[str, BaseRouter]],
    ) -> list[EvaluationResult]:
        """Compare multiple routers.

        Args:
            routers: List of (name, router) tuples.

        Returns:
            List of EvaluationResult for each router.
        """
        results = []
        for name, router in routers:
            result = self.evaluate(router, name)
            results.append(result)
        return results

    def print_comparison_table(
        self,
        results: list[EvaluationResult],
        metrics: Optional[list[str]] = None,
    ) -> str:
        """Print formatted comparison table.

        Args:
            results: List of evaluation results.
            metrics: Metrics to include (default: key metrics).

        Returns:
            Formatted table string.
        """
        if metrics is None:
            metrics = [
                "frac_oracle",
                "top1_accuracy",
                "ndcg_at_3",
                "ece",
                "parameter_count",
                "latency_ms",
            ]

        # Build table
        headers = ["Router"] + [m.replace("_", " ").title() for m in metrics]

        rows = []
        for result in results:
            row = [result.name]
            for metric in metrics:
                if metric == "parameter_count":
                    value = f"{result.parameter_count:,}"
                elif metric == "latency_ms":
                    value = f"{result.inference_latency_ms:.2f}"
                elif metric == "ece":
                    value = f"{result.expected_calibration_error:.4f}"
                else:
                    value = f"{getattr(result, metric, 0):.4f}"
                row.append(value)
            rows.append(row)

        # Format table
        col_widths = [max(len(str(row[i])) for row in [headers] + rows)
                      for i in range(len(headers))]

        def format_row(row: list) -> str:
            return " | ".join(str(v).ljust(w) for v, w in zip(row, col_widths))

        separator = "-+-".join("-" * w for w in col_widths)

        lines = [
            format_row(headers),
            separator,
        ] + [format_row(row) for row in rows]

        table = "\n".join(lines)
        print(table)
        return table

    def _compute_frac_oracle(
        self,
        scores: torch.Tensor,
        targets: torch.Tensor,
    ) -> float:
        """Compute fraction of oracle performance."""
        pred_best = scores.argmax(dim=-1)
        oracle_perf = targets.max(dim=-1).values
        random_perf = targets.mean(dim=-1)
        selected_perf = targets.gather(1, pred_best.unsqueeze(1)).squeeze(1)

        denom = (oracle_perf - random_perf).clamp_min(1e-8)
        frac_oracle = ((selected_perf - random_perf) / denom).clamp(0, 1)
        return frac_oracle.mean().item()

    def _compute_topk_accuracy(
        self,
        scores: torch.Tensor,
        targets: torch.Tensor,
        k: int,
    ) -> float:
        """Compute top-k accuracy."""
        true_best = targets.argmax(dim=-1)
        _, top_k_pred = scores.topk(k, dim=-1)

        # Check if true best is in top-k predictions
        correct = (top_k_pred == true_best.unsqueeze(1)).any(dim=-1)
        return correct.float().mean().item()

    def _compute_ndcg(
        self,
        scores: torch.Tensor,
        targets: torch.Tensor,
        k: int,
    ) -> float:
        """Compute NDCG@k."""
        num_models = scores.shape[1]
        k = min(k, num_models)

        _, pred_order = scores.sort(dim=-1, descending=True)
        pred_order_k = pred_order[:, :k]

        relevance = targets.gather(1, pred_order_k)

        positions = torch.arange(1, k + 1, device=scores.device, dtype=torch.float)
        discounts = torch.log2(positions + 1)
        dcg = (relevance / discounts).sum(dim=-1)

        ideal_relevance, _ = targets.sort(dim=-1, descending=True)
        ideal_relevance_k = ideal_relevance[:, :k]
        idcg = (ideal_relevance_k / discounts).sum(dim=-1)

        ndcg = dcg / idcg.clamp_min(1e-8)
        return ndcg.mean().item()

    def _compute_mrr(
        self,
        scores: torch.Tensor,
        targets: torch.Tensor,
    ) -> float:
        """Compute Mean Reciprocal Rank."""
        true_best = targets.argmax(dim=-1)
        _, pred_order = scores.sort(dim=-1, descending=True)

        ranks = (pred_order == true_best.unsqueeze(1)).nonzero(as_tuple=True)[1] + 1
        return (1.0 / ranks.float()).mean().item()

    def _compute_kendall_tau(
        self,
        scores: torch.Tensor,
        targets: torch.Tensor,
    ) -> float:
        """Compute Kendall's Tau rank correlation."""
        # Compute for each sample and average
        taus = []
        for i in range(scores.shape[0]):
            pred_ranks = scores[i].argsort(descending=True).argsort()
            true_ranks = targets[i].argsort(descending=True).argsort()

            # Concordant and discordant pairs
            n = len(pred_ranks)
            concordant = 0
            discordant = 0
            for j in range(n):
                for k in range(j + 1, n):
                    pred_diff = pred_ranks[j] - pred_ranks[k]
                    true_diff = true_ranks[j] - true_ranks[k]
                    if pred_diff * true_diff > 0:
                        concordant += 1
                    elif pred_diff * true_diff < 0:
                        discordant += 1

            total_pairs = n * (n - 1) / 2
            if total_pairs > 0:
                tau = (concordant - discordant) / total_pairs
                taus.append(tau)

        return np.mean(taus) if taus else 0.0

    def _compute_ece(
        self,
        probs: torch.Tensor,
        targets: torch.Tensor,
        num_bins: int = 10,
    ) -> float:
        """Compute Expected Calibration Error."""
        max_probs, pred_classes = probs.max(dim=-1)
        true_best = targets.argmax(dim=-1)
        correctness = (pred_classes == true_best).float()

        bin_boundaries = torch.linspace(0, 1, num_bins + 1, device=probs.device)
        ece = 0.0
        total_samples = probs.shape[0]

        for i in range(num_bins):
            in_bin = (max_probs > bin_boundaries[i]) & (max_probs <= bin_boundaries[i + 1])
            bin_size = in_bin.sum().item()

            if bin_size > 0:
                bin_accuracy = correctness[in_bin].mean().item()
                bin_confidence = max_probs[in_bin].mean().item()
                ece += (bin_size / total_samples) * abs(bin_accuracy - bin_confidence)

        return ece

    def _compute_mce(
        self,
        probs: torch.Tensor,
        targets: torch.Tensor,
        num_bins: int = 10,
    ) -> float:
        """Compute Maximum Calibration Error."""
        max_probs, pred_classes = probs.max(dim=-1)
        true_best = targets.argmax(dim=-1)
        correctness = (pred_classes == true_best).float()

        bin_boundaries = torch.linspace(0, 1, num_bins + 1, device=probs.device)
        max_error = 0.0

        for i in range(num_bins):
            in_bin = (max_probs > bin_boundaries[i]) & (max_probs <= bin_boundaries[i + 1])
            bin_size = in_bin.sum().item()

            if bin_size > 0:
                bin_accuracy = correctness[in_bin].mean().item()
                bin_confidence = max_probs[in_bin].mean().item()
                max_error = max(max_error, abs(bin_accuracy - bin_confidence))

        return max_error

    def _compute_brier(
        self,
        probs: torch.Tensor,
        targets: torch.Tensor,
    ) -> float:
        """Compute Brier score."""
        true_best = targets.argmax(dim=-1)
        one_hot = F.one_hot(true_best, num_classes=probs.shape[-1]).float()
        brier = ((probs - one_hot) ** 2).sum(dim=-1).mean().item()
        return brier

    def _compute_per_difficulty(
        self,
        scores: torch.Tensor,
        targets: torch.Tensor,
    ) -> dict[str, float]:
        """Compute frac_oracle per difficulty level."""
        # Compute target entropy as difficulty proxy
        target_probs = F.softmax(targets, dim=-1)
        entropy = -(target_probs * (target_probs + 1e-10).log()).sum(dim=-1)
        max_entropy = np.log(targets.shape[-1])
        normalized_entropy = entropy / max_entropy

        # Split by difficulty
        easy_mask = normalized_entropy < 0.3
        hard_mask = normalized_entropy > 0.7
        medium_mask = ~easy_mask & ~hard_mask

        results = {}
        for name, mask in [("easy", easy_mask), ("medium", medium_mask), ("hard", hard_mask)]:
            if mask.sum() > 0:
                masked_scores = scores[mask]
                masked_targets = targets[mask]
                results[name] = self._compute_frac_oracle(masked_scores, masked_targets)
            else:
                results[name] = 0.0

        return results

    def _measure_latency(self, router: BaseRouter) -> float:
        """Measure inference latency in milliseconds."""
        router.eval()

        # Get a sample batch
        sample = next(iter(self.dataloader))
        query_embs = sample[0].to(self.device)
        model_embs = sample[1].to(self.device)

        # Warmup
        with torch.no_grad():
            for _ in range(10):
                _ = router(query_embs, model_embs)

        # Measure
        if torch.cuda.is_available():
            torch.cuda.synchronize()

        times = []
        with torch.no_grad():
            for _ in range(self.num_latency_runs):
                start = time.perf_counter()
                _ = router(query_embs, model_embs)
                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                end = time.perf_counter()
                times.append((end - start) * 1000)  # Convert to ms

        return np.mean(times)

    def _measure_memory(self, router: BaseRouter) -> float:
        """Measure peak GPU memory usage in MB."""
        if not torch.cuda.is_available():
            return 0.0

        router.eval()

        # Get a sample batch
        sample = next(iter(self.dataloader))
        query_embs = sample[0].to(self.device)
        model_embs = sample[1].to(self.device)

        torch.cuda.reset_peak_memory_stats()

        with torch.no_grad():
            _ = router(query_embs, model_embs)

        peak_memory = torch.cuda.max_memory_allocated() / (1024 * 1024)  # MB
        return peak_memory


def run_baseline_comparison(
    train_data: RouterDataset,
    test_data: RouterDataset,
    config: Optional[dict] = None,
) -> dict[str, EvaluationResult]:
    """Run complete baseline comparison.

    Convenience function that trains all three baselines and compares them.

    Args:
        train_data: Training dataset.
        test_data: Test dataset.
        config: Optional configuration overrides.

    Returns:
        Dictionary mapping router name to EvaluationResult.
    """
    from .base import RouterConfig
    from .gnn_baseline import SimplifiedGNNRouter
    from .multi_head_mlp import MultiHeadMLPRouter
    from .trainer import BaselineTrainer, TrainerConfig
    from .two_tower import TwoTowerRouter

    # Default configuration
    router_config = RouterConfig(
        query_dim=config.get("query_dim", 512) if config else 512,
        model_dim=config.get("model_dim", 512) if config else 512,
        hidden_dim=config.get("hidden_dim", 384) if config else 384,
        num_models=config.get("num_models", 7) if config else 7,
    )

    trainer_config = TrainerConfig(
        num_epochs=config.get("num_epochs", 50) if config else 50,
        early_stopping_patience=config.get("patience", 10) if config else 10,
    )

    # Initialize routers
    routers = [
        ("Two Tower", TwoTowerRouter(router_config)),
        ("Multi-Head MLP", MultiHeadMLPRouter(router_config)),
        ("GNN (Simplified)", SimplifiedGNNRouter(router_config)),
    ]

    # Train each router
    trained_routers = []
    for name, router in routers:
        print(f"\n{'='*50}")
        print(f"Training {name}...")
        print(f"{'='*50}")

        trainer = BaselineTrainer(router, trainer_config)
        metrics = trainer.train(train_data)

        print(f"Best epoch: {metrics.best_epoch}")
        print(f"Best frac_oracle: {metrics.best_metrics().get('frac_oracle', 0):.4f}")

        trained_routers.append((name, router))

    # Evaluate
    evaluator = BaselineEvaluator(test_data)
    results = evaluator.compare(trained_routers)

    print("\n" + "=" * 60)
    print("BASELINE COMPARISON RESULTS")
    print("=" * 60)
    evaluator.print_comparison_table(results)

    return {r.name: r for r in results}
