#!/usr/bin/env python3
"""Analyze how the EB1-GNN Conductor makes decisions compared to baseline.

This script provides detailed analysis of:
1. How often the GNN is called and for what purpose
2. What the conductor does with GNN responses
3. Comparison of routing decisions: baseline vs conductor
4. Cost/quality tradeoff analysis

Usage:
    python scripts/analyze_conductor_decisions.py \\
        --checkpoint registry/eb1-gnn-preview-12-19

    # With sample queries
    python scripts/analyze_conductor_decisions.py \\
        --checkpoint registry/eb1-gnn-preview-12-19 \\
        --queries "What is 2+2?" "Explain quantum computing"
"""

from __future__ import annotations

import argparse
import sys
from dataclasses import dataclass
from pathlib import Path

# 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


@dataclass
class ComparisonResult:
    """Comparison of baseline vs conductor decisions."""

    query: str
    query_id: int

    # Baseline (pure GNN argmax)
    baseline_model: str
    baseline_confidence: float
    baseline_entropy: float
    baseline_tier: str

    # Conductor decision
    conductor_path: str  # trivial, easy, hard
    conductor_would_differ: bool  # Would conductor route differently?
    conductor_reason: str

    # GNN utilization
    gnn_calls_baseline: int = 1
    gnn_calls_conductor: int = 1  # At minimum, may be more for subtasks


def analyze_single_query(
    query: str,
    query_id: int,
    router: "EB1Router",
    thresholds: "RoutingThresholds",
) -> ComparisonResult:
    """Analyze how baseline vs conductor would handle a query.

    Args:
        query: The query text.
        query_id: Query identifier.
        router: GNN router instance.
        thresholds: Routing thresholds.

    Returns:
        ComparisonResult with analysis.
    """
    from ember.takumi.eb1_gnn_conductor import RoutingSignals, RoutingTier

    # Get GNN prediction
    result = router.route(query)

    # Compute signals
    scores = list(result.scores.values())
    mean_logit = sum(scores) / len(scores)
    sorted_scores = sorted(scores, reverse=True)
    logit_gap = sorted_scores[0] - sorted_scores[1] if len(sorted_scores) > 1 else 0.0

    signals = RoutingSignals(
        confidence=result.confidence,
        entropy=result.entropy,
        difficulty=1.0 - result.confidence,
        mean_logit=mean_logit,
        logit_gap=logit_gap,
        chosen_model=result.best_model,
        chosen_model_idx=result.best_model_idx,
        probabilities=result.probabilities,
    )

    # Classify tier
    tier = signals.classify_with_thresholds(thresholds)

    # Determine conductor path and whether it would differ
    if tier == RoutingTier.TRIVIAL:
        conductor_path = "trivial"
        conductor_would_differ = True  # Would use cheap model instead
        conductor_reason = "High confidence + low entropy → use cheap model"
        gnn_calls_conductor = 1
    elif tier == RoutingTier.EASY:
        conductor_path = "easy"
        conductor_would_differ = False  # Would use same model
        conductor_reason = f"Moderate confidence → trust GNN: {result.best_model}"
        gnn_calls_conductor = 1
    else:  # HARD
        conductor_path = "hard"
        conductor_would_differ = True  # Conductor orchestrates
        conductor_reason = "Low confidence → LLM conductor orchestrates with subtasks"
        # Estimate: 1 initial + potential subtask calls
        gnn_calls_conductor = 3  # Estimate

    return ComparisonResult(
        query=query,
        query_id=query_id,
        baseline_model=result.best_model,
        baseline_confidence=result.confidence,
        baseline_entropy=result.entropy,
        baseline_tier=tier.value,
        conductor_path=conductor_path,
        conductor_would_differ=conductor_would_differ,
        conductor_reason=conductor_reason,
        gnn_calls_baseline=1,
        gnn_calls_conductor=gnn_calls_conductor,
    )


def print_comparison_report(results: list[ComparisonResult]) -> None:
    """Print a detailed comparison report."""
    print("\n" + "=" * 80)
    print("BASELINE VS CONDUCTOR: DECISION COMPARISON")
    print("=" * 80)

    # Summary statistics
    total = len(results)
    trivial = sum(1 for r in results if r.conductor_path == "trivial")
    easy = sum(1 for r in results if r.conductor_path == "easy")
    hard = sum(1 for r in results if r.conductor_path == "hard")
    would_differ = sum(1 for r in results if r.conductor_would_differ)

    print(f"\nTotal Queries: {total}")
    print(f"\nPath Distribution:")
    print(f"  Trivial (cheap model): {trivial} ({trivial/total:.1%})")
    print(f"  Easy (GNN direct): {easy} ({easy/total:.1%})")
    print(f"  Hard (conductor): {hard} ({hard/total:.1%})")

    print(f"\nRouting Differences:")
    print(f"  Queries where conductor would differ: {would_differ} ({would_differ/total:.1%})")
    print(f"  Queries with same routing: {total - would_differ} ({(total-would_differ)/total:.1%})")

    # GNN utilization
    baseline_gnn_calls = sum(r.gnn_calls_baseline for r in results)
    conductor_gnn_calls = sum(r.gnn_calls_conductor for r in results)
    print(f"\nGNN Utilization:")
    print(f"  Baseline total calls: {baseline_gnn_calls}")
    print(f"  Conductor total calls: {conductor_gnn_calls}")
    print(f"  Conductor overhead: {(conductor_gnn_calls/baseline_gnn_calls - 1):.1%} more calls")

    # Confidence analysis
    confidences = [r.baseline_confidence for r in results]
    print(f"\nGNN Confidence Analysis:")
    print(f"  Mean: {np.mean(confidences):.3f}")
    print(f"  Std: {np.std(confidences):.3f}")
    print(f"  Range: [{min(confidences):.3f}, {max(confidences):.3f}]")

    # Model distribution
    print(f"\nBaseline Model Selection:")
    model_counts = {}
    for r in results:
        model_counts[r.baseline_model] = model_counts.get(r.baseline_model, 0) + 1
    for model, count in sorted(model_counts.items(), key=lambda x: x[1], reverse=True):
        print(f"  {model}: {count} ({count/total:.1%})")

    # Per-query details (for small sets)
    if len(results) <= 10:
        print("\n" + "-" * 80)
        print("PER-QUERY ANALYSIS")
        print("-" * 80)

        for r in results:
            print(f"\nQuery {r.query_id}: {r.query[:60]}...")
            print(f"  GNN prediction: {r.baseline_model} (conf={r.baseline_confidence:.2f})")
            print(f"  Conductor path: {r.conductor_path}")
            print(f"  Would differ: {'Yes' if r.conductor_would_differ else 'No'}")
            print(f"  Reason: {r.conductor_reason}")

    print("\n" + "=" * 80)


def main():
    parser = argparse.ArgumentParser(
        description="Analyze conductor vs baseline routing decisions"
    )
    parser.add_argument(
        "--checkpoint",
        type=str,
        required=True,
        help="Path to GNN checkpoint",
    )
    parser.add_argument(
        "--queries",
        type=str,
        nargs="*",
        help="Specific queries to analyze",
    )
    parser.add_argument(
        "--random-samples",
        type=int,
        default=20,
        help="Number of random embedding samples (if no queries provided)",
    )

    args = parser.parse_args()

    # Import here after path setup
    try:
        from embergraphrouter.inference.eb1_router import EB1Router
    except ImportError:
        from embergraphrouter.inference import EB1Router

    from ember.takumi.eb1_gnn_conductor import RoutingThresholds

    print("Loading GNN router...")
    router = EB1Router.from_registry(args.checkpoint, registry_root="embergraphrouter/registry")
    thresholds = RoutingThresholds()
    print(f"✓ Loaded. LLM pool: {router.llm_names}")

    results = []

    if args.queries:
        # Use provided queries
        for i, query in enumerate(args.queries):
            result = analyze_single_query(query, i, router, thresholds)
            results.append(result)
    else:
        # Use random embeddings for analysis
        print(f"\nUsing {args.random_samples} random embedding samples...")
        np.random.seed(42)

        for i in range(args.random_samples):
            # Generate random embedding
            embedding = np.random.randn(512).astype(np.float32)

            # Create a synthetic routing result
            routing_result = router.route_from_embedding(embedding)

            # Compute signals manually
            from ember.takumi.eb1_gnn_conductor import RoutingSignals

            scores = list(routing_result.scores.values())
            mean_logit = sum(scores) / len(scores)
            sorted_scores = sorted(scores, reverse=True)
            logit_gap = sorted_scores[0] - sorted_scores[1]

            signals = RoutingSignals(
                confidence=routing_result.confidence,
                entropy=routing_result.entropy,
                difficulty=1.0 - routing_result.confidence,
                mean_logit=mean_logit,
                logit_gap=logit_gap,
                chosen_model=routing_result.best_model,
                chosen_model_idx=routing_result.best_model_idx,
                probabilities=routing_result.probabilities,
            )

            tier = signals.classify_with_thresholds(thresholds)

            # Create comparison result
            if tier.value == "trivial":
                would_differ = True
                reason = "High confidence + low entropy → cheap model"
                gnn_calls = 1
            elif tier.value == "easy":
                would_differ = False
                reason = f"Trust GNN: {routing_result.best_model}"
                gnn_calls = 1
            else:
                would_differ = True
                reason = "Low confidence → conductor orchestration"
                gnn_calls = 3

            result = ComparisonResult(
                query=f"[Random embedding #{i}]",
                query_id=i,
                baseline_model=routing_result.best_model,
                baseline_confidence=routing_result.confidence,
                baseline_entropy=routing_result.entropy,
                baseline_tier=tier.value,
                conductor_path=tier.value,
                conductor_would_differ=would_differ,
                conductor_reason=reason,
                gnn_calls_baseline=1,
                gnn_calls_conductor=gnn_calls,
            )
            results.append(result)

    print_comparison_report(results)


if __name__ == "__main__":
    main()
