"""Comparison runner for evaluating multiple strategies on the same benchmark.

This module provides infrastructure for running multiple strategies on the
same benchmark examples and computing paired statistical tests to determine
significant differences. Supports checkpointing for resumption.
"""

import asyncio
import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Sequence

import numpy as np

from embergraphrouter.evals.benchmarks.base import Benchmark
from embergraphrouter.evals.metrics.aggregation import (
    compute_efficiency_metrics,
    compute_quality_metrics,
    compute_routing_metrics,
)
from embergraphrouter.evals.runners.base import (
    AggregateMetrics,
    EvaluationResult,
    ExampleResult,
    Runner,
)
from embergraphrouter.evals.runners.checkpoint import (
    AdaptiveRateLimiter,
    CheckpointManager,
    ProgressTracker,
    compute_config_hash,
)
from embergraphrouter.evals.strategies.base import Strategy

logger = logging.getLogger(__name__)


@dataclass
class PairedTest:
    """Result of a paired statistical test.

    Attributes:
        strategy_a: Name of first strategy.
        strategy_b: Name of second strategy.
        metric: Metric being compared.
        mean_diff: Mean difference (A - B).
        p_value: Statistical significance.
        significant: Whether difference is significant at alpha=0.05.
        test_type: Type of test used (mcnemar, paired_t, wilcoxon).
        effect_size: Cohen's d or similar effect size measure.
    """

    strategy_a: str
    strategy_b: str
    metric: str
    mean_diff: float
    p_value: float
    significant: bool
    test_type: str
    effect_size: float = 0.0
    n_samples: int = 0

    def to_dict(self) -> dict:
        """Convert to JSON-serializable dict."""
        return {
            "strategy_a": self.strategy_a,
            "strategy_b": self.strategy_b,
            "metric": self.metric,
            "mean_diff": float(self.mean_diff),
            "p_value": float(self.p_value),
            "significant": self.significant,
            "test_type": self.test_type,
            "effect_size": float(self.effect_size),
            "n_samples": self.n_samples,
        }


@dataclass
class ComparisonResult:
    """Results from comparing multiple strategies.

    Attributes:
        benchmark_name: Name of the benchmark evaluated.
        strategy_names: List of strategy names compared.
        n_examples: Number of examples evaluated.
        strategy_results: Individual EvaluationResult per strategy.
        paired_tests: Pairwise statistical test results.
        summary: High-level comparison summary.
    """

    benchmark_name: str
    strategy_names: list[str]
    n_examples: int
    strategy_results: dict[str, EvaluationResult]
    paired_tests: list[PairedTest] = field(default_factory=list)
    summary: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        """Convert to JSON-serializable dict."""
        return {
            "benchmark_name": self.benchmark_name,
            "strategy_names": self.strategy_names,
            "n_examples": self.n_examples,
            "strategy_results": {
                name: result.to_dict()
                for name, result in self.strategy_results.items()
            },
            "paired_tests": [t.to_dict() for t in self.paired_tests],
            "summary": self.summary,
        }

    def save(self, path: Path | str) -> None:
        """Save results to JSON file."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump(self.to_dict(), f, indent=2)


def mcnemar_test(correct_a: Sequence[bool], correct_b: Sequence[bool]) -> PairedTest:
    """McNemar's test for paired binary outcomes.

    Tests whether two strategies have significantly different accuracy.

    Args:
        correct_a: Boolean correctness for strategy A.
        correct_b: Boolean correctness for strategy B.

    Returns:
        PairedTest with test results.
    """
    # Build contingency table
    # b=0  b=1
    # [n00, n01]  a=0
    # [n10, n11]  a=1
    n01 = sum(1 for a, b in zip(correct_a, correct_b, strict=True) if not a and b)
    n10 = sum(1 for a, b in zip(correct_a, correct_b, strict=True) if a and not b)

    # McNemar's statistic with continuity correction
    if n01 + n10 == 0:
        p_value = 1.0
        statistic = 0.0
    else:
        statistic = ((abs(n01 - n10) - 1) ** 2) / (n01 + n10)
        # Chi-square distribution with 1 df
        from scipy import stats

        p_value = 1 - stats.chi2.cdf(statistic, df=1)

    mean_diff = (sum(correct_a) - sum(correct_b)) / len(correct_a)

    return PairedTest(
        strategy_a="",  # Filled by caller
        strategy_b="",
        metric="accuracy",
        mean_diff=mean_diff,
        p_value=p_value,
        significant=p_value < 0.05,
        test_type="mcnemar",
        n_samples=len(correct_a),
    )


def paired_t_test(values_a: Sequence[float], values_b: Sequence[float]) -> PairedTest:
    """Paired t-test for continuous metrics.

    Args:
        values_a: Metric values for strategy A.
        values_b: Metric values for strategy B.

    Returns:
        PairedTest with test results.
    """
    from scipy import stats

    differences = [a - b for a, b in zip(values_a, values_b, strict=True)]
    mean_diff = np.mean(differences)
    std_diff = np.std(differences, ddof=1)

    if std_diff == 0 or len(differences) < 2:
        p_value = 1.0 if mean_diff == 0 else 0.0
        effect_size = 0.0
    else:
        t_stat, p_value = stats.ttest_rel(values_a, values_b)
        effect_size = mean_diff / std_diff  # Cohen's d

    return PairedTest(
        strategy_a="",
        strategy_b="",
        metric="",
        mean_diff=float(mean_diff),
        p_value=float(p_value),
        significant=p_value < 0.05,
        test_type="paired_t",
        effect_size=float(effect_size),
        n_samples=len(differences),
    )


def wilcoxon_test(values_a: Sequence[float], values_b: Sequence[float]) -> PairedTest:
    """Wilcoxon signed-rank test for non-normal continuous metrics.

    Non-parametric alternative to paired t-test.

    Args:
        values_a: Metric values for strategy A.
        values_b: Metric values for strategy B.

    Returns:
        PairedTest with test results.
    """
    from scipy import stats

    differences = [a - b for a, b in zip(values_a, values_b, strict=True)]
    non_zero_diffs = [d for d in differences if d != 0]

    if len(non_zero_diffs) < 10:
        # Not enough non-tied pairs for Wilcoxon
        return paired_t_test(values_a, values_b)

    try:
        statistic, p_value = stats.wilcoxon(values_a, values_b)
    except ValueError:
        p_value = 1.0

    mean_diff = float(np.mean(differences))

    return PairedTest(
        strategy_a="",
        strategy_b="",
        metric="",
        mean_diff=mean_diff,
        p_value=float(p_value),
        significant=p_value < 0.05,
        test_type="wilcoxon",
        n_samples=len(differences),
    )


class ComparisonRunner(Runner):
    """Runner that evaluates multiple strategies on the same benchmark.

    Runs each strategy on identical examples, enabling paired statistical
    tests to determine significant differences. Supports checkpointing for
    resumption after interruption.

    Example:
        # Basic usage
        strategies = [
            TieredRouterStrategy.from_registry("eb1-gnn-preview-latest"),
            SingleModelStrategy("gpt-4"),
            SingleModelStrategy("claude-opus-4.5"),
        ]
        runner = ComparisonRunner(strategies, benchmark)
        result = await runner.run()

        # With checkpointing
        runner = ComparisonRunner(
            strategies, benchmark,
            checkpoint_dir="checkpoints",
            parallelism=10,
        )
        result = await runner.run()

        # Resume interrupted run
        runner = ComparisonRunner.resume(
            "checkpoints/run_abc123.pkl",
            strategies, benchmark,
        )
        result = await runner.run()
    """

    def __init__(
        self,
        strategies: list[Strategy],
        benchmark: Benchmark,
        max_examples: int | None = None,
        system_prompt: str | None = None,
        subset: str | None = None,
        parallel: bool = False,
        checkpoint_dir: str | Path | None = None,
        parallelism: int = 1,
        checkpoint_interval: int = 50,
        timeout_per_example: float = 120.0,
        checkpoint_manager: CheckpointManager | None = None,
    ):
        """Initialize comparison runner.

        Args:
            strategies: List of strategies to compare.
            benchmark: Benchmark to run against.
            max_examples: Maximum examples to evaluate.
            system_prompt: Optional system prompt for all queries.
            subset: Benchmark subset to use.
            parallel: If True, run strategies concurrently per example (deprecated).
            checkpoint_dir: Directory for checkpoints (None = no checkpointing).
            parallelism: Number of concurrent example evaluations.
            checkpoint_interval: Save checkpoint every N examples.
            timeout_per_example: Timeout per example in seconds.
            checkpoint_manager: Pre-configured checkpoint manager for resumption.
        """
        self._strategies = strategies
        self._benchmark = benchmark
        self._max_examples = max_examples
        self._system_prompt = system_prompt
        self._subset = subset
        self._parallel = parallel or (parallelism > 1)
        self._checkpoint_dir = Path(checkpoint_dir) if checkpoint_dir else None
        self._parallelism = parallelism
        self._checkpoint_interval = checkpoint_interval
        self._timeout = timeout_per_example
        self._checkpoint_manager = checkpoint_manager
        self._rate_limiter: AdaptiveRateLimiter | None = None

    @classmethod
    def resume(
        cls,
        checkpoint_path: str | Path,
        strategies: list[Strategy],
        benchmark: Benchmark,
        **kwargs: Any,
    ) -> "ComparisonRunner":
        """Resume runner from checkpoint.

        Args:
            checkpoint_path: Path to checkpoint file.
            strategies: Strategy instances.
            benchmark: Benchmark instance.
            **kwargs: Additional runner arguments.

        Returns:
            ComparisonRunner configured to resume.
        """
        manager = CheckpointManager.resume(checkpoint_path)

        return cls(
            strategies=strategies,
            benchmark=benchmark,
            checkpoint_manager=manager,
            **kwargs,
        )

    def _get_config_hash(self) -> str:
        """Compute hash of current configuration."""
        return compute_config_hash({
            "strategies": [s.name for s in self._strategies],
            "benchmark": self._benchmark.name,
            "max_examples": self._max_examples,
            "subset": self._subset,
        })

    def setup(self) -> None:
        """Pre-evaluation setup: warm up strategies and initialize infrastructure."""
        for strategy in self._strategies:
            strategy.warm_up()

        # Initialize checkpointing if enabled
        if self._checkpoint_dir and not self._checkpoint_manager:
            self._checkpoint_manager = CheckpointManager(
                checkpoint_dir=self._checkpoint_dir,
                config_hash=self._get_config_hash(),
                save_interval=self._checkpoint_interval,
            )

        # Initialize rate limiter for parallel execution
        if self._parallelism > 1:
            self._rate_limiter = AdaptiveRateLimiter(
                initial=self._parallelism,
                min_level=1,
                max_level=self._parallelism * 2,
            )

    async def run(self) -> ComparisonResult:
        """Run comparison evaluation.

        Supports both sequential and parallel execution with optional
        checkpointing for resumption.

        Returns:
            ComparisonResult with per-strategy results and paired tests.
        """
        # Load examples
        examples = list(self._benchmark.load(self._subset))
        if self._max_examples:
            examples = examples[: self._max_examples]

        n_examples = len(examples)
        strategy_names = [s.name for s in self._strategies]

        # Check for resumed results
        resumed_count = 0
        if self._checkpoint_manager:
            resumed_count = self._checkpoint_manager.n_completed
            if resumed_count > 0:
                logger.info(
                    f"Resuming from checkpoint: {resumed_count} examples already done"
                )

        logger.info(
            "Running comparison: %d strategies on %d examples from %s "
            "(parallelism=%d)",
            len(self._strategies),
            n_examples,
            self._benchmark.name,
            self._parallelism,
        )

        # Initialize result containers per strategy
        strategy_example_results: dict[str, list[ExampleResult]] = {
            s.name: [] for s in self._strategies
        }
        errors: list[str] = []

        # Run evaluation with appropriate execution mode
        if self._parallelism > 1 and self._rate_limiter:
            await self._run_parallel_examples(
                examples, strategy_example_results, errors
            )
        else:
            await self._run_sequential_examples(
                examples, strategy_example_results, errors
            )

        # Build EvaluationResult per strategy
        strategy_results: dict[str, EvaluationResult] = {}
        for strategy in self._strategies:
            example_results = strategy_example_results[strategy.name]
            eval_result = EvaluationResult.create(
                strategy_name=strategy.name,
                benchmark_name=self._benchmark.name,
            )
            eval_result.n_examples = n_examples
            eval_result.example_results = example_results
            eval_result.errors = errors
            eval_result.quality = compute_quality_metrics(example_results)
            eval_result.efficiency = compute_efficiency_metrics(example_results)
            eval_result.routing = compute_routing_metrics(example_results)
            strategy_results[strategy.name] = eval_result

        # Compute paired tests
        paired_tests = self._compute_paired_tests(strategy_example_results)

        # Build summary
        summary = self._build_summary(strategy_results, paired_tests)

        # Final checkpoint save
        if self._checkpoint_manager:
            self._checkpoint_manager.save()

        logger.info("Comparison complete: %d strategies evaluated", len(self._strategies))

        return ComparisonResult(
            benchmark_name=self._benchmark.name,
            strategy_names=strategy_names,
            n_examples=n_examples,
            strategy_results=strategy_results,
            paired_tests=paired_tests,
            summary=summary,
        )

    async def _run_sequential_examples(
        self,
        examples: list,
        strategy_example_results: dict[str, list[ExampleResult]],
        errors: list[str],
    ) -> None:
        """Run evaluation on examples sequentially with checkpointing."""
        progress = ProgressTracker(total=len(examples))

        for i, example in enumerate(examples):
            # Skip if already completed (resumption)
            if self._checkpoint_manager and self._checkpoint_manager.is_completed(
                example.id
            ):
                cached = self._checkpoint_manager.get_result(example.id)
                if cached:
                    for strategy in self._strategies:
                        if strategy.name in cached:
                            strategy_example_results[strategy.name].append(
                                ExampleResult(**cached[strategy.name])
                            )
                progress.increment()
                continue

            try:
                if self._parallel:
                    results = await self._evaluate_example_parallel(example)
                else:
                    results = await self._evaluate_example_sequential(example)

                result_cache = {}
                for strategy, result in zip(self._strategies, results, strict=True):
                    strategy_example_results[strategy.name].append(result)
                    result_cache[strategy.name] = result.__dict__

                # Checkpoint
                if self._checkpoint_manager:
                    self._checkpoint_manager.mark_completed(example.id, result_cache)

                progress.increment()

                if (i + 1) % 10 == 0:
                    logger.info(progress.format_status())

            except Exception as e:
                error_msg = f"Error on example {example.id}: {e}"
                logger.warning(error_msg)
                errors.append(error_msg)
                if self._checkpoint_manager:
                    self._checkpoint_manager.record_error(error_msg)

                # Add error results for each strategy
                for strategy in self._strategies:
                    strategy_example_results[strategy.name].append(
                        ExampleResult(
                            example_id=example.id,
                            strategy_name=strategy.name,
                            prediction="",
                            reference=example.reference,
                            correct=False,
                            score=0.0,
                            model_used="error",
                            error=str(e),
                            metadata=example.metadata,
                        )
                    )
                progress.increment()

    async def _run_parallel_examples(
        self,
        examples: list,
        strategy_example_results: dict[str, list[ExampleResult]],
        errors: list[str],
    ) -> None:
        """Run evaluation on examples with adaptive parallelism."""
        assert self._rate_limiter is not None

        results_lock = asyncio.Lock()
        progress = ProgressTracker(total=len(examples))

        async def process_one(example) -> None:
            # Skip if already completed
            if self._checkpoint_manager and self._checkpoint_manager.is_completed(
                example.id
            ):
                cached = self._checkpoint_manager.get_result(example.id)
                if cached:
                    async with results_lock:
                        for strategy in self._strategies:
                            if strategy.name in cached:
                                strategy_example_results[strategy.name].append(
                                    ExampleResult(**cached[strategy.name])
                                )
                progress.increment()
                return

            ctx = await self._rate_limiter.acquire()
            try:
                async with ctx:
                    try:
                        results = await asyncio.wait_for(
                            self._evaluate_example_parallel(example),
                            timeout=self._timeout,
                        )
                        await self._rate_limiter.record_success()
                    except asyncio.TimeoutError:
                        await self._rate_limiter.record_error()
                        results = [
                            ExampleResult(
                                example_id=example.id,
                                strategy_name=strategy.name,
                                prediction="",
                                reference=example.reference,
                                correct=False,
                                score=0.0,
                                model_used="timeout",
                                error="Timeout",
                                metadata=example.metadata,
                            )
                            for strategy in self._strategies
                        ]

                    result_cache = {}
                    async with results_lock:
                        for strategy, result in zip(
                            self._strategies, results, strict=True
                        ):
                            strategy_example_results[strategy.name].append(result)
                            result_cache[strategy.name] = result.__dict__

                        if self._checkpoint_manager:
                            self._checkpoint_manager.mark_completed(
                                example.id, result_cache
                            )

                    progress.increment()

            except Exception as e:
                if "rate" in str(e).lower():
                    await self._rate_limiter.record_rate_limit()
                else:
                    await self._rate_limiter.record_error()

                error_msg = f"Error on example {example.id}: {e}"
                async with results_lock:
                    errors.append(error_msg)
                    if self._checkpoint_manager:
                        self._checkpoint_manager.record_error(error_msg)

                    for strategy in self._strategies:
                        strategy_example_results[strategy.name].append(
                            ExampleResult(
                                example_id=example.id,
                                strategy_name=strategy.name,
                                prediction="",
                                reference=example.reference,
                                correct=False,
                                score=0.0,
                                model_used="error",
                                error=str(e),
                                metadata=example.metadata,
                            )
                        )
                progress.increment()

        # Run all examples
        tasks = [process_one(ex) for ex in examples]
        await asyncio.gather(*tasks)

        # Log final rate limiter stats
        logger.info(f"Rate limiter stats: {self._rate_limiter.stats}")

    async def _evaluate_example_sequential(
        self, example
    ) -> list[ExampleResult]:
        """Evaluate example with each strategy sequentially."""
        results = []
        for strategy in self._strategies:
            result = await self._evaluate_single(strategy, example)
            results.append(result)
        return results

    async def _evaluate_example_parallel(
        self, example
    ) -> list[ExampleResult]:
        """Evaluate example with all strategies in parallel."""
        tasks = [
            self._evaluate_single(strategy, example)
            for strategy in self._strategies
        ]
        return await asyncio.gather(*tasks)

    async def _evaluate_single(
        self, strategy: Strategy, example
    ) -> ExampleResult:
        """Evaluate single example with a strategy."""
        try:
            strategy_result = await strategy.infer(
                example.query,
                system_prompt=self._system_prompt or example.system_prompt,
            )

            eval_score = self._benchmark.evaluate(
                strategy_result.prediction,
                example,
            )

            routing_signals = {}
            if strategy_result.routing:
                routing_signals = dict(strategy_result.routing.signals)

            return ExampleResult(
                example_id=example.id,
                strategy_name=strategy.name,
                prediction=strategy_result.prediction,
                reference=example.reference,
                correct=eval_score.correct,
                score=eval_score.score,
                model_used=strategy_result.model_used,
                tier=strategy_result.tier,
                latency_ms=strategy_result.latency_ms,
                cost_usd=strategy_result.cost_usd,
                tokens_in=strategy_result.tokens_in,
                tokens_out=strategy_result.tokens_out,
                metadata=example.metadata,
                routing_signals=routing_signals,
            )

        except Exception as e:
            return ExampleResult(
                example_id=example.id,
                strategy_name=strategy.name,
                prediction="",
                reference=example.reference,
                correct=False,
                score=0.0,
                model_used="error",
                error=str(e),
                metadata=example.metadata,
            )

    def _compute_paired_tests(
        self,
        strategy_results: dict[str, list[ExampleResult]],
    ) -> list[PairedTest]:
        """Compute pairwise statistical tests between strategies."""
        tests = []
        strategy_names = list(strategy_results.keys())

        for i, name_a in enumerate(strategy_names):
            for name_b in strategy_names[i + 1:]:
                results_a = strategy_results[name_a]
                results_b = strategy_results[name_b]

                # Accuracy comparison (McNemar's test)
                correct_a = [r.correct for r in results_a]
                correct_b = [r.correct for r in results_b]
                acc_test = mcnemar_test(correct_a, correct_b)
                acc_test.strategy_a = name_a
                acc_test.strategy_b = name_b
                acc_test.metric = "accuracy"
                tests.append(acc_test)

                # Score comparison (paired t-test)
                scores_a = [r.score for r in results_a]
                scores_b = [r.score for r in results_b]
                score_test = paired_t_test(scores_a, scores_b)
                score_test.strategy_a = name_a
                score_test.strategy_b = name_b
                score_test.metric = "score"
                tests.append(score_test)

                # Cost comparison (Wilcoxon - typically non-normal)
                costs_a = [r.cost_usd for r in results_a]
                costs_b = [r.cost_usd for r in results_b]
                if any(c > 0 for c in costs_a + costs_b):
                    cost_test = wilcoxon_test(costs_a, costs_b)
                    cost_test.strategy_a = name_a
                    cost_test.strategy_b = name_b
                    cost_test.metric = "cost_usd"
                    tests.append(cost_test)

                # Latency comparison
                lat_a = [r.latency_ms for r in results_a]
                lat_b = [r.latency_ms for r in results_b]
                if any(lat > 0 for lat in lat_a + lat_b):
                    lat_test = wilcoxon_test(lat_a, lat_b)
                    lat_test.strategy_a = name_a
                    lat_test.strategy_b = name_b
                    lat_test.metric = "latency_ms"
                    tests.append(lat_test)

        return tests

    def _build_summary(
        self,
        strategy_results: dict[str, EvaluationResult],
        paired_tests: list[PairedTest],
    ) -> dict:
        """Build high-level comparison summary."""
        # Rank strategies by accuracy
        rankings: list[tuple[str, float]] = []
        for name, result in strategy_results.items():
            acc = result.quality.get("accuracy", AggregateMetrics(mean=0.0))
            rankings.append((name, acc.mean))

        rankings.sort(key=lambda x: -x[1])

        # Find significant differences
        significant_pairs = [
            t for t in paired_tests if t.significant and t.metric == "accuracy"
        ]

        return {
            "accuracy_ranking": [
                {"strategy": name, "accuracy": acc} for name, acc in rankings
            ],
            "best_strategy": rankings[0][0] if rankings else None,
            "n_significant_differences": len(significant_pairs),
            "significant_comparisons": [
                {
                    "comparison": f"{t.strategy_a} vs {t.strategy_b}",
                    "p_value": t.p_value,
                    "diff": t.mean_diff,
                }
                for t in significant_pairs
            ],
        }


async def run_comparison(
    strategies: list[Strategy] | list[str],
    benchmark: Benchmark | str,
    max_examples: int | None = None,
    system_prompt: str | None = None,
    subset: str | None = None,
    parallel: bool = False,
    output_path: Path | str | None = None,
) -> ComparisonResult:
    """Run strategy comparison and optionally save results.

    High-level convenience function for running comparisons.

    Args:
        strategies: List of Strategy instances or spec strings.
        benchmark: Benchmark instance or name.
        max_examples: Maximum examples to evaluate.
        system_prompt: Optional system prompt.
        subset: Benchmark subset.
        parallel: Run strategies in parallel per example.
        output_path: Optional path to save JSON results.

    Returns:
        ComparisonResult with all metrics and paired tests.
    """
    # Resolve strategies
    resolved_strategies = []
    for s in strategies:
        if isinstance(s, str):
            from embergraphrouter.evals.runners.simple import _resolve_strategy

            resolved_strategies.append(_resolve_strategy(s))
        else:
            resolved_strategies.append(s)

    # Resolve benchmark
    if isinstance(benchmark, str):
        from embergraphrouter.evals.benchmarks import get_benchmark

        benchmark = get_benchmark(benchmark)

    # Create and run runner
    runner = ComparisonRunner(
        strategies=resolved_strategies,
        benchmark=benchmark,
        max_examples=max_examples,
        system_prompt=system_prompt,
        subset=subset,
        parallel=parallel,
    )

    runner.setup()
    result = await runner.run()

    # Save results
    if output_path:
        result.save(output_path)
        logger.info("Comparison results saved to %s", output_path)

    return result


__all__ = [
    "ComparisonResult",
    "ComparisonRunner",
    "PairedTest",
    "mcnemar_test",
    "paired_t_test",
    "run_comparison",
    "wilcoxon_test",
]
