#!/usr/bin/env python3
# This script should be run from the project root directory using:
# cd /path/to/local-deep-research
# python -m examples.optimization.strategy_benchmark_plan
"""
Strategy Benchmark Plan - Comprehensive Optuna-based optimization for search strategies

This benchmark specifically focuses on comparing the iterdrag and source_based strategies
with 500 examples per experiment to ensure statistically significant results.
"""

import json
import os
import random
import sys
import time
from datetime import datetime, UTC
from pathlib import Path
from typing import Any, Dict, Tuple

from loguru import logger

# Add the src directory to the Python path before local imports
project_root = str(Path(__file__).parent.parent.parent.resolve())
sys.path.insert(0, str(Path(project_root) / "src"))

# Now we can import from the local project
from local_deep_research.benchmarks.optimization.optuna_optimizer import (  # noqa: E402
    OptunaOptimizer,
)

# Logger is already imported from loguru at the top

# Number of examples to use in each benchmark experiment
NUM_EXAMPLES = 500


def progress_callback(trial_num, total_trials, data):
    """Progress callback for optimization"""
    print(f"Progress: {trial_num}/{total_trials} - {data}")


def run_strategy_comparison():
    """
    Run a comprehensive comparison between iterdrag and source_based strategies.
    Uses a large sample size (500 examples) for statistical significance.
    """
    # Verify LLM and search database settings before proceeding
    try:
        from local_deep_research.config.llm_config import get_llm
        from local_deep_research.config.search_config import get_search
        from local_deep_research.utilities.db_utils import get_db_setting

        # Try to initialize LLM and search engine to check configuration
        llm = get_llm()
        search = get_search()

        # Get relevant DB settings
        try:
            iterations = get_db_setting("search.iterations") or 3
            questions_per_iteration = (
                get_db_setting("search.questions_per_iteration") or 3
            )
        except Exception as e:
            logger.warning(f"Error getting DB settings: {e}")
            iterations = 3
            questions_per_iteration = 3

        logger.info("Successfully connected to database")
        logger.info(f"Using LLM: {llm.__class__.__name__}")
        logger.info(f"Using search engine: {search.__class__.__name__}")
        logger.info(f"Default iterations from DB: {iterations}")
        logger.info(
            f"Default questions per iteration from DB: {questions_per_iteration}"
        )
    except Exception as e:
        logger.exception("Error initializing LLM or search settings")
        logger.info("Please check your database configuration")
        return {"error": str(e)}
    timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")
    base_output_dir = f"strategy_benchmark_results_{timestamp}"
    os.makedirs(base_output_dir, exist_ok=True)

    # Define test query
    query = "What are the latest developments in fusion energy research?"

    # Track execution stats
    execution_stats = {"start_time": time.time(), "experiments": []}

    # Define parameter space specific to strategy comparison
    strategy_param_space = {
        "search_strategy": {
            "type": "categorical",
            "choices": ["iterdrag", "source_based"],
        },
        "iterations": {
            "type": "int",
            "low": 1,
            "high": 3,
            "step": 1,
        },
        "questions_per_iteration": {
            "type": "int",
            "low": 1,
            "high": 5,
            "step": 1,
        },
        "max_results": {
            "type": "int",
            "low": 10,
            "high": 50,
            "step": 10,
        },
    }

    # Common settings for all experiments
    common_settings = {
        "query": query,
        "n_trials": 30,  # Optuna trials per experiment
        "n_jobs": 1,  # Run one job at a time for consistent resource measurement
        "timeout": 3600,  # 1 hour timeout per experiment
        "progress_callback": progress_callback,
    }

    # ====== EXPERIMENT 1: Quality-focused optimization ======
    logger.info("Starting quality-focused benchmark with 500 examples")
    quality_output_dir = str(Path(base_output_dir) / "quality_focused")
    Path(quality_output_dir).mkdir(parents=True, exist_ok=True)

    # Create optimizer for quality
    quality_optimizer = OptunaOptimizer(
        base_query=query,
        output_dir=quality_output_dir,
        n_trials=common_settings["n_trials"],
        timeout=common_settings["timeout"],
        n_jobs=common_settings["n_jobs"],
        progress_callback=common_settings["progress_callback"],
        study_name="strategy_quality_benchmark",
        optimization_metrics=["quality", "speed"],
        metric_weights={"quality": 0.9, "speed": 0.1},
        num_examples=NUM_EXAMPLES,  # Use 500 examples for robust evaluation
    )

    # Run quality optimization
    quality_start = time.time()
    best_quality_params, best_quality_score = quality_optimizer.optimize(
        strategy_param_space
    )
    quality_end = time.time()

    quality_result = {
        "experiment": "quality_focused",
        "best_params": best_quality_params,
        "best_score": best_quality_score,
        "duration_seconds": quality_end - quality_start,
    }
    execution_stats["experiments"].append(quality_result)

    # Log and save results
    logger.info(f"Quality benchmark complete: {best_quality_params}")
    logger.info(f"Best quality score: {best_quality_score}")
    logger.info(f"Duration: {quality_end - quality_start} seconds")

    with open(
        Path(quality_output_dir) / "results.json", "w", encoding="utf-8"
    ) as f:
        json.dump(quality_result, f, indent=2)

    # ====== EXPERIMENT 2: Speed-focused optimization ======
    logger.info("Starting speed-focused benchmark with 500 examples")
    speed_output_dir = str(Path(base_output_dir) / "speed_focused")
    Path(speed_output_dir).mkdir(parents=True, exist_ok=True)

    # Create optimizer for speed
    speed_optimizer = OptunaOptimizer(
        base_query=query,
        output_dir=speed_output_dir,
        n_trials=common_settings["n_trials"],
        timeout=common_settings["timeout"],
        n_jobs=common_settings["n_jobs"],
        progress_callback=common_settings["progress_callback"],
        study_name="strategy_speed_benchmark",
        optimization_metrics=["quality", "speed"],
        metric_weights={"quality": 0.2, "speed": 0.8},
        num_examples=NUM_EXAMPLES,  # Use 500 examples for robust evaluation
    )

    # Run speed optimization
    speed_start = time.time()
    best_speed_params, best_speed_score = speed_optimizer.optimize(
        strategy_param_space
    )
    speed_end = time.time()

    speed_result = {
        "experiment": "speed_focused",
        "best_params": best_speed_params,
        "best_score": best_speed_score,
        "duration_seconds": speed_end - speed_start,
    }
    execution_stats["experiments"].append(speed_result)

    # Log and save results
    logger.info(f"Speed benchmark complete: {best_speed_params}")
    logger.info(f"Best speed score: {best_speed_score}")
    logger.info(f"Duration: {speed_end - speed_start} seconds")

    with open(
        Path(speed_output_dir) / "results.json", "w", encoding="utf-8"
    ) as f:
        json.dump(speed_result, f, indent=2)

    # ====== EXPERIMENT 3: Balanced optimization ======
    logger.info("Starting balanced benchmark with 500 examples")
    balanced_output_dir = str(Path(base_output_dir) / "balanced")
    Path(balanced_output_dir).mkdir(parents=True, exist_ok=True)

    # Create optimizer for balanced approach
    balanced_optimizer = OptunaOptimizer(
        base_query=query,
        output_dir=balanced_output_dir,
        n_trials=common_settings["n_trials"],
        timeout=common_settings["timeout"],
        n_jobs=common_settings["n_jobs"],
        progress_callback=common_settings["progress_callback"],
        study_name="strategy_balanced_benchmark",
        optimization_metrics=["quality", "speed", "resource"],
        metric_weights={"quality": 0.4, "speed": 0.3, "resource": 0.3},
        num_examples=NUM_EXAMPLES,  # Use 500 examples for robust evaluation
    )

    # Run balanced optimization
    balanced_start = time.time()
    best_balanced_params, best_balanced_score = balanced_optimizer.optimize(
        strategy_param_space
    )
    balanced_end = time.time()

    balanced_result = {
        "experiment": "balanced",
        "best_params": best_balanced_params,
        "best_score": best_balanced_score,
        "duration_seconds": balanced_end - balanced_start,
    }
    execution_stats["experiments"].append(balanced_result)

    # Log and save results
    logger.info(f"Balanced benchmark complete: {best_balanced_params}")
    logger.info(f"Best balanced score: {best_balanced_score}")
    logger.info(f"Duration: {balanced_end - balanced_start} seconds")

    with open(
        Path(balanced_output_dir) / "results.json", "w", encoding="utf-8"
    ) as f:
        json.dump(balanced_result, f, indent=2)

    # ====== EXPERIMENT 4: Multi-Benchmark (SimpleQA + BrowseComp) ======
    logger.info("Starting multi-benchmark optimization with 500 examples")
    multi_output_dir = str(Path(base_output_dir) / "multi_benchmark")
    Path(multi_output_dir).mkdir(parents=True, exist_ok=True)

    # Create optimizer with multi-benchmark weights
    multi_optimizer = OptunaOptimizer(
        base_query=query,
        output_dir=multi_output_dir,
        n_trials=common_settings["n_trials"],
        timeout=common_settings["timeout"],
        n_jobs=common_settings["n_jobs"],
        progress_callback=common_settings["progress_callback"],
        study_name="strategy_multi_benchmark",
        optimization_metrics=["quality", "speed"],
        metric_weights={"quality": 0.6, "speed": 0.4},
        benchmark_weights={"simpleqa": 0.6, "browsecomp": 0.4},
        num_examples=NUM_EXAMPLES,  # Use 500 examples for robust evaluation
    )

    # Run multi-benchmark optimization
    multi_start = time.time()
    best_multi_params, best_multi_score = multi_optimizer.optimize(
        strategy_param_space
    )
    multi_end = time.time()

    multi_result = {
        "experiment": "multi_benchmark",
        "best_params": best_multi_params,
        "best_score": best_multi_score,
        "duration_seconds": multi_end - multi_start,
    }
    execution_stats["experiments"].append(multi_result)

    # Log and save results
    logger.info(f"Multi-benchmark complete: {best_multi_params}")
    logger.info(f"Best multi-benchmark score: {best_multi_score}")
    logger.info(f"Duration: {multi_end - multi_start} seconds")

    with open(
        Path(multi_output_dir) / "results.json", "w", encoding="utf-8"
    ) as f:
        json.dump(multi_result, f, indent=2)

    # ====== Save summary of all executions ======
    execution_stats["total_duration"] = (
        time.time() - execution_stats["start_time"]
    )
    execution_stats["timestamp"] = timestamp

    with open(
        Path(base_output_dir) / "summary.json", "w", encoding="utf-8"
    ) as f:
        json.dump(execution_stats, f, indent=2)

    # Generate summary report
    generate_summary_report(base_output_dir, execution_stats)

    return execution_stats


def generate_summary_report(base_dir, stats):
    """Generate a human-readable summary report of all benchmarks"""
    summary_text = f"""
# Strategy Benchmark Results Summary

## Overview
- **Date:** {datetime.fromtimestamp(stats["start_time"]).strftime("%Y-%m-%d %H:%M:%S")}
- **Total Duration:** {stats["total_duration"] / 3600:.2f} hours
- **Number of Examples per Experiment:** {NUM_EXAMPLES}

## Experiment Results

"""
    # Add detailed results for each experiment
    for exp in stats["experiments"]:
        summary_text += f"""### {exp["experiment"].replace("_", " ").title()}
- **Best Parameters:** {json.dumps(exp["best_params"], indent=2)}
- **Best Score:** {exp["best_score"]:.4f}
- **Duration:** {exp["duration_seconds"] / 60:.2f} minutes

"""

    summary_text += """
## Strategy Comparison

| Metric Focus | Best Strategy | Other Parameters | Score |
|--------------|--------------|------------------|-------|
"""

    for exp in stats["experiments"]:
        best_strategy = exp["best_params"].get("search_strategy", "unknown")
        other_params = {
            k: v
            for k, v in exp["best_params"].items()
            if k != "search_strategy"
        }
        summary_text += f"| {exp['experiment'].replace('_', ' ').title()} | {best_strategy} | {other_params} | {exp['best_score']:.4f} |\n"

    summary_text += """
## Analysis

This benchmark compared the performance of iterdrag and source_based strategies across different optimization goals:
- Quality-focused: Prioritizes result quality (90%) over speed (10%)
- Speed-focused: Prioritizes execution speed (80%) over quality (20%)
- Balanced: Balances quality (40%), speed (30%), and resource usage (30%)
- Multi-benchmark: Uses weighted combination of SimpleQA (60%) and BrowseComp (40%)

The results indicate which strategy is better suited for each optimization goal when using a statistically
significant sample size of 500 examples per experiment.
"""

    # Write summary to file
    with open(Path(base_dir) / "summary_report.md", "w", encoding="utf-8") as f:
        f.write(summary_text)


def run_strategy_simulation(num_examples=10):
    """
    Run a smaller simulation of the strategy benchmark with fewer examples
    for testing purposes or quick comparisons.

    This fallback simulation mode doesn't require actual database or LLM access,
    making it useful for testing the script structure.
    """
    timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")
    sim_output_dir = f"strategy_sim_results_{timestamp}"
    os.makedirs(sim_output_dir, exist_ok=True)

    # Define test query
    query = "What are the latest developments in fusion energy research?"

    # Define parameter space limited to strategies
    strategy_param_space = {
        "search_strategy": {
            "type": "categorical",
            "choices": ["iterdrag", "source_based"],
        },
        "iterations": {
            "type": "int",
            "low": 1,
            "high": 2,
            "step": 1,
        },
    }

    try:
        # Try to use real optimizer if available
        logger.info("Attempting to use real optimizer...")

        # Check if we can access necessary components
        from local_deep_research.config.llm_config import get_llm
        from local_deep_research.config.search_config import get_search

        # Try to initialize LLM and search engine to check configuration
        llm = get_llm()
        search = get_search()

        logger.info(
            f"Connected to LLM ({llm.__class__.__name__}) and search ({search.__class__.__name__})"
        )

        # Create optimizer for simulation
        sim_optimizer = OptunaOptimizer(
            base_query=query,
            output_dir=sim_output_dir,
            n_trials=5,  # Just a few trials for simulation
            timeout=600,  # 10 minutes timeout
            n_jobs=1,
            study_name="strategy_simulation",
            optimization_metrics=["quality", "speed"],
            metric_weights={"quality": 0.5, "speed": 0.5},
            num_examples=num_examples,  # Use fewer examples for simulation
        )

        # Run simulation
        best_params, best_score = sim_optimizer.optimize(strategy_param_space)

    except Exception as e:
        logger.warning(f"Could not initialize real optimizer: {e!s}")
        logger.warning(
            "Falling back to pure simulation mode (no real benchmarks)"
        )

        # Simulate optimization if real system is unavailable
        logger.info(
            "Running purely simulated optimization (no real benchmarks)"
        )
        best_params, best_score = simulate_optimization(
            strategy_param_space,
            n_trials=5,
            metric_weights={"quality": 0.5, "speed": 0.5},
        )

    # Log and save results
    logger.info(f"Simulation complete: {best_params}")
    logger.info(f"Best simulation score: {best_score}")

    sim_result = {
        "best_params": best_params,
        "best_score": best_score,
    }

    with open(
        Path(sim_output_dir) / "simulation_results.json", "w", encoding="utf-8"
    ) as f:
        json.dump(sim_result, f, indent=2)

    return sim_result


def simulate_optimization(
    param_space: Dict[str, Any],
    n_trials: int = 5,
    metric_weights: Dict[str, float] = None,
) -> Tuple[Dict[str, Any], float]:
    """
    Simulate an optimization process without actually running benchmarks.
    This is just for demonstration/testing purposes when the real system is unavailable.

    Args:
        param_space: Dictionary defining parameter search spaces
        n_trials: Number of simulated trials
        metric_weights: Weights for quality vs speed metrics

    Returns:
        Tuple of (best_parameters, best_score)
    """
    if metric_weights is None:
        metric_weights = {"quality": 0.5, "speed": 0.5}

    logger.info(f"Starting simulated optimization with {n_trials} trials")
    logger.info(f"Parameter space: {param_space}")
    logger.info(f"Metric weights: {metric_weights}")

    # Generate random trials
    best_score = 0.0
    best_params = {}

    for i in range(n_trials):
        # Generate random parameters
        params = {}
        for param_name, param_config in param_space.items():
            if param_config.get("type") == "int":
                params[param_name] = random.randint(
                    param_config.get("low", 1), param_config.get("high", 5)
                )
            elif param_config.get("type") == "categorical":
                params[param_name] = random.choice(
                    param_config.get("choices", ["standard"])
                )

        logger.info(f"Trial {i + 1}: Testing parameters: {params}")

        # Simulate execution delay
        time.sleep(0.5)

        # Simulate metrics for different strategies
        quality_score = 0.0
        speed_score = 0.0

        # Generate strategy-specific simulated scores
        if params.get("search_strategy") == "iterdrag":
            # IterDRAG typically has higher quality but lower speed
            quality_score = random.uniform(0.7, 0.95)
            speed_score = random.uniform(0.4, 0.7)
        elif params.get("search_strategy") == "source_based":
            # Source-based typically has medium quality but higher speed
            quality_score = random.uniform(0.6, 0.85)
            speed_score = random.uniform(0.6, 0.9)
        else:
            # Other strategies
            quality_score = random.uniform(0.5, 0.9)
            speed_score = random.uniform(0.5, 0.9)

        # More iterations generally means higher quality but lower speed
        iterations = params.get("iterations", 1)
        quality_score += (
            iterations * 0.05
        )  # More iterations slightly improves quality
        speed_score -= (
            iterations * 0.15
        )  # More iterations significantly reduces speed

        # Normalize scores to 0-1 range
        quality_score = max(0.0, min(1.0, quality_score))
        speed_score = max(0.0, min(1.0, speed_score))

        # Calculate weighted score based on metric weights
        combined_score = quality_score * metric_weights.get(
            "quality", 0.5
        ) + speed_score * metric_weights.get("speed", 0.5)

        logger.info(
            f"Trial {i + 1}: Quality: {quality_score:.2f}, Speed: {speed_score:.2f}, Score: {combined_score:.2f}"
        )

        # Update best parameters if this trial is better
        if combined_score > best_score:
            best_score = combined_score
            best_params = params.copy()
            logger.info(
                f"New best parameters found: {best_params} with score: {best_score:.2f}"
            )

    return best_params, best_score


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Run strategy benchmarks")
    parser.add_argument(
        "--simulate",
        action="store_true",
        help="Run a quick simulation instead of full benchmark",
    )
    parser.add_argument(
        "--examples",
        type=int,
        default=NUM_EXAMPLES,
        help=f"Number of examples to use (default: {NUM_EXAMPLES})",
    )

    args = parser.parse_args()

    if args.simulate:
        logger.info(f"Running simulation with {args.examples} examples")
        run_strategy_simulation(args.examples)
    else:
        logger.info(f"Running full benchmark with {args.examples} examples")
        NUM_EXAMPLES = args.examples  # Override global constant

        # Just run the benchmark function directly
        run_strategy_comparison()