#!/usr/bin/env python3
"""Evaluate EB1-GNN Conductor against baseline routing strategies.

This script provides comprehensive comparison between:
1. Baseline: Pure GNN argmax routing (no conductor)
2. Conductor: EB1-GNN Conductor with LLM orchestration for hard cases
3. EB1 Ensemble: Full multi-model ensemble (reference)
4. Oracle: Perfect routing (upper bound)

Metrics computed:
- Accuracy: Correctness on benchmark tasks
- Cost: Estimated dollar cost based on model pricing
- Latency: End-to-end response time
- Routing distribution: Breakdown by path (trivial/fast/conductor)
- Per-tier accuracy: Performance stratified by difficulty

Usage:
    python scripts/eval_conductor_vs_baseline.py \\
        --checkpoint registry/eb1-gnn-preview-12-19 \\
        --dataset unified_qa \\
        --samples 100

    # Compare with EB1 ensemble
    python scripts/eval_conductor_vs_baseline.py \\
        --checkpoint registry/eb1-gnn-preview-12-19 \\
        --include-eb1-ensemble \\
        --samples 50
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

# Ensure imports work when running as script
_eb1_dir = Path(__file__).resolve().parent.parent.parent
if str(_eb1_dir) not in sys.path:
    sys.path.insert(0, str(_eb1_dir))

import numpy as np
import pandas as pd

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


# =============================================================================
# Cost Model
# =============================================================================

# Relative cost per 1K tokens (normalized to cheapest = 1.0)
MODEL_COSTS = {
    "GPT-5 Nano": 1.0,
    "GPT-5": 5.0,
    "GPT-5 High": 15.0,
    "Claude Opus 4.5": 25.0,
    "Claude Sonnet 4": 8.0,
    "Gemini 3 Flash": 0.5,
    "Gemini 3 Pro": 10.0,
    # Conductor overhead (per call)
    "conductor_overhead": 8.0,  # Claude Sonnet for orchestration
}


def estimate_cost(model_name: str | list[str], conductor_calls: int = 0) -> float:
    """Estimate relative cost for a routing decision.

    Args:
        model_name: Single model or list of models used.
        conductor_calls: Number of conductor LLM invocations.

    Returns:
        Relative cost (normalized to cheapest model = 1.0).
    """
    cost = 0.0

    if isinstance(model_name, str):
        cost += MODEL_COSTS.get(model_name, 5.0)
    else:
        for name in model_name:
            cost += MODEL_COSTS.get(name, 5.0)

    cost += conductor_calls * MODEL_COSTS["conductor_overhead"]
    return cost


# =============================================================================
# Result Dataclasses
# =============================================================================


@dataclass
class SingleQueryResult:
    """Result from evaluating a single query."""

    query_id: int
    query_text: str
    ground_truth: str
    routing_path: str
    model_used: str | list[str]
    response_text: str
    is_correct: bool
    cost: float
    latency_ms: float
    conductor_calls: int = 0
    subtask_count: int = 0
    routing_confidence: float = 0.0
    routing_entropy: float = 0.0


@dataclass
class StrategyResult:
    """Aggregate result for a routing strategy."""

    strategy_name: str
    query_results: list[SingleQueryResult] = field(default_factory=list)

    @property
    def accuracy(self) -> float:
        if not self.query_results:
            return 0.0
        return sum(r.is_correct for r in self.query_results) / len(self.query_results)

    @property
    def total_cost(self) -> float:
        return sum(r.cost for r in self.query_results)

    @property
    def avg_cost(self) -> float:
        if not self.query_results:
            return 0.0
        return self.total_cost / len(self.query_results)

    @property
    def avg_latency_ms(self) -> float:
        if not self.query_results:
            return 0.0
        return sum(r.latency_ms for r in self.query_results) / len(self.query_results)

    @property
    def path_distribution(self) -> dict[str, float]:
        if not self.query_results:
            return {}
        counts = defaultdict(int)
        for r in self.query_results:
            counts[r.routing_path] += 1
        total = len(self.query_results)
        return {k: v / total for k, v in counts.items()}

    def accuracy_by_path(self) -> dict[str, float]:
        by_path = defaultdict(list)
        for r in self.query_results:
            by_path[r.routing_path].append(r.is_correct)
        return {k: sum(v) / len(v) if v else 0.0 for k, v in by_path.items()}


@dataclass
class EvaluationSummary:
    """Summary comparing all strategies."""

    baseline: StrategyResult
    conductor: StrategyResult
    eb1_ensemble: StrategyResult | None = None
    oracle: StrategyResult | None = None

    def to_dict(self) -> dict[str, Any]:
        results = {
            "baseline": {
                "accuracy": self.baseline.accuracy,
                "avg_cost": self.baseline.avg_cost,
                "avg_latency_ms": self.baseline.avg_latency_ms,
                "path_distribution": self.baseline.path_distribution,
                "accuracy_by_path": self.baseline.accuracy_by_path(),
            },
            "conductor": {
                "accuracy": self.conductor.accuracy,
                "avg_cost": self.conductor.avg_cost,
                "avg_latency_ms": self.conductor.avg_latency_ms,
                "path_distribution": self.conductor.path_distribution,
                "accuracy_by_path": self.conductor.accuracy_by_path(),
            },
        }

        if self.eb1_ensemble:
            results["eb1_ensemble"] = {
                "accuracy": self.eb1_ensemble.accuracy,
                "avg_cost": self.eb1_ensemble.avg_cost,
                "avg_latency_ms": self.eb1_ensemble.avg_latency_ms,
            }

        # Compute deltas
        results["deltas"] = {
            "conductor_vs_baseline_accuracy": (
                self.conductor.accuracy - self.baseline.accuracy
            ),
            "conductor_vs_baseline_cost": (
                self.conductor.avg_cost - self.baseline.avg_cost
            ),
            "conductor_cost_efficiency": (
                (self.conductor.accuracy / self.conductor.avg_cost)
                / (self.baseline.accuracy / self.baseline.avg_cost)
                if self.baseline.accuracy > 0 and self.baseline.avg_cost > 0
                else 0.0
            ),
        }

        return results


# =============================================================================
# Baseline Strategy
# =============================================================================


def evaluate_baseline(
    gnn_router: Any,
    queries: list[dict[str, Any]],
    model_pool: dict[str, Any],
    ground_truth_fn: Any,
) -> StrategyResult:
    """Evaluate pure GNN argmax routing.

    Args:
        gnn_router: EB1Router instance.
        queries: List of query dictionaries with 'text' key.
        model_pool: Available models.
        ground_truth_fn: Function to check correctness.

    Returns:
        StrategyResult for baseline strategy.
    """
    result = StrategyResult(strategy_name="baseline")

    for i, query_data in enumerate(queries):
        query_text = query_data["text"]
        ground_truth = query_data.get("answer", "")

        start = time.perf_counter()

        # Route with GNN
        routing_result = gnn_router.route(query_text)
        model_name = routing_result.best_model

        # Execute (simulated for evaluation)
        model = model_pool.get(model_name)
        if model:
            response = model.response(query_text)
            response_text = response.text
        else:
            response_text = "[Model not available]"

        latency = (time.perf_counter() - start) * 1000

        # Check correctness
        is_correct = ground_truth_fn(response_text, ground_truth, query_data)

        # Determine path based on confidence
        if routing_result.confidence > 0.7:
            path = "trivial"
        elif routing_result.confidence > 0.5:
            path = "easy"
        else:
            path = "hard"

        query_result = SingleQueryResult(
            query_id=i,
            query_text=query_text,
            ground_truth=ground_truth,
            routing_path=path,
            model_used=model_name,
            response_text=response_text,
            is_correct=is_correct,
            cost=estimate_cost(model_name),
            latency_ms=latency,
            routing_confidence=routing_result.confidence,
            routing_entropy=routing_result.entropy,
        )
        result.query_results.append(query_result)

        if (i + 1) % 10 == 0:
            logger.info(f"Baseline: {i + 1}/{len(queries)} processed")

    return result


# =============================================================================
# Conductor Strategy
# =============================================================================


def evaluate_conductor(
    conductor: Any,
    queries: list[dict[str, Any]],
    ground_truth_fn: Any,
) -> StrategyResult:
    """Evaluate EB1-GNN Conductor routing.

    Args:
        conductor: EB1GNNInference instance.
        queries: List of query dictionaries.
        ground_truth_fn: Function to check correctness.

    Returns:
        StrategyResult for conductor strategy.
    """
    result = StrategyResult(strategy_name="conductor")

    for i, query_data in enumerate(queries):
        query_text = query_data["text"]
        ground_truth = query_data.get("answer", "")

        start = time.perf_counter()

        # Route and execute with conductor
        inference_result = conductor(query_text)

        latency = (time.perf_counter() - start) * 1000

        # Extract response
        response_text = inference_result.response.text if inference_result.response else ""

        # Check correctness
        is_correct = ground_truth_fn(response_text, ground_truth, query_data)

        # Get conductor calls count
        conductor_calls = 0
        subtask_count = 0
        if inference_result.conductor_trace:
            conductor_calls = inference_result.conductor_trace.conductor_calls
            subtask_count = len(inference_result.conductor_trace.subtask_results)

        query_result = SingleQueryResult(
            query_id=i,
            query_text=query_text,
            ground_truth=ground_truth,
            routing_path=inference_result.routing_path,
            model_used=inference_result.model_used,
            response_text=response_text,
            is_correct=is_correct,
            cost=estimate_cost(inference_result.model_used, conductor_calls),
            latency_ms=latency,
            conductor_calls=conductor_calls,
            subtask_count=subtask_count,
            routing_confidence=inference_result.routing_signals.confidence,
            routing_entropy=inference_result.routing_signals.entropy,
        )
        result.query_results.append(query_result)

        if (i + 1) % 10 == 0:
            logger.info(f"Conductor: {i + 1}/{len(queries)} processed")

    return result


# =============================================================================
# Ground Truth Functions
# =============================================================================


def exact_match(response: str, ground_truth: str, query_data: dict) -> bool:
    """Check if response exactly matches ground truth."""
    return response.strip().lower() == ground_truth.strip().lower()


def contains_answer(response: str, ground_truth: str, query_data: dict) -> bool:
    """Check if response contains the ground truth answer."""
    return ground_truth.strip().lower() in response.strip().lower()


def multiple_choice_match(response: str, ground_truth: str, query_data: dict) -> bool:
    """Check if response contains the correct multiple choice letter."""
    # Extract letter from response (A, B, C, D, etc.)
    import re

    response_lower = response.lower()
    gt_lower = ground_truth.lower().strip()

    # Check for explicit "the answer is X" pattern
    pattern = r"(?:answer|correct|option).*?([a-d])"
    match = re.search(pattern, response_lower)
    if match:
        return match.group(1) == gt_lower

    # Fallback to checking if the letter appears prominently
    return gt_lower in response_lower


# =============================================================================
# Data Loading
# =============================================================================


def load_unified_qa_data(
    data_path: Path,
    n_samples: int | None = None,
) -> list[dict[str, Any]]:
    """Load unified QA dataset.

    Args:
        data_path: Path to CSV file.
        n_samples: Maximum samples to load.

    Returns:
        List of query dictionaries.
    """
    df = pd.read_csv(data_path)

    if n_samples:
        df = df.head(n_samples)

    queries = []
    for _, row in df.iterrows():
        queries.append({
            "text": row["query"],
            "answer": row.get("answer", ""),
            "task_id": row.get("task_id", "unknown"),
        })

    return queries


# =============================================================================
# Main Evaluation
# =============================================================================


def run_evaluation(
    checkpoint_path: str,
    dataset_path: str,
    n_samples: int = 100,
    include_eb1_ensemble: bool = False,
    output_path: str | None = None,
) -> EvaluationSummary:
    """Run full evaluation comparing strategies.

    Args:
        checkpoint_path: Path to GNN checkpoint.
        dataset_path: Path to evaluation dataset.
        n_samples: Number of samples to evaluate.
        include_eb1_ensemble: Whether to include full EB1 comparison.
        output_path: Optional path to save results.

    Returns:
        EvaluationSummary with all results.
    """
    logger.info("Loading GNN router...")
    try:
        from embergraphrouter.inference.eb1_router import EB1Router
    except ImportError:
        from embergraphrouter.inference import EB1Router

    gnn_router = EB1Router.from_registry(checkpoint_path)

    logger.info("Loading EB1-GNN Conductor...")
    try:
        from ember.takumi.eb1_gnn_conductor import EB1GNNInference
    except ImportError:
        # Add src to path
        src_path = Path(__file__).parent.parent.parent / "src"
        if str(src_path) not in sys.path:
            sys.path.insert(0, str(src_path))
        from ember.takumi.eb1_gnn_conductor import EB1GNNInference

    conductor = EB1GNNInference.from_registry(checkpoint_path)

    logger.info("Loading dataset...")
    queries = load_unified_qa_data(Path(dataset_path), n_samples)
    logger.info(f"Loaded {len(queries)} queries")

    # Set up model pool (for baseline)
    from ember.api import models as ember_models

    model_pool = {
        "GPT-5 Nano": ember_models.instance("gpt-5-nano"),
        "GPT-5": ember_models.instance("gpt-5"),
        "Claude Sonnet 4": ember_models.instance("claude-sonnet-4-20250514"),
        "Gemini 3 Flash": ember_models.instance("gemini-2.0-flash"),
    }

    # Choose ground truth function based on dataset
    ground_truth_fn = contains_answer  # Default

    logger.info("Evaluating baseline strategy...")
    baseline_result = evaluate_baseline(
        gnn_router, queries, model_pool, ground_truth_fn
    )

    logger.info("Evaluating conductor strategy...")
    conductor_result = evaluate_conductor(conductor, queries, ground_truth_fn)

    eb1_result = None
    if include_eb1_ensemble:
        logger.info("Evaluating EB1 ensemble (this may take a while)...")
        # TODO: Implement EB1 ensemble evaluation
        pass

    summary = EvaluationSummary(
        baseline=baseline_result,
        conductor=conductor_result,
        eb1_ensemble=eb1_result,
    )

    # Log summary
    logger.info("=" * 60)
    logger.info("EVALUATION SUMMARY")
    logger.info("=" * 60)
    logger.info(f"Baseline Accuracy: {baseline_result.accuracy:.2%}")
    logger.info(f"Conductor Accuracy: {conductor_result.accuracy:.2%}")
    logger.info(f"Accuracy Delta: {summary.to_dict()['deltas']['conductor_vs_baseline_accuracy']:+.2%}")
    logger.info("-" * 60)
    logger.info(f"Baseline Avg Cost: {baseline_result.avg_cost:.2f}")
    logger.info(f"Conductor Avg Cost: {conductor_result.avg_cost:.2f}")
    logger.info(f"Cost Delta: {summary.to_dict()['deltas']['conductor_vs_baseline_cost']:+.2f}")
    logger.info("-" * 60)
    logger.info(f"Baseline Avg Latency: {baseline_result.avg_latency_ms:.0f}ms")
    logger.info(f"Conductor Avg Latency: {conductor_result.avg_latency_ms:.0f}ms")
    logger.info("-" * 60)
    logger.info("Conductor Path Distribution:")
    for path, pct in conductor_result.path_distribution.items():
        logger.info(f"  {path}: {pct:.1%}")
    logger.info("-" * 60)
    logger.info("Conductor Accuracy by Path:")
    for path, acc in conductor_result.accuracy_by_path().items():
        logger.info(f"  {path}: {acc:.2%}")

    if output_path:
        with open(output_path, "w") as f:
            json.dump(summary.to_dict(), f, indent=2)
        logger.info(f"Results saved to: {output_path}")

    return summary


def main():
    parser = argparse.ArgumentParser(
        description="Evaluate EB1-GNN Conductor vs baseline routing"
    )
    parser.add_argument(
        "--checkpoint",
        type=str,
        required=True,
        help="Path to GNN checkpoint (e.g., registry/eb1-gnn-preview-12-19)",
    )
    parser.add_argument(
        "--dataset",
        type=str,
        default="data/unified_qa_data.csv",
        help="Path to evaluation dataset",
    )
    parser.add_argument(
        "--samples",
        type=int,
        default=100,
        help="Number of samples to evaluate",
    )
    parser.add_argument(
        "--include-eb1-ensemble",
        action="store_true",
        help="Include full EB1 ensemble in comparison",
    )
    parser.add_argument(
        "--output",
        type=str,
        help="Path to save JSON results",
    )

    args = parser.parse_args()

    run_evaluation(
        checkpoint_path=args.checkpoint,
        dataset_path=args.dataset,
        n_samples=args.samples,
        include_eb1_ensemble=args.include_eb1_ensemble,
        output_path=args.output,
    )


if __name__ == "__main__":
    main()
