"""
Example of multi-benchmark optimization using weighted benchmarks.

This script demonstrates how to use the optimization system with both
SimpleQA and BrowseComp benchmarks with custom weights.
"""

import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict


# Print current directory and python path for debugging
print(f"Current directory: {os.getcwd()}")
print(f"Python path: {sys.path}")

# Add appropriate paths
sys.path.insert(0, str(Path(__file__).parent.parent.resolve()))

try:
    # Try to import from the local module structure
    from src.local_deep_research.benchmarks.optimization.optuna_optimizer import (
        optimize_for_quality,
        optimize_for_speed,
        optimize_parameters,
    )

    print("Successfully imported using src.local_deep_research path")
except ImportError:
    print("First import attempt failed, trying with direct import...")
    try:
        # Try to import directly
        from local_deep_research.benchmarks.optimization.optuna_optimizer import (
            optimize_for_quality,
            optimize_for_speed,
            optimize_parameters,
        )

        print("Successfully imported using local_deep_research path")
    except ImportError as e:
        print(f"Import error: {e}")
        print("Creating simulation functions for demonstration only...")

        # Create simulation functions if imports fail
        def optimize_parameters(*args, **kwargs):
            benchmark_weights = kwargs.get(
                "benchmark_weights", {"simpleqa": 1.0}
            )
            print(
                f"SIMULATION: optimize_parameters called with benchmark_weights={benchmark_weights}"
            )

            # Return different results based on the benchmark weights
            if (
                "browsecomp" in benchmark_weights
                and benchmark_weights["browsecomp"] >= 1.0
            ):
                # BrowseComp only
                return {
                    "iterations": 4,
                    "questions_per_iteration": 5,
                    "search_strategy": "parallel",
                }, 0.78
            if (
                "browsecomp" in benchmark_weights
                and benchmark_weights["browsecomp"] > 0
            ):
                # Mixed weights
                return {
                    "iterations": 2,
                    "questions_per_iteration": 2,
                    "search_strategy": "iterdrag",
                }, 0.81
            # SimpleQA only (default)
            return {
                "iterations": 3,
                "questions_per_iteration": 2,
                "search_strategy": "standard",
            }, 0.75

        def optimize_for_quality(*args, **kwargs):
            benchmark_weights = kwargs.get(
                "benchmark_weights", {"simpleqa": 1.0}
            )
            print(
                f"SIMULATION: optimize_for_quality called with benchmark_weights={benchmark_weights}"
            )
            return {
                "iterations": 4,
                "questions_per_iteration": 1,
                "search_strategy": "iterdrag",
            }, 0.85

        def optimize_for_speed(*args, **kwargs):
            benchmark_weights = kwargs.get(
                "benchmark_weights", {"simpleqa": 1.0}
            )
            print(
                f"SIMULATION: optimize_for_speed called with benchmark_weights={benchmark_weights}"
            )
            return {
                "iterations": 2,
                "questions_per_iteration": 2,
                "search_strategy": "rapid",
            }, 0.67


# Loguru automatically handles logging configuration


def print_optimization_results(params: Dict[str, Any], score: float):
    """Print optimization results in a nicely formatted way."""
    print("\n" + "=" * 50)
    print(" OPTIMIZATION RESULTS ")
    print("=" * 50)
    print(f"SCORE: {score:.4f}")
    print("\nBest Parameters:")
    for param, value in params.items():
        print(f"  {param}: {value}")
    print("=" * 50 + "\n")


def main():
    """Run the multi-benchmark optimization examples."""
    # Create a timestamp-based directory for results
    from datetime import timezone

    timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
    output_dir = f"optimization_demo_{timestamp}"
    os.makedirs(output_dir, exist_ok=True)

    # Research query for optimization examples
    query = "Recent advancements in renewable energy"

    # Example 1: SimpleQA only (default)
    print("\n🔍 Running optimization with SimpleQA benchmark only...")
    params1, score1 = optimize_parameters(
        query=query,
        n_trials=3,  # Using a small number for quick demonstration
        output_dir=str(Path(output_dir) / "simpleqa_only"),
    )
    print_optimization_results(params1, score1)

    # Example 2: BrowseComp only
    print("\n🔍 Running optimization with BrowseComp benchmark only...")
    params2, score2 = optimize_parameters(
        query=query,
        n_trials=3,  # Using a small number for quick demonstration
        output_dir=str(Path(output_dir) / "browsecomp_only"),
        benchmark_weights={"browsecomp": 1.0},
    )
    print_optimization_results(params2, score2)

    # Example 3: 60/40 weighted combination (SimpleQA/BrowseComp)
    print("\n🔍 Running optimization with 60% SimpleQA and 40% BrowseComp...")
    params3, score3 = optimize_parameters(
        query=query,
        n_trials=5,  # Using a small number for quick demonstration
        output_dir=str(Path(output_dir) / "weighted_combination"),
        benchmark_weights={
            "simpleqa": 0.6,  # 60% weight for SimpleQA
            "browsecomp": 0.4,  # 40% weight for BrowseComp
        },
    )
    print_optimization_results(params3, score3)

    # Example 4: Quality-focused with both benchmarks
    print("\n🔍 Running quality-focused optimization with both benchmarks...")
    params4, score4 = optimize_for_quality(
        query=query,
        n_trials=3,
        output_dir=str(Path(output_dir) / "quality_focused"),
        benchmark_weights={"simpleqa": 0.6, "browsecomp": 0.4},
    )
    print_optimization_results(params4, score4)

    # Example 5: Speed-focused with both benchmarks
    print("\n🔍 Running speed-focused optimization with both benchmarks...")
    params5, score5 = optimize_for_speed(
        query=query,
        n_trials=3,
        output_dir=str(Path(output_dir) / "speed_focused"),
        benchmark_weights={"simpleqa": 0.5, "browsecomp": 0.5},
    )
    print_optimization_results(params5, score5)

    print(f"\nAll optimization results saved to: {output_dir}")
    print("View the results directory for detailed logs and visualizations.")


if __name__ == "__main__":
    main()