#!/usr/bin/env python3
"""Deep error analysis for the needle-weighted GNN router.

This script performs comprehensive analysis of routing errors to understand:
1. Which queries are consistently misrouted
2. What patterns exist in the errors (task type, LLM characteristics, etc.)
3. Why there's a gap between current performance and theoretical ceiling

The needle-weighted model achieved 69.14% frac_oracle. This script investigates
the remaining 30.86% of queries where routing was suboptimal.

Usage:
    python scripts/analyze_needle_model_errors.py \
        --config configs/v7_gat4_needle_weighted.yaml \
        --model model_path/v7_gat4_needle_weighted_best.pth

Output:
    - Detailed error breakdown by category
    - Confusion matrix: which LLMs are confused with which
    - Analysis of "needle query" routing accuracy
    - Identification of systematic failure modes
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import NamedTuple

import numpy as np
import pandas as pd
import torch
import yaml

# Add parent directories to path
_SCRIPT_DIR = Path(__file__).resolve().parent
_MODULE_DIR = _SCRIPT_DIR.parent
if str(_MODULE_DIR.parent) not in sys.path:
    sys.path.insert(0, str(_MODULE_DIR.parent))

from embergraphrouter.model.multi_task_graph_router import GraphRouterTrainer

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


class QueryResult(NamedTuple):
    """Result for a single query's routing decision."""

    query_idx: int
    task_id: str
    model_choice: int
    oracle_choice: int
    is_correct: bool
    model_confidence: float
    oracle_reward: float
    chosen_reward: float
    regret: float
    num_correct_llms: int  # How many LLMs got it right
    performance_vector: tuple[float, ...]  # Per-LLM performance


@dataclass
class ErrorAnalysis:
    """Comprehensive error analysis results."""

    total_queries: int = 0
    correct_queries: int = 0
    frac_oracle: float = 0.0

    # By task
    task_errors: dict[str, dict] = field(default_factory=dict)

    # By number of correct LLMs (needle analysis)
    by_num_correct: dict[int, dict] = field(default_factory=dict)

    # Confusion matrix: model_choice -> oracle_choice -> count
    confusion: dict[int, Counter] = field(default_factory=lambda: defaultdict(Counter))

    # Most confused pairs
    confused_pairs: list[tuple[int, int, int]] = field(default_factory=list)

    # Queries grouped by error pattern
    errors_by_pattern: dict[str, list[int]] = field(default_factory=dict)


def load_model_and_data(
    config_path: Path,
    model_path: Path | None = None,
) -> tuple[torch.nn.Module, GraphRouterTrainer]:
    """Load trained model and data.

    Args:
        config_path: Path to config YAML.
        model_path: Optional override for model checkpoint path.

    Returns:
        Tuple of (model, trainer) where trainer has loaded data.
    """
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    # Resolve relative paths
    for key in (
        "unified_qa_data_path",
        "llm_description_path",
        "saved_router_data_path",
        "llm_embedding_path",
        "model_path",
    ):
        if key in config and not Path(config[key]).is_absolute():
            config[key] = str(_MODULE_DIR / config[key])

    # Override model path if provided
    if model_path:
        config["model_path"] = str(model_path)

    logger.info("Loading data and model from config: %s", config_path)

    # Use a minimal wandb stub
    class WandbStub:
        run = None

        def log(self, *args, **kwargs):
            pass

    trainer = GraphRouterTrainer(
        router_data_path=config["saved_router_data_path"],
        llm_path=config["llm_description_path"],
        llm_embedding_path=config["llm_embedding_path"],
        config=config,
        wandb=WandbStub(),
    )

    # Load model weights
    model_weights_path = Path(config["model_path"])
    if model_weights_path.exists():
        logger.info("Loading model weights from %s", model_weights_path)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        state_dict = torch.load(model_weights_path, map_location=device, weights_only=True)
        trainer.gnn.model.load_state_dict(state_dict)
        trainer.gnn.model.eval()
    else:
        logger.warning("Model weights not found at %s", model_weights_path)

    return trainer.gnn.model, trainer


def run_inference(
    model: torch.nn.Module,
    trainer: GraphRouterTrainer,
) -> list[QueryResult]:
    """Run inference on test set and collect per-query results.

    Args:
        model: Trained GNN model.
        trainer: Trainer with loaded data.

    Returns:
        List of QueryResult for each test query.
    """
    model.eval()
    device = next(model.parameters()).device

    # Get test data
    test_data = trainer.form_data.create_data(
        trainer.query_embeddings,
        trainer.task_embeddings,
        trainer.llm_embeddings,
        trainer.test_idx,
        trainer.perf_list,
        trainer.cost_list,
        trainer.latency_list,
        trainer.num_queries,
        trainer.num_llms,
    )
    test_data = test_data.to(device)

    results = []
    with torch.no_grad():
        # Forward pass
        logits = model(
            test_data.query_x,
            test_data.llm_x,
            test_data.edge_index,
            test_data.edge_attr,
            test_data.combined_edge,
        )

        # Get mask for test edges
        mask = test_data.test_mask if hasattr(test_data, "test_mask") else test_data.train_mask
        logits = logits[mask]

        # Reshape to (num_queries, num_llms)
        num_llms = trainer.num_llms
        scores = logits.view(-1, num_llms)

        # Get rewards (ground truth performance)
        rewards = test_data.combined_edge[mask, 1].view(-1, num_llms)

        # Compute softmax probabilities for confidence
        probs = torch.softmax(scores, dim=1)

        # Per-query analysis
        for i in range(scores.size(0)):
            query_scores = scores[i]
            query_rewards = rewards[i]
            query_probs = probs[i]

            model_choice = query_scores.argmax().item()
            oracle_choice = query_rewards.argmax().item()

            model_confidence = query_probs[model_choice].item()
            oracle_reward = query_rewards[oracle_choice].item()
            chosen_reward = query_rewards[model_choice].item()

            # Number of LLMs that got this query correct (effect > 0.5)
            num_correct = (query_rewards > 0.5).sum().item()

            # Get test query index
            test_query_idx = trainer.test_idx[i] if i < len(trainer.test_idx) else i
            task_id = trainer.query_task_labels[test_query_idx]

            result = QueryResult(
                query_idx=test_query_idx,
                task_id=task_id,
                model_choice=model_choice,
                oracle_choice=oracle_choice,
                is_correct=(model_choice == oracle_choice),
                model_confidence=model_confidence,
                oracle_reward=oracle_reward,
                chosen_reward=chosen_reward,
                regret=oracle_reward - chosen_reward,
                num_correct_llms=num_correct,
                performance_vector=tuple(query_rewards.cpu().numpy()),
            )
            results.append(result)

    return results


def analyze_errors(
    results: list[QueryResult],
    llm_names: list[str],
) -> ErrorAnalysis:
    """Perform comprehensive error analysis.

    Args:
        results: Per-query results from inference.
        llm_names: Names of LLMs.

    Returns:
        ErrorAnalysis with detailed breakdown.
    """
    analysis = ErrorAnalysis()
    analysis.total_queries = len(results)
    analysis.correct_queries = sum(1 for r in results if r.is_correct)
    analysis.frac_oracle = analysis.correct_queries / analysis.total_queries

    # Group by task
    task_results = defaultdict(list)
    for r in results:
        task_results[r.task_id].append(r)

    for task_id, task_res in task_results.items():
        correct = sum(1 for r in task_res if r.is_correct)
        total = len(task_res)
        avg_regret = np.mean([r.regret for r in task_res])
        avg_confidence = np.mean([r.model_confidence for r in task_res])

        analysis.task_errors[task_id] = {
            "total": total,
            "correct": correct,
            "accuracy": correct / total,
            "avg_regret": avg_regret,
            "avg_confidence": avg_confidence,
        }

    # Group by number of correct LLMs (needle analysis)
    num_correct_results = defaultdict(list)
    for r in results:
        num_correct_results[r.num_correct_llms].append(r)

    for num_correct, group_res in num_correct_results.items():
        correct = sum(1 for r in group_res if r.is_correct)
        total = len(group_res)
        avg_regret = np.mean([r.regret for r in group_res])
        avg_confidence = np.mean([r.model_confidence for r in group_res])

        analysis.by_num_correct[num_correct] = {
            "total": total,
            "correct": correct,
            "accuracy": correct / total if total > 0 else 0,
            "pct_of_data": total / len(results) * 100,
            "avg_regret": avg_regret,
            "avg_confidence": avg_confidence,
        }

    # Build confusion matrix
    for r in results:
        if not r.is_correct:
            analysis.confusion[r.model_choice][r.oracle_choice] += 1

    # Find most confused pairs
    pairs = []
    for model_choice, oracle_counts in analysis.confusion.items():
        for oracle_choice, count in oracle_counts.items():
            pairs.append((model_choice, oracle_choice, count))

    analysis.confused_pairs = sorted(pairs, key=lambda x: -x[2])[:10]

    # Categorize errors by pattern
    # Pattern 1: Confident but wrong
    confident_wrong = [
        r.query_idx for r in results
        if not r.is_correct and r.model_confidence > 0.5
    ]
    analysis.errors_by_pattern["confident_wrong"] = confident_wrong

    # Pattern 2: Needle queries (only 1 correct LLM) that were wrong
    needle_wrong = [
        r.query_idx for r in results
        if not r.is_correct and r.num_correct_llms == 1
    ]
    analysis.errors_by_pattern["needle_wrong"] = needle_wrong

    # Pattern 3: Zero-oracle queries (no LLM correct)
    zero_oracle = [
        r.query_idx for r in results
        if r.num_correct_llms == 0
    ]
    analysis.errors_by_pattern["zero_oracle"] = zero_oracle

    # Pattern 4: High regret (chose very wrong LLM)
    high_regret = [
        r.query_idx for r in results
        if r.regret > 0.5
    ]
    analysis.errors_by_pattern["high_regret"] = high_regret

    return analysis


def print_analysis_report(
    analysis: ErrorAnalysis,
    llm_names: list[str],
) -> None:
    """Print formatted analysis report."""
    print("\n" + "=" * 80)
    print("NEEDLE-WEIGHTED MODEL ERROR ANALYSIS")
    print("=" * 80)

    print(f"\nOverall: {analysis.correct_queries}/{analysis.total_queries} "
          f"({analysis.frac_oracle:.1%}) queries correctly routed")

    # By number of correct LLMs (the core insight)
    print("\n" + "-" * 80)
    print("ROUTING ACCURACY BY NUMBER OF CORRECT LLMs (Needle Analysis)")
    print("-" * 80)
    print(f"{'#Correct LLMs':<15} {'Count':>8} {'%Data':>8} {'Accuracy':>10} "
          f"{'AvgRegret':>10} {'AvgConf':>10}")
    print("-" * 80)

    for num_correct in sorted(analysis.by_num_correct.keys()):
        stats = analysis.by_num_correct[num_correct]
        label = f"{num_correct}"
        if num_correct == 0:
            label += " (zero-oracle)"
        elif num_correct == 1:
            label += " (NEEDLE)"
        elif num_correct == 7:
            label += " (all correct)"

        print(f"{label:<15} {stats['total']:>8} {stats['pct_of_data']:>7.1f}% "
              f"{stats['accuracy']:>9.1%} {stats['avg_regret']:>10.3f} "
              f"{stats['avg_confidence']:>10.3f}")

    # By task
    print("\n" + "-" * 80)
    print("ROUTING ACCURACY BY TASK")
    print("-" * 80)
    print(f"{'Task':<25} {'Count':>8} {'Accuracy':>10} {'AvgRegret':>10} {'AvgConf':>10}")
    print("-" * 80)

    for task_id, stats in sorted(
        analysis.task_errors.items(),
        key=lambda x: -x[1]["accuracy"],
    ):
        task_name = task_id[:24]
        print(f"{task_name:<25} {stats['total']:>8} {stats['accuracy']:>9.1%} "
              f"{stats['avg_regret']:>10.3f} {stats['avg_confidence']:>10.3f}")

    # Confusion matrix (most confused pairs)
    print("\n" + "-" * 80)
    print("MOST CONFUSED LLM PAIRS (model_choice -> oracle_choice)")
    print("-" * 80)
    for model_choice, oracle_choice, count in analysis.confused_pairs:
        model_name = llm_names[model_choice] if model_choice < len(llm_names) else str(model_choice)
        oracle_name = llm_names[oracle_choice] if oracle_choice < len(llm_names) else str(oracle_choice)
        print(f"  {model_name:<25} -> {oracle_name:<25} ({count:>4} times)")

    # Error patterns
    print("\n" + "-" * 80)
    print("ERROR PATTERNS")
    print("-" * 80)
    for pattern, query_idxs in analysis.errors_by_pattern.items():
        pct = len(query_idxs) / analysis.total_queries * 100
        print(f"  {pattern:<20}: {len(query_idxs):>5} queries ({pct:.1f}%)")

    # Key insights
    print("\n" + "=" * 80)
    print("KEY INSIGHTS")
    print("=" * 80)

    # Needle query analysis
    needle_stats = analysis.by_num_correct.get(1, {})
    if needle_stats:
        print(f"\n1. NEEDLE QUERIES (1 correct LLM):")
        print(f"   - These are {needle_stats.get('pct_of_data', 0):.1f}% of test data")
        print(f"   - Routing accuracy: {needle_stats.get('accuracy', 0):.1%}")
        print(f"   - This is where routing matters MOST")

    # Zero-oracle analysis
    zero_stats = analysis.by_num_correct.get(0, {})
    if zero_stats:
        print(f"\n2. ZERO-ORACLE QUERIES (no LLM correct):")
        print(f"   - These are {zero_stats.get('pct_of_data', 0):.1f}% of test data")
        print(f"   - Cannot improve on these with better routing")
        print(f"   - Consider: network-of-networks approach")

    # All-correct analysis
    all_correct_stats = analysis.by_num_correct.get(7, {})
    if all_correct_stats:
        print(f"\n3. ALL-CORRECT QUERIES (all 7 LLMs correct):")
        print(f"   - These are {all_correct_stats.get('pct_of_data', 0):.1f}% of test data")
        print(f"   - Routing doesn't matter here (any choice is optimal)")

    # Ceiling calculation
    print(f"\n4. THEORETICAL CEILING:")
    non_zero_oracle = [
        stats for num, stats in analysis.by_num_correct.items()
        if num > 0
    ]
    if non_zero_oracle:
        total_routable = sum(s["total"] for s in non_zero_oracle)
        routable_pct = total_routable / analysis.total_queries * 100
        print(f"   - Routable queries (>=1 LLM correct): {routable_pct:.1f}%")
        print(f"   - Perfect routing on these would give: {routable_pct:.1f}% frac_oracle")

        # Current performance gap
        current = analysis.frac_oracle * 100
        ceiling = routable_pct
        gap = ceiling - current
        print(f"   - Current: {current:.1f}%, Gap to ceiling: {gap:.1f}%")


def save_analysis_results(
    analysis: ErrorAnalysis,
    output_path: Path,
) -> None:
    """Save analysis results to JSON."""
    results = {
        "summary": {
            "total_queries": analysis.total_queries,
            "correct_queries": analysis.correct_queries,
            "frac_oracle": analysis.frac_oracle,
        },
        "by_num_correct_llms": {
            str(k): v for k, v in analysis.by_num_correct.items()
        },
        "by_task": analysis.task_errors,
        "confused_pairs": [
            {"model": m, "oracle": o, "count": c}
            for m, o, c in analysis.confused_pairs
        ],
        "error_pattern_counts": {
            pattern: len(idxs)
            for pattern, idxs in analysis.errors_by_pattern.items()
        },
    }

    with open(output_path, "w") as f:
        json.dump(results, f, indent=2)

    logger.info("Saved analysis results to %s", output_path)


def main():
    parser = argparse.ArgumentParser(description="Analyze needle-weighted model errors")
    parser.add_argument(
        "--config",
        type=Path,
        default=Path(_MODULE_DIR / "configs" / "v7_gat4_needle_weighted.yaml"),
        help="Path to config YAML",
    )
    parser.add_argument(
        "--model",
        type=Path,
        default=None,
        help="Path to model checkpoint (overrides config)",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path(_MODULE_DIR / "results" / "needle_model_error_analysis.json"),
        help="Output path for JSON results",
    )
    args = parser.parse_args()

    # Load model and data
    model, trainer = load_model_and_data(args.config, args.model)

    # Run inference
    logger.info("Running inference on test set...")
    results = run_inference(model, trainer)
    logger.info("Collected results for %d queries", len(results))

    # Analyze errors
    logger.info("Analyzing errors...")
    analysis = analyze_errors(results, trainer.llm_names)

    # Print report
    print_analysis_report(analysis, trainer.llm_names)

    # Save results
    args.output.parent.mkdir(parents=True, exist_ok=True)
    save_analysis_results(analysis, args.output)


if __name__ == "__main__":
    main()
