#!/usr/bin/env python3
"""Analyze v7_gat4 experiments comparing hard query optimization strategies.

This script compares three experiments:
1. v7_gat4_optimal: Baseline 4-layer GAT without hard query optimization
2. v7_gat4_hard_mining: Hard negative mining (batch-level weighting)
3. v7_gat4_needle_weighted: Needle query weighting (loss-level weighting)

The key insight is that 12% of queries are "needle queries" where only 1 LLM succeeds.
These queries are where routing adds the most value but are hardest to learn.

Usage:
    python scripts/analyze_hard_query_experiments.py
"""

from __future__ import annotations

import json
import logging
import sys
from pathlib import Path
from typing import TypedDict

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


class ExperimentResult(TypedDict):
    """Results from a single experiment."""

    name: str
    frac_oracle: float
    frac_oracle_q0: float  # Easy queries (lowest difficulty quartile)
    frac_oracle_q1: float
    frac_oracle_q2: float
    frac_oracle_q3: float  # Hard queries (highest difficulty quartile)
    avg_confidence: float
    avg_uncertainty: float


def load_experiment_results(results_dir: Path) -> dict[str, ExperimentResult]:
    """Load experiment results from JSON files or parse from logs.

    Args:
        results_dir: Directory containing experiment result files.

    Returns:
        Dictionary mapping experiment name to results.
    """
    results = {}

    experiments = [
        "v7_gat4_optimal",
        "v7_gat4_hard_mining",
        "v7_gat4_needle_weighted",
    ]

    for exp_name in experiments:
        json_path = results_dir / f"{exp_name}_results.json"
        if json_path.exists():
            with open(json_path, "r") as f:
                results[exp_name] = json.load(f)
        else:
            logger.warning("No results file for %s at %s", exp_name, json_path)

    return results


def print_comparison_table(results: dict[str, ExperimentResult]) -> None:
    """Print a formatted comparison table.

    Args:
        results: Dictionary of experiment results.
    """
    if not results:
        logger.error("No results to compare")
        return

    # Header
    print("\n" + "=" * 80)
    print("v7_gat4 Hard Query Optimization Experiment Comparison")
    print("=" * 80)
    print()

    # Table header
    print(f"{'Metric':<25} ", end="")
    for name in sorted(results.keys()):
        short_name = name.replace("v7_gat4_", "")
        print(f"{short_name:>15} ", end="")
    print()
    print("-" * 80)

    # Metrics
    metrics = [
        ("frac_oracle", "Overall"),
        ("frac_oracle_q0", "Easy (q0)"),
        ("frac_oracle_q1", "Medium-Easy (q1)"),
        ("frac_oracle_q2", "Medium-Hard (q2)"),
        ("frac_oracle_q3", "Hard (q3)"),
        ("avg_confidence", "Avg Confidence"),
        ("avg_uncertainty", "Avg Uncertainty"),
    ]

    for metric_key, metric_label in metrics:
        print(f"{metric_label:<25} ", end="")
        for name in sorted(results.keys()):
            value = results[name].get(metric_key, float("nan"))
            if "frac" in metric_key:
                print(f"{value:>14.1%} ", end="")
            else:
                print(f"{value:>15.3f} ", end="")
        print()

    print()


def compute_relative_improvements(results: dict[str, ExperimentResult]) -> None:
    """Compute and print relative improvements over baseline.

    Args:
        results: Dictionary of experiment results.
    """
    if "v7_gat4_optimal" not in results:
        logger.warning("No baseline (v7_gat4_optimal) found for comparison")
        return

    baseline = results["v7_gat4_optimal"]

    print("\n" + "-" * 80)
    print("Relative Improvement over Baseline (v7_gat4_optimal)")
    print("-" * 80)

    for name, exp_results in results.items():
        if name == "v7_gat4_optimal":
            continue

        short_name = name.replace("v7_gat4_", "")
        print(f"\n{short_name}:")

        # Overall
        if baseline["frac_oracle"] > 0:
            delta = exp_results["frac_oracle"] - baseline["frac_oracle"]
            pct = delta / baseline["frac_oracle"] * 100
            print(f"  Overall: {delta:+.1%} ({pct:+.1f}% relative)")

        # Hard queries (q3) - most important
        if baseline["frac_oracle_q3"] > 0:
            delta = exp_results["frac_oracle_q3"] - baseline["frac_oracle_q3"]
            pct = delta / baseline["frac_oracle_q3"] * 100
            print(f"  Hard queries (q3): {delta:+.1%} ({pct:+.1f}% relative)")

        # Easy queries (q0) - check for regression
        if baseline["frac_oracle_q0"] > 0:
            delta = exp_results["frac_oracle_q0"] - baseline["frac_oracle_q0"]
            pct = delta / baseline["frac_oracle_q0"] * 100
            print(f"  Easy queries (q0): {delta:+.1%} ({pct:+.1f}% relative)")


def analyze_needle_query_impact(results: dict[str, ExperimentResult]) -> None:
    """Analyze the specific impact on needle queries.

    Needle queries have only 1 LLM that succeeds (12% of data).
    These queries have the highest potential for routing value.

    Args:
        results: Dictionary of experiment results.
    """
    print("\n" + "=" * 80)
    print("Analysis: Needle Query Impact")
    print("=" * 80)
    print()
    print("Background:")
    print("  - Needle queries: queries where only 1 LLM succeeds (~12% of data)")
    print("  - These are where routing adds the most value but are hardest to learn")
    print("  - Needle weighting gives these queries 7x loss weight")
    print()

    if "v7_gat4_needle_weighted" in results and "v7_gat4_optimal" in results:
        baseline_q3 = results["v7_gat4_optimal"]["frac_oracle_q3"]
        needle_q3 = results["v7_gat4_needle_weighted"]["frac_oracle_q3"]

        print(f"Hard query performance:")
        print(f"  Baseline:        {baseline_q3:.1%}")
        print(f"  Needle weighted: {needle_q3:.1%}")
        print(f"  Improvement:     {needle_q3 - baseline_q3:+.1%}")
        print()

        if needle_q3 > baseline_q3 + 0.02:
            print("Conclusion: Needle weighting SIGNIFICANTLY improves hard query routing")
        elif needle_q3 > baseline_q3:
            print("Conclusion: Needle weighting provides MODEST improvement on hard queries")
        else:
            print("Conclusion: Needle weighting did NOT improve hard query routing")
            print("            Consider alternative approaches (curriculum learning, etc.)")


def main() -> int:
    """Main entry point."""
    # Look for results in standard locations
    repo_root = Path(__file__).parent.parent
    results_dir = repo_root / "results"
    results_dir.mkdir(parents=True, exist_ok=True)

    results = load_experiment_results(results_dir)

    if not results:
        logger.error("No experiment results found in %s", results_dir)
        return 1

    print_comparison_table(results)
    compute_relative_improvements(results)
    analyze_needle_query_impact(results)

    return 0


if __name__ == "__main__":
    sys.exit(main())
