#!/usr/bin/env python3
"""End-to-end test of EB1-GNN Conductor with real GNN weights and real LLM.

This script:
1. Loads real GNN weights from the registry
2. Loads real queries from router_data.parquet
3. Runs GNN inference to classify query difficulty
4. Uses a real LLM conductor for hard queries
5. Compares conductor decisions against baseline GNN argmax
6. Reports detailed metrics

Usage:
    # Quick test (10 queries)
    python scripts/run_conductor_e2e.py --checkpoint registry/eb1-gnn-preview-12-19 --n-queries 10

    # Full test with metrics
    python scripts/run_conductor_e2e.py --checkpoint registry/eb1-gnn-preview-12-19 --n-queries 50 --verbose

    # Use specific conductor model
    python scripts/run_conductor_e2e.py --checkpoint registry/eb1-gnn-preview-12-19 --conductor-model claude-haiku-4.5
"""

from __future__ import annotations

import argparse
import json
import logging
import pickle
import sys
import time
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
import torch
import torch.nn.functional as F
import yaml

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


# =============================================================================
# Data Classes
# =============================================================================


@dataclass
class QuerySample:
    """A single query with its ground truth and GNN prediction."""

    query_id: int
    query_text: str
    task_id: str

    # GNN outputs
    gnn_recommendation: str
    gnn_confidence: float
    gnn_entropy: float
    gnn_probabilities: dict[str, float]
    gnn_tier: str  # trivial, easy, hard

    # Ground truth
    ground_truth_effects: dict[str, float]
    oracle_model: str
    oracle_effect: float


@dataclass
class ConductorDecision:
    """Conductor's decision for a query."""

    query_id: int
    action: str  # direct, decompose, route, delegate
    reasoning: str
    selected_model: str | list[str]
    subtasks: list[dict] | None = None
    final_answer: str | None = None
    latency_ms: float = 0.0


@dataclass
class E2EResult:
    """End-to-end evaluation result."""

    query: QuerySample
    baseline_model: str
    baseline_effect: float
    conductor_model: str | list[str]
    conductor_effect: float
    conductor_decision: ConductorDecision | None
    path_taken: str  # trivial, easy, hard
    would_differ: bool
    improvement_over_baseline: float


@dataclass
class E2EMetrics:
    """Aggregate metrics from E2E evaluation."""

    n_queries: int = 0
    n_trivial: int = 0
    n_easy: int = 0
    n_hard: int = 0

    # Baseline metrics
    baseline_accuracy: float = 0.0  # % matching oracle
    baseline_avg_effect: float = 0.0
    baseline_oracle_gap: float = 0.0  # avg gap to oracle

    # Conductor metrics
    conductor_avg_effect: float = 0.0
    conductor_oracle_gap: float = 0.0
    conductor_improvements: int = 0  # queries where conductor beat baseline
    conductor_degradations: int = 0  # queries where conductor was worse

    # LLM usage
    conductor_llm_calls: int = 0
    conductor_avg_latency_ms: float = 0.0

    # Model selection distribution
    baseline_model_counts: dict[str, int] = field(default_factory=dict)
    conductor_model_counts: dict[str, int] = field(default_factory=dict)

    def print_summary(self) -> None:
        """Print metrics summary."""
        print("\n" + "=" * 70)
        print("END-TO-END EVALUATION RESULTS")
        print("=" * 70)

        print(f"\nTotal Queries: {self.n_queries}")
        print(f"Path Distribution:")
        print(f"  Trivial (cheap model): {self.n_trivial} ({self.n_trivial/max(1,self.n_queries):.1%})")
        print(f"  Easy (GNN direct):     {self.n_easy} ({self.n_easy/max(1,self.n_queries):.1%})")
        print(f"  Hard (conductor):      {self.n_hard} ({self.n_hard/max(1,self.n_queries):.1%})")

        print(f"\nBaseline (GNN Argmax) Performance:")
        print(f"  Accuracy (matches oracle): {self.baseline_accuracy:.1%}")
        print(f"  Avg Effect: {self.baseline_avg_effect:.4f}")
        print(f"  Gap to Oracle: {self.baseline_oracle_gap:.4f}")

        print(f"\nConductor Performance:")
        print(f"  Avg Effect: {self.conductor_avg_effect:.4f}")
        print(f"  Gap to Oracle: {self.conductor_oracle_gap:.4f}")
        print(f"  Improvements over baseline: {self.conductor_improvements}")
        print(f"  Degradations vs baseline: {self.conductor_degradations}")

        if self.conductor_llm_calls > 0:
            print(f"\nConductor LLM Usage:")
            print(f"  Total LLM calls: {self.conductor_llm_calls}")
            print(f"  Avg latency: {self.conductor_avg_latency_ms:.1f}ms")

        print(f"\nBaseline Model Selection:")
        for model, count in sorted(
            self.baseline_model_counts.items(), key=lambda x: x[1], reverse=True
        ):
            print(f"  {model}: {count} ({count/max(1,self.n_queries):.1%})")

        if self.conductor_model_counts:
            print(f"\nConductor Model Selection:")
            for model, count in sorted(
                self.conductor_model_counts.items(), key=lambda x: x[1], reverse=True
            ):
                print(f"  {model}: {count}")

        print("=" * 70)


# =============================================================================
# GNN Router Wrapper
# =============================================================================


class E2EGNNRouter:
    """GNN router wrapper for E2E testing."""

    def __init__(
        self,
        checkpoint_path: Path,
        device: str = "cpu",
    ):
        self.checkpoint_path = checkpoint_path
        self.device = device
        self.model = None
        self.llm_names: list[str] = []
        self._cache = None
        self._load_model()

    def _load_model(self) -> None:
        """Load GNN model and training cache."""
        # Load training cache
        cache_path = self.checkpoint_path / "training_graph_cache.pkl"
        logger.info("Loading training graph cache from %s", cache_path)
        with open(cache_path, "rb") as f:
            self._cache = pickle.load(f)

        self.llm_names = self._cache.get(
            "llm_order", [f"LLM_{i}" for i in range(self._cache["num_llms"])]
        )

        # Load config
        config_path = self.checkpoint_path / "config.yaml"
        raw_config = {}
        if config_path.exists():
            with open(config_path) as f:
                raw_config = yaml.safe_load(f)
            config = {
                "hidden_features": raw_config.get("hidden_dim", 384),
                "num_gnn_layers": raw_config.get("num_gat_layers", 3),
                "num_attention_heads": raw_config.get("num_gat_heads", 4),
                "dropout": raw_config.get("dropout", 0.4),
                "use_difficulty_head": raw_config.get("use_difficulty_head", True),
                "hidden_multiplier": 2,
            }
        else:
            config = {
                "hidden_features": 384,
                "num_gnn_layers": 3,
                "num_attention_heads": 4,
                "dropout": 0.4,
                "use_difficulty_head": True,
                "hidden_multiplier": 2,
            }

        # Load LLM embeddings from config-specified path if available
        llm_embedding_path = raw_config.get("llm_embedding_path")
        if llm_embedding_path:
            # Resolve relative to embergraphrouter directory
            full_llm_path = self.checkpoint_path.parent.parent / llm_embedding_path
            if full_llm_path.exists():
                logger.info("Loading LLM embeddings from %s", full_llm_path)
                with open(full_llm_path, "rb") as f:
                    llm_emb_data = pickle.load(f)
                # Handle different formats
                if isinstance(llm_emb_data, dict):
                    self._llm_embeddings = np.array([
                        llm_emb_data[name] for name in self.llm_names
                    ])
                else:
                    self._llm_embeddings = np.array(llm_emb_data)
                logger.info("LLM embeddings shape: %s", self._llm_embeddings.shape)
            else:
                self._llm_embeddings = None
        else:
            self._llm_embeddings = None

        # Load model weights
        model_path = self.checkpoint_path / "model.pth"
        state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
        state_keys = set(state_dict.keys())

        config["use_contrastive"] = any(k.startswith("contrastive_encoder.") for k in state_keys)
        config["use_uncertainty_head"] = any(k.startswith("uncertainty_proj.") for k in state_keys)

        # Infer dimensions
        if "model_align.query_transform.weight" in state_dict:
            config["hidden_features"] = state_dict["model_align.query_transform.weight"].shape[0]
            config["query_dim"] = state_dict["model_align.query_transform.weight"].shape[1]
        if "model_align.llm_transform.weight" in state_dict:
            config["llm_dim"] = state_dict["model_align.llm_transform.weight"].shape[1]

        self._config = config

        # Build model
        from embergraphrouter.model.graph_nn import GATEncoderDecoderNet

        self.model = GATEncoderDecoderNet(
            query_feature_dim=config.get("query_dim", 512),
            llm_feature_dim=config.get("llm_dim", 512),
            hidden_features=config["hidden_features"],
            in_edges=6,
            dropout=config["dropout"],
            num_gnn_layers=config["num_gnn_layers"],
            use_difficulty_head=config["use_difficulty_head"],
            use_uncertainty_head=config["use_uncertainty_head"],
            num_attention_heads=config["num_attention_heads"],
            hidden_multiplier=config["hidden_multiplier"],
            use_contrastive=config["use_contrastive"],
        )

        # Load weights (strict=False for deprecated heads)
        self.model.load_state_dict(state_dict, strict=False)
        self.model.eval()
        self.model.to(self.device)

        logger.info("Loaded GNN model with %d LLMs: %s", len(self.llm_names), self.llm_names)

    def route_from_embedding(
        self, query_embedding: np.ndarray
    ) -> dict[str, Any]:
        """Route a query from its embedding.

        Args:
            query_embedding: Query embedding (dimension must match model)

        Returns:
            Dict with best_model, confidence, entropy, probabilities, scores
        """
        # Use LLM embeddings from config path or cache
        if self._llm_embeddings is not None:
            llm_embeddings = torch.from_numpy(self._llm_embeddings).float()
        else:
            llm_embeddings = torch.from_numpy(self._cache["llm_embeddings"]).float()

        # Build minimal graph for single query
        query_emb = torch.from_numpy(query_embedding).float().unsqueeze(0)
        task_emb = query_emb.clone()  # Use query as task for simplicity
        n_llms = len(self.llm_names)

        # Edge index: query 0 -> LLMs 1..n_llms
        source = torch.zeros(n_llms, dtype=torch.long)
        dest = torch.arange(n_llms, dtype=torch.long) + 1
        edge_index = torch.stack([source, dest])

        # Edge weight (default)
        edge_weight = torch.zeros(n_llms, 6)
        edge_weight[:, 3] = 0.7  # w_perf
        edge_weight[:, 4] = 0.2  # w_cost
        edge_weight[:, 5] = 0.1  # w_latency

        # Masks
        edge_mask = torch.ones(n_llms, dtype=torch.bool)
        edge_can_see = torch.ones(n_llms, dtype=torch.bool)

        # Move to device
        query_emb = query_emb.to(self.device)
        task_emb = task_emb.to(self.device)
        llm_embeddings = llm_embeddings.to(self.device)
        edge_index = edge_index.to(self.device)
        edge_weight = edge_weight.to(self.device)
        edge_mask = edge_mask.to(self.device)
        edge_can_see = edge_can_see.to(self.device)

        with torch.no_grad():
            logits = self.model(
                task_features=task_emb,
                query_features=query_emb,
                llm_features=llm_embeddings,
                edge_index=edge_index,
                edge_mask=edge_mask,
                edge_can_see=edge_can_see,
                edge_weight=edge_weight,
            )

        logits = logits.cpu().numpy()
        probs = F.softmax(torch.from_numpy(logits), dim=0).numpy()

        # Build result
        best_idx = int(np.argmax(probs))
        best_model = self.llm_names[best_idx]
        confidence = float(probs[best_idx])

        # Mean logit - THE KEY DIFFICULTY SIGNAL
        # High mean_logit (+0.5 to +1.0): Trivial queries (all LLMs likely correct)
        # Near-zero: Medium/hard queries
        # Negative (-0.3 to -0.6): Very hard/impossible queries
        mean_logit = float(logits.mean())

        # Normalized entropy (secondary signal)
        entropy = -np.sum(probs * np.log(probs + 1e-8))
        max_entropy = np.log(n_llms)
        normalized_entropy = float(entropy / max_entropy)

        return {
            "best_model": best_model,
            "best_model_idx": best_idx,
            "confidence": confidence,
            "mean_logit": mean_logit,  # Primary difficulty signal
            "entropy": normalized_entropy,
            "probabilities": {name: float(probs[i]) for i, name in enumerate(self.llm_names)},
            "scores": {name: float(logits[i]) for i, name in enumerate(self.llm_names)},
        }


# =============================================================================
# Conductor Prompts
# =============================================================================

CONDUCTOR_SYSTEM_PROMPT = """You are an AI routing conductor. Your job is to analyze user queries and decide the best way to handle them using available LLM models.

You have access to a GNN router that provides per-model scores based on query characteristics. The mean_logit signal indicates query difficulty:
- High mean_logit (> +0.3): TRIVIAL - all models likely correct, use cheapest
- Near-zero mean_logit: MEDIUM - trust GNN's top recommendation
- Negative mean_logit (< -0.1): HARD - few models will succeed, choose carefully

## Available Models (with capabilities):

**Gemini 2.5 Flash Lite** ($0.10/$0.40 per 1M tokens)
- BEST FOR: High-volume classification, simple QA, content filtering, cost-sensitive workloads
- WEAKNESSES: Lower accuracy on complex reasoning, math, multi-step logic

**Claude Haiku 4.5** ($1.00/$5.00 per 1M tokens)
- BEST FOR: High-volume coding assistants, agentic workflows, latency-sensitive tasks
- BENCHMARKS: SWE-bench 73.3%, 90% of Sonnet performance at 1/3 cost

**GPT-5 Nano** ($0.05/$0.40 per 1M tokens)
- BEST FOR: Simple code snippets, data extraction, instruction-following, drafting
- WEAKNESSES: Weaker on complex reasoning, advanced coding, hard benchmarks

**GPT-5.1** ($1.25/$10.00 per 1M tokens)
- BEST FOR: Complex analysis, research, multi-step reasoning, math, general capability
- BENCHMARKS: MMLU 91%, SWE-bench 76.3%, AIME 94.6%

**GPT-5.1 Codex** ($1.25/$10.00 per 1M tokens)
- BEST FOR: Complex software engineering, debugging, multi-file refactoring
- BENCHMARKS: SWE-bench 77-78% (top tier)

**Claude Opus 4.5** ($5.00/$25.00 per 1M tokens)
- BEST FOR: Hard coding, complex debugging, systematic reasoning, repository-scale work
- BENCHMARKS: SWE-bench 80.9% (industry-leading), MMLU 90.8%

**Gemini 3 Pro Preview** ($2.00/$12.00 per 1M tokens)
- BEST FOR: Complex multimodal, long-document analysis, advanced math, web development
- BENCHMARKS: MMLU 91.8%, GPQA 91.9%, AIME 95%, 1M token context

## Decision Framework:

1. **DIRECT**: Query is simple OR GNN is confident (high mean_logit) → route to GNN's top choice or cheapest
2. **ROUTE**: You disagree with GNN based on query content → specify better model
3. **DECOMPOSE**: Query is complex with multiple parts → break into subtasks with model per subtask
4. **DELEGATE**: Needs human judgment → explain why

## Response Format (JSON):
{
    "action": "direct|decompose|route|delegate",
    "reasoning": "Brief explanation citing the mean_logit difficulty and model strengths",
    "selected_model": "Model name",
    "subtasks": [{"id": "s1", "description": "...", "preferred_model": "..."}]
}"""


def create_conductor_prompt(
    query: str,
    gnn_recommendation: str,
    mean_logit: float,
    confidence: float,
    probabilities: dict[str, float],
) -> str:
    """Create the conductor prompt for a query.

    Args:
        query: The user query
        gnn_recommendation: GNN's top-1 model recommendation
        mean_logit: Mean of all logits (key difficulty signal)
            - High (+0.5 to +1.0): Trivial (all LLMs likely correct)
            - Near zero: Medium difficulty
            - Negative (-0.3 to -0.6): Very hard/impossible
        confidence: Max probability (secondary signal)
        probabilities: Model probabilities
    """
    top_models = sorted(probabilities.items(), key=lambda x: x[1], reverse=True)[:3]
    top_models_str = ", ".join(f"{m}: {p:.2f}" for m, p in top_models)

    # Interpret difficulty
    if mean_logit > 0.3:
        difficulty_desc = "TRIVIAL (high mean_logit suggests all models likely correct)"
    elif mean_logit > -0.1:
        difficulty_desc = "MEDIUM (near-zero mean_logit suggests mixed model performance)"
    else:
        difficulty_desc = "HARD (negative mean_logit suggests few models will succeed)"

    return f"""Query: {query}

GNN Routing Signals:
- Recommended model: {gnn_recommendation}
- Mean logit (difficulty signal): {mean_logit:+.2f} → {difficulty_desc}
- Confidence (max prob): {confidence:.2f}
- Top models: {top_models_str}

Based on the mean_logit difficulty signal and query content, what is your routing decision?"""


# =============================================================================
# Conductor Runner
# =============================================================================


def run_conductor(
    query: str,
    gnn_result: dict[str, Any],
    conductor_model: str,
) -> ConductorDecision:
    """Run the conductor LLM to make a routing decision.

    Args:
        query: The user query
        gnn_result: GNN routing result
        conductor_model: Model to use as conductor (e.g., "claude-opus-4-5-20251101")

    Returns:
        ConductorDecision with the routing choice
    """
    from ember.api import models

    start_time = time.time()

    # Create conductor prompt
    prompt = create_conductor_prompt(
        query=query,
        gnn_recommendation=gnn_result["best_model"],
        mean_logit=gnn_result["mean_logit"],
        confidence=gnn_result["confidence"],
        probabilities=gnn_result["probabilities"],
    )

    # Full prompt with system context
    full_prompt = f"{CONDUCTOR_SYSTEM_PROMPT}\n\n{prompt}"

    # Call conductor LLM
    try:
        response = models.response(
            conductor_model,
            full_prompt,
            temperature=0.0,
            max_tokens=500,
        )

        latency_ms = (time.time() - start_time) * 1000

        # Parse JSON response
        response_text = response.text if hasattr(response, "text") else str(response)

        # Extract JSON from response (handle markdown code blocks)
        json_str = response_text
        if "```json" in json_str:
            json_str = json_str.split("```json")[1].split("```")[0]
        elif "```" in json_str:
            json_str = json_str.split("```")[1].split("```")[0]

        decision_data = json.loads(json_str.strip())

        return ConductorDecision(
            query_id=0,  # Set by caller
            action=decision_data.get("action", "direct"),
            reasoning=decision_data.get("reasoning", ""),
            selected_model=decision_data.get("selected_model", gnn_result["best_model"]),
            subtasks=decision_data.get("subtasks"),
            latency_ms=latency_ms,
        )

    except Exception as e:
        logger.warning("Conductor error: %s. Falling back to GNN.", e)
        latency_ms = (time.time() - start_time) * 1000
        return ConductorDecision(
            query_id=0,
            action="direct",
            reasoning=f"Fallback due to error: {e}",
            selected_model=gnn_result["best_model"],
            latency_ms=latency_ms,
        )


# =============================================================================
# Routing Tiers
# =============================================================================


def classify_tier(
    mean_logit: float,
    trivial_threshold: float = 0.3,
    hard_threshold: float = -0.1,
) -> str:
    """Classify query into routing tier based on mean logit.

    Mean logit is the key difficulty signal (6x better correlation than confidence):
    - High mean_logit (+0.5 to +1.0): Trivial queries (all LLMs likely correct)
    - Near-zero: Medium/hard queries (use GNN top-1)
    - Negative (-0.3 to -0.6): Very hard/impossible queries (conductor)

    Args:
        mean_logit: Mean of all logits for this query (difficulty signal)
        trivial_threshold: Above this -> trivial (route to cheap model)
        hard_threshold: Below this -> hard (conductor orchestration)

    Returns:
        "trivial", "easy", or "hard"
    """
    if mean_logit > trivial_threshold:
        return "trivial"
    elif mean_logit > hard_threshold:
        return "easy"
    else:
        return "hard"


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


def parse_embedding(emb_value: Any) -> np.ndarray:
    """Parse embedding from various formats.

    Args:
        emb_value: Embedding as string, list, or numpy array

    Returns:
        Numpy array of floats
    """
    import ast

    if emb_value is None:
        raise ValueError("Embedding value is None")
    if isinstance(emb_value, str):
        return np.array(ast.literal_eval(emb_value), dtype=np.float32)
    elif isinstance(emb_value, (list, np.ndarray)):
        return np.array(emb_value, dtype=np.float32)
    else:
        raise ValueError(f"Unknown embedding type: {type(emb_value)}")


def load_queries(
    data_path: Path,
    n_queries: int,
    random_seed: int = 42,
    embedding_col: str = "query_embedding_qwen3_instruct_router_v1_512",
) -> pd.DataFrame:
    """Load random sample of queries from router data.

    Args:
        data_path: Path to router_data.parquet
        n_queries: Number of queries to sample
        random_seed: Random seed for reproducibility
        embedding_col: Name of the embedding column to use

    Returns:
        DataFrame with unique queries and their effects
    """
    logger.info("Loading queries from %s", data_path)
    df = pd.read_parquet(data_path)

    # Check if the preferred embedding column exists
    if embedding_col not in df.columns:
        logger.warning("Embedding column %s not found, using query_embedding", embedding_col)
        embedding_col = "query_embedding"

    logger.info("Using embedding column: %s", embedding_col)

    # Get unique queries with the proper embedding
    agg_dict = {
        "task_id": "first",
        embedding_col: "first",
    }
    unique_queries = df.groupby("query").agg(agg_dict).reset_index()

    # Rename embedding column for consistency
    if embedding_col != "query_embedding":
        unique_queries = unique_queries.rename(columns={embedding_col: "query_embedding"})

    # Sample
    np.random.seed(random_seed)
    if len(unique_queries) > n_queries:
        sample_idx = np.random.choice(len(unique_queries), n_queries, replace=False)
        unique_queries = unique_queries.iloc[sample_idx]

    logger.info("Sampled %d unique queries", len(unique_queries))
    return unique_queries, df


def get_ground_truth_effects(
    query: str,
    full_df: pd.DataFrame,
    llm_names: list[str],
) -> dict[str, float]:
    """Get ground truth effects for all LLMs for a query.

    Args:
        query: Query text
        full_df: Full router data DataFrame
        llm_names: List of LLM names

    Returns:
        Dict mapping LLM name to effect
    """
    query_rows = full_df[full_df["query"] == query]
    effects = {}
    for _, row in query_rows.iterrows():
        llm = row["llm"]
        if llm in llm_names:
            effects[llm] = float(row["effect"])
    return effects


def run_e2e_evaluation(
    checkpoint_path: Path,
    data_path: Path,
    n_queries: int,
    conductor_model: str,
    embedding_col: str = "query_embedding_qwen3_instruct_router_v1_256",
    run_conductor_for_hard: bool = True,
    verbose: bool = False,
) -> tuple[list[E2EResult], E2EMetrics]:
    """Run end-to-end evaluation.

    Args:
        checkpoint_path: Path to GNN checkpoint
        data_path: Path to router data
        n_queries: Number of queries to evaluate
        conductor_model: Model to use as conductor
        embedding_col: Embedding column to use from data
        run_conductor_for_hard: Whether to actually call conductor for hard queries
        verbose: Print per-query details

    Returns:
        Tuple of (results list, aggregate metrics)
    """
    # Load GNN router
    logger.info("Loading GNN router from %s", checkpoint_path)
    router = E2EGNNRouter(checkpoint_path)

    # Load queries from parquet with proper embeddings
    queries_df, full_df = load_queries(
        data_path, n_queries, embedding_col=embedding_col
    )

    results: list[E2EResult] = []
    metrics = E2EMetrics()
    metrics.n_queries = len(queries_df)

    conductor_latencies = []

    for idx, row in queries_df.iterrows():
        query_text = row["query"]
        task_id = row["task_id"]
        query_embedding = parse_embedding(row["query_embedding"])

        # Get GNN prediction
        gnn_result = router.route_from_embedding(query_embedding)

        # Classify tier using mean_logit (the key difficulty signal)
        tier = classify_tier(gnn_result["mean_logit"])

        # Get ground truth
        ground_truth = get_ground_truth_effects(query_text, full_df, router.llm_names)
        if not ground_truth:
            continue  # Skip if no ground truth

        oracle_model = max(ground_truth, key=ground_truth.get)
        oracle_effect = ground_truth[oracle_model]

        # Baseline decision (GNN argmax)
        baseline_model = gnn_result["best_model"]
        baseline_effect = ground_truth.get(baseline_model, 0.0)

        # Conductor decision
        conductor_decision = None
        conductor_model_used = baseline_model
        conductor_effect = baseline_effect

        if tier == "trivial":
            # Use cheapest model
            conductor_model_used = "Gemini 2.5 Flash Lite"
            conductor_effect = ground_truth.get(conductor_model_used, 0.0)
            metrics.n_trivial += 1
        elif tier == "easy":
            # Trust GNN
            conductor_model_used = baseline_model
            conductor_effect = baseline_effect
            metrics.n_easy += 1
        else:  # hard
            metrics.n_hard += 1
            if run_conductor_for_hard:
                # Run actual conductor
                conductor_decision = run_conductor(
                    query=query_text,
                    gnn_result=gnn_result,
                    conductor_model=conductor_model,
                )
                conductor_decision.query_id = idx
                conductor_latencies.append(conductor_decision.latency_ms)
                metrics.conductor_llm_calls += 1

                # Get conductor's selected model effect
                if isinstance(conductor_decision.selected_model, list):
                    # For decomposed tasks, average effects
                    selected = conductor_decision.selected_model
                    conductor_effect = np.mean([
                        ground_truth.get(m, 0.0) for m in selected if m in ground_truth
                    ])
                    conductor_model_used = selected
                else:
                    conductor_model_used = conductor_decision.selected_model
                    conductor_effect = ground_truth.get(conductor_model_used, 0.0)

        # Track model usage
        if isinstance(baseline_model, str):
            metrics.baseline_model_counts[baseline_model] = (
                metrics.baseline_model_counts.get(baseline_model, 0) + 1
            )
        if isinstance(conductor_model_used, str):
            metrics.conductor_model_counts[conductor_model_used] = (
                metrics.conductor_model_counts.get(conductor_model_used, 0) + 1
            )
        elif isinstance(conductor_model_used, list):
            for m in conductor_model_used:
                metrics.conductor_model_counts[m] = (
                    metrics.conductor_model_counts.get(m, 0) + 1
                )

        # Track improvements
        improvement = conductor_effect - baseline_effect
        if improvement > 0.01:
            metrics.conductor_improvements += 1
        elif improvement < -0.01:
            metrics.conductor_degradations += 1

        # Build result
        result = E2EResult(
            query=QuerySample(
                query_id=idx,
                query_text=query_text,
                task_id=task_id,
                gnn_recommendation=gnn_result["best_model"],
                gnn_confidence=gnn_result["confidence"],
                gnn_entropy=gnn_result["entropy"],
                gnn_probabilities=gnn_result["probabilities"],
                gnn_tier=tier,
                ground_truth_effects=ground_truth,
                oracle_model=oracle_model,
                oracle_effect=oracle_effect,
            ),
            baseline_model=baseline_model,
            baseline_effect=baseline_effect,
            conductor_model=conductor_model_used,
            conductor_effect=conductor_effect,
            conductor_decision=conductor_decision,
            path_taken=tier,
            would_differ=conductor_model_used != baseline_model,
            improvement_over_baseline=improvement,
        )
        results.append(result)

        if verbose:
            print(f"\n[{idx}] {query_text[:60]}...")
            print(f"  GNN: {baseline_model} (mean_logit={gnn_result['mean_logit']:+.2f}, conf={gnn_result['confidence']:.2f})")
            print(f"  Tier: {tier}")
            print(f"  Oracle: {oracle_model} (effect={oracle_effect:.3f})")
            if conductor_decision:
                print(f"  Conductor: {conductor_decision.action} -> {conductor_model_used}")
                print(f"  Reasoning: {conductor_decision.reasoning[:80]}...")

    # Compute aggregate metrics
    if results:
        metrics.baseline_accuracy = sum(
            1 for r in results if r.baseline_model == r.query.oracle_model
        ) / len(results)
        metrics.baseline_avg_effect = np.mean([r.baseline_effect for r in results])
        metrics.baseline_oracle_gap = np.mean([
            r.query.oracle_effect - r.baseline_effect for r in results
        ])
        metrics.conductor_avg_effect = np.mean([r.conductor_effect for r in results])
        metrics.conductor_oracle_gap = np.mean([
            r.query.oracle_effect - r.conductor_effect for r in results
        ])

    if conductor_latencies:
        metrics.conductor_avg_latency_ms = np.mean(conductor_latencies)

    return results, metrics


def main():
    parser = argparse.ArgumentParser(description="End-to-end conductor evaluation")
    parser.add_argument(
        "--checkpoint",
        type=str,
        default="embergraphrouter/registry/anticollapse_ablation_baseline",
        help="Path to GNN checkpoint (must match embedding type)",
    )
    parser.add_argument(
        "--data",
        type=str,
        default="embergraphrouter/data/router_data_qwen3_instruct_router_v1.parquet",
        help="Path to router data with qwen3 instruct embeddings",
    )
    parser.add_argument(
        "--embedding-col",
        type=str,
        default="query_embedding_qwen3_instruct_router_v1_256",
        help="Embedding column name to use (must match GNN training)",
    )
    parser.add_argument(
        "--n-queries",
        type=int,
        default=20,
        help="Number of queries to evaluate",
    )
    parser.add_argument(
        "--conductor-model",
        type=str,
        default="claude-opus-4-5-20251101",
        help="Model to use as conductor (e.g., claude-opus-4-5-20251101)",
    )
    parser.add_argument(
        "--skip-conductor",
        action="store_true",
        help="Skip actual conductor calls (just simulate)",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print per-query details",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output path for results (JSON)",
    )

    args = parser.parse_args()

    checkpoint_path = Path(args.checkpoint)
    data_path = Path(args.data)

    if not checkpoint_path.exists():
        logger.error("Checkpoint not found: %s", checkpoint_path)
        sys.exit(1)

    if not data_path.exists():
        logger.error("Data file not found: %s", data_path)
        sys.exit(1)

    # Run evaluation
    logger.info(
        "Running E2E evaluation: %d queries, conductor=%s",
        args.n_queries,
        args.conductor_model,
    )

    results, metrics = run_e2e_evaluation(
        checkpoint_path=checkpoint_path,
        data_path=data_path,
        n_queries=args.n_queries,
        conductor_model=args.conductor_model,
        embedding_col=args.embedding_col,
        run_conductor_for_hard=not args.skip_conductor,
        verbose=args.verbose,
    )

    # Print results
    metrics.print_summary()

    # Save results
    if args.output:
        output_path = Path(args.output)
        output_data = {
            "metrics": {
                "n_queries": metrics.n_queries,
                "n_trivial": metrics.n_trivial,
                "n_easy": metrics.n_easy,
                "n_hard": metrics.n_hard,
                "baseline_accuracy": metrics.baseline_accuracy,
                "baseline_avg_effect": metrics.baseline_avg_effect,
                "conductor_avg_effect": metrics.conductor_avg_effect,
                "conductor_improvements": metrics.conductor_improvements,
                "conductor_degradations": metrics.conductor_degradations,
            },
            "results": [
                {
                    "query": r.query.query_text[:100],
                    "tier": r.path_taken,
                    "baseline_model": r.baseline_model,
                    "conductor_model": r.conductor_model,
                    "improvement": r.improvement_over_baseline,
                }
                for r in results
            ],
        }
        with open(output_path, "w") as f:
            json.dump(output_data, f, indent=2)
        logger.info("Results saved to %s", output_path)


if __name__ == "__main__":
    main()
