#!/usr/bin/env python
"""
Multi-benchmark optimization with speed metrics demonstration.

This script shows how the multi-benchmark API can be used with speed optimization
without actually running the benchmarks (simulation only).

Usage:
    # Run from project root with venv activated
    cd /path/to/local-deep-research
    source .venv/bin/activate
    cd src
    python ../examples/optimization/multi_benchmark_speed_demo.py
"""

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

# Add src directory to Python path
src_dir = str((Path(__file__).parent.parent / "src").resolve())
if src_dir not in sys.path:
    sys.path.insert(0, src_dir)


class SimulatedBenchmarkEvaluator:
    """Simulated benchmark evaluator that doesn't run actual benchmarks."""

    def __init__(self, name, quality_score=0.75, speed_score=0.65):
        self.name = name
        self.quality_score = quality_score
        self.speed_score = speed_score

    def evaluate(self, system_config, num_examples=1, output_dir=None):
        """Simulate benchmark evaluation with predefined scores."""
        print(f"[SIM] Running {self.name} benchmark simulation...")
        print(f"[SIM] System config: {system_config}")

        # Return simulated results
        return {
            "quality_score": self.quality_score,
            "speed_score": self.speed_score,
            "component_timing": {
                "search": 0.5,
                "processing": 0.3,
                "llm": 1.2,
                "total": 2.0,
            },
            "resource_usage": {"memory_mb": 500, "cpu_percent": 30},
        }


class SimulatedCompositeBenchmarkEvaluator:
    """Simulated composite benchmark evaluator that combines multiple benchmarks."""

    def __init__(self, benchmark_weights=None):
        self.benchmark_weights = benchmark_weights or {"simpleqa": 1.0}
        print(
            f"[SIM] Created composite evaluator with weights: {self.benchmark_weights}"
        )

        # Normalize weights
        total = sum(self.benchmark_weights.values())
        self.normalized_weights = {
            k: v / total for k, v in self.benchmark_weights.items()
        }
        print(f"[SIM] Normalized weights: {self.normalized_weights}")

        # Create evaluators with slightly different characteristics
        self.evaluators = {
            "simpleqa": SimulatedBenchmarkEvaluator(
                "SimpleQA", quality_score=0.80, speed_score=0.70
            ),
            "browsecomp": SimulatedBenchmarkEvaluator(
                "BrowseComp", quality_score=0.85, speed_score=0.60
            ),
        }

    def evaluate(self, system_config, num_examples=1, output_dir=None):
        """Run evaluation for all benchmarks with weights."""
        print(
            f"[SIM] Running composite evaluation with {num_examples} examples"
        )

        # Run each benchmark
        benchmark_results = {}
        for name, evaluator in self.evaluators.items():
            if name in self.benchmark_weights:
                benchmark_results[name] = evaluator.evaluate(
                    system_config, num_examples, output_dir
                )

        # Calculate combined quality score
        quality_score = sum(
            self.normalized_weights[name] * results["quality_score"]
            for name, results in benchmark_results.items()
        )

        # Calculate combined speed score
        speed_score = sum(
            self.normalized_weights[name] * results["speed_score"]
            for name, results in benchmark_results.items()
        )

        return {
            "quality_score": quality_score,
            "speed_score": speed_score,
            "benchmark_weights": self.benchmark_weights,
            "benchmark_results": benchmark_results,
        }


class SimulatedOptimizer:
    """Simulated optimizer that demonstrates the API structure without running actual optimization."""

    def __init__(
        self,
        base_query: str = "Example query",
        output_dir: str = "./results",
        metric_weights: Dict[str, float] = None,
        benchmark_weights: Dict[str, float] = None,
    ):
        self.base_query = base_query
        self.output_dir = output_dir
        self.metric_weights = metric_weights or {"quality": 0.6, "speed": 0.4}
        self.benchmark_weights = benchmark_weights or {"simpleqa": 1.0}

        # Create evaluator
        self.evaluator = SimulatedCompositeBenchmarkEvaluator(
            self.benchmark_weights
        )

        print("[SIM] Created optimizer with:")
        print(f"[SIM]   - Metric weights: {self.metric_weights}")
        print(f"[SIM]   - Benchmark weights: {self.benchmark_weights}")

    def optimize(self, param_space=None):
        """Simulate optimization process."""
        # Simulate a few trials
        print("[SIM] Running optimization with parameter space:", param_space)
        print("[SIM] Using metric weights:", self.metric_weights)

        # Simulate trials
        trials = [
            {"iterations": 1, "search_strategy": "rapid"},
            {"iterations": 2, "search_strategy": "standard"},
            {"iterations": 3, "search_strategy": "iterdrag"},
        ]

        # Simulate scores based on trials and weights
        trial_scores = []
        for trial in trials:
            # Get benchmark scores
            results = self.evaluator.evaluate(trial, num_examples=1)

            # Calculate combined score based on metric weights
            combined_score = (
                self.metric_weights.get("quality", 0) * results["quality_score"]
                + self.metric_weights.get("speed", 0) * results["speed_score"]
            )

            trial_scores.append((trial, combined_score))
            print(f"[SIM] Trial {trial}: Score {combined_score:.4f}")

        # Return best parameters and score
        best_trial, best_score = max(trial_scores, key=lambda x: x[1])
        print(f"[SIM] Best trial: {best_trial} with score {best_score:.4f}")

        return best_trial, best_score


def optimize_for_quality(
    query: str, benchmark_weights: Dict[str, float] = None
):
    """Simulate quality-focused optimization."""
    print("\n🔍 Simulating quality-focused optimization...")

    # Quality-focused weights: 90% quality, 10% speed
    metric_weights = {"quality": 0.9, "speed": 0.1}

    optimizer = SimulatedOptimizer(
        base_query=query,
        metric_weights=metric_weights,
        benchmark_weights=benchmark_weights,
    )

    return optimizer.optimize()


def optimize_for_speed(query: str, benchmark_weights: Dict[str, float] = None):
    """Simulate speed-focused optimization."""
    print("\n🔍 Simulating speed-focused optimization...")

    # Speed-focused weights: 20% quality, 80% speed
    metric_weights = {"quality": 0.2, "speed": 0.8}

    optimizer = SimulatedOptimizer(
        base_query=query,
        metric_weights=metric_weights,
        benchmark_weights=benchmark_weights,
    )

    return optimizer.optimize()


def optimize_for_efficiency(
    query: str, benchmark_weights: Dict[str, float] = None
):
    """Simulate efficiency-focused optimization."""
    print("\n🔍 Simulating efficiency-focused optimization...")

    # Balanced weights: 40% quality, 30% speed, 30% resource
    metric_weights = {"quality": 0.4, "speed": 0.3, "resource": 0.3}

    optimizer = SimulatedOptimizer(
        base_query=query,
        metric_weights=metric_weights,
        benchmark_weights=benchmark_weights,
    )

    return optimizer.optimize()


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 simulated multi-benchmark optimization examples."""
    query = "Fusion energy research developments"

    # Run 1: SimpleQA benchmark only with quality focus
    print("\n🔬 DEMO: SimpleQA-only optimization (quality focus)")
    params1, score1 = optimize_for_quality(
        query=query, benchmark_weights={"simpleqa": 1.0}
    )
    print_optimization_results(params1, score1)

    # Run 2: BrowseComp benchmark only with quality focus
    print("\n🔬 DEMO: BrowseComp-only optimization (quality focus)")
    params2, score2 = optimize_for_quality(
        query=query, benchmark_weights={"browsecomp": 1.0}
    )
    print_optimization_results(params2, score2)

    # Run 3: Combined benchmarks with quality focus
    print("\n🔬 DEMO: Combined benchmarks with weights (quality focus)")
    params3, score3 = optimize_for_quality(
        query=query, benchmark_weights={"simpleqa": 0.6, "browsecomp": 0.4}
    )
    print_optimization_results(params3, score3)

    # Run 4: Combined benchmarks with speed focus
    print("\n🔬 DEMO: Combined benchmarks with weights (speed focus)")
    params4, score4 = optimize_for_speed(
        query=query, benchmark_weights={"simpleqa": 0.6, "browsecomp": 0.4}
    )
    print_optimization_results(params4, score4)
    print("Speed metrics weighting: Quality (20%), Speed (80%)")

    # Run 5: Combined benchmarks with efficiency focus
    print("\n🔬 DEMO: Combined benchmarks with weights (efficiency focus)")
    params5, score5 = optimize_for_efficiency(
        query=query, benchmark_weights={"simpleqa": 0.6, "browsecomp": 0.4}
    )
    print_optimization_results(params5, score5)
    print(
        "Efficiency metrics weighting: Quality (40%), Speed (30%), Resource (30%)"
    )


if __name__ == "__main__":
    main()