"""
Optuna-based parameter optimizer for Local Deep Research.

This module provides the core optimization functionality using Optuna
to find optimal parameters for the research system, balancing quality
and performance metrics.
"""

import os
from pathlib import Path
import time
from datetime import datetime, UTC
from functools import partial
from typing import Any, Callable, Dict, List, Optional, Tuple

import joblib
import numpy as np
import optuna
from optuna.visualization import (
    plot_contour,
    plot_optimization_history,
    plot_param_importances,
    plot_slice,
)

from local_deep_research.benchmarks.efficiency.speed_profiler import (
    SpeedProfiler,
)
from local_deep_research.security import sanitize_data
from loguru import logger

from local_deep_research.benchmarks.evaluators import (
    CompositeBenchmarkEvaluator,
)

# Import benchmark evaluator components

# Try to import visualization libraries, but don't fail if not available
try:
    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D

    # We'll use matplotlib for plotting visualization results

    PLOTTING_AVAILABLE = True
except ImportError:
    PLOTTING_AVAILABLE = False
    logger.warning("Matplotlib not available, visualization will be limited")


class OptunaOptimizer:
    """
    Optimize parameters for Local Deep Research using Optuna.

    This class provides functionality to:
    1. Define search spaces for parameter optimization
    2. Evaluate parameter combinations using objective functions
    3. Find optimal parameters via Optuna
    4. Visualize and analyze optimization results
    """

    def __init__(
        self,
        base_query: str,
        output_dir: str = "optimization_results",
        model_name: Optional[str] = None,
        provider: Optional[str] = None,
        search_tool: Optional[str] = None,
        temperature: float = 0.7,
        n_trials: int = 30,
        timeout: Optional[int] = None,
        n_jobs: int = 1,
        study_name: Optional[str] = None,
        optimization_metrics: Optional[List[str]] = None,
        metric_weights: Optional[Dict[str, float]] = None,
        progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
        benchmark_weights: Optional[Dict[str, float]] = None,
    ):
        """
        Initialize the optimizer.

        Args:
            base_query: The research query to use for all experiments
            output_dir: Directory to save optimization results
            model_name: Name of the LLM model to use
            provider: LLM provider
            search_tool: Search engine to use
            temperature: LLM temperature
            n_trials: Number of parameter combinations to try
            timeout: Maximum seconds to run optimization (None for no limit)
            n_jobs: Number of parallel jobs for optimization
            study_name: Name of the Optuna study
            optimization_metrics: List of metrics to optimize (default: ["quality", "speed"])
            metric_weights: Dictionary of weights for each metric (e.g., {"quality": 0.6, "speed": 0.4})
            progress_callback: Optional callback for progress updates
            benchmark_weights: Dictionary mapping benchmark types to weights
                (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
                If None, only SimpleQA is used with weight 1.0
        """
        self.base_query = base_query
        self.output_dir = output_dir
        self.model_name = model_name
        self.provider = provider
        self.search_tool = search_tool
        self.temperature = temperature
        self.n_trials = n_trials
        self.timeout = timeout
        self.n_jobs = n_jobs
        self.optimization_metrics = optimization_metrics or ["quality", "speed"]
        self.metric_weights = metric_weights or {"quality": 0.6, "speed": 0.4}
        self.progress_callback = progress_callback

        # Initialize benchmark evaluator with weights
        self.benchmark_weights = benchmark_weights or {"simpleqa": 1.0}
        self.benchmark_evaluator = CompositeBenchmarkEvaluator(
            self.benchmark_weights
        )

        # Normalize weights to sum to 1.0
        total_weight = sum(self.metric_weights.values())
        if total_weight > 0:
            self.metric_weights = {
                k: v / total_weight for k, v in self.metric_weights.items()
            }

        # Generate a unique study name if not provided
        timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")
        self.study_name = study_name or f"ldr_opt_{timestamp}"

        # Create output directory
        os.makedirs(output_dir, exist_ok=True)

        # Store the trial history for analysis
        self.trials_history = []

        # Storage for the best parameters and study
        self.best_params = None
        self.study = None

    def optimize(
        self, param_space: Optional[Dict[str, Any]] = None
    ) -> Tuple[Dict[str, Any], float]:
        """
        Run the optimization process using Optuna.

        Args:
            param_space: Dictionary defining parameter search spaces
                         (if None, use default spaces)

        Returns:
            Tuple containing (best_parameters, best_score)
        """
        param_space = param_space or self._get_default_param_space()

        # Create a study object
        storage_name = f"sqlite:///{self.output_dir}/{self.study_name}.db"
        self.study = optuna.create_study(
            study_name=self.study_name,
            storage=storage_name,
            load_if_exists=True,
            direction="maximize",
            sampler=optuna.samplers.TPESampler(seed=42),
        )

        # Create partial function with param_space
        objective = partial(self._objective, param_space=param_space)

        # Log optimization start
        logger.info(
            f"Starting optimization with {self.n_trials} trials, {self.n_jobs} parallel jobs"
        )
        logger.info(f"Parameter space: {param_space}")
        logger.info(f"Metric weights: {self.metric_weights}")
        logger.info(f"Benchmark weights: {self.benchmark_weights}")

        # Initialize progress tracking
        if self.progress_callback:
            self.progress_callback(
                0,
                self.n_trials,
                {
                    "status": "starting",
                    "stage": "initialization",
                    "trials_completed": 0,
                    "total_trials": self.n_trials,
                },
            )

        try:
            # Run optimization
            self.study.optimize(
                objective,
                n_trials=self.n_trials,
                timeout=self.timeout,
                n_jobs=self.n_jobs,
                callbacks=[self._optimization_callback],
                show_progress_bar=True,
            )

            # Store best parameters
            self.best_params = self.study.best_params

            # Save the results
            self._save_results()

            # Create visualizations
            self._create_visualizations()

            logger.info(
                f"Optimization complete. Best parameters: {self.best_params}"
            )
            logger.info(f"Best value: {self.study.best_value}")

            # Report completion
            if self.progress_callback:
                self.progress_callback(
                    self.n_trials,
                    self.n_trials,
                    {
                        "status": "completed",
                        "stage": "finished",
                        "trials_completed": len(self.study.trials),
                        "total_trials": self.n_trials,
                        "best_params": self.best_params,
                        "best_value": self.study.best_value,
                    },
                )

            return self.best_params, self.study.best_value

        except KeyboardInterrupt:
            logger.info("Optimization interrupted by user")
            # Still save what we have
            self._save_results()
            self._create_visualizations()

            # Report interruption
            if self.progress_callback:
                self.progress_callback(
                    len(self.study.trials),
                    self.n_trials,
                    {
                        "status": "interrupted",
                        "stage": "interrupted",
                        "trials_completed": len(self.study.trials),
                        "total_trials": self.n_trials,
                        "best_params": self.study.best_params,
                        "best_value": self.study.best_value,
                    },
                )

            return self.study.best_params, self.study.best_value

    def _get_default_param_space(self) -> Dict[str, Any]:
        """
        Get default parameter search space.

        Returns:
            Dictionary defining the default parameter search spaces
        """
        return {
            "iterations": {
                "type": "int",
                "low": 1,
                "high": 5,
                "step": 1,
            },
            "questions_per_iteration": {
                "type": "int",
                "low": 1,
                "high": 5,
                "step": 1,
            },
            "search_strategy": {
                "type": "categorical",
                "choices": [
                    "iterdrag",
                    "standard",
                    "rapid",
                    "parallel",
                    "source_based",
                ],
            },
            "max_results": {
                "type": "int",
                "low": 10,
                "high": 100,
                "step": 10,
            },
        }

    def _objective(
        self, trial: optuna.Trial, param_space: Dict[str, Any]
    ) -> float:
        """
        Objective function for Optuna optimization.

        Args:
            trial: Optuna trial object
            param_space: Dictionary defining parameter search spaces

        Returns:
            Score to maximize
        """
        # Generate parameters for this trial
        params = {}
        for param_name, param_config in param_space.items():
            param_type = param_config["type"]

            if param_type == "int":
                params[param_name] = trial.suggest_int(
                    param_name,
                    param_config["low"],
                    param_config["high"],
                    step=param_config.get("step", 1),
                )
            elif param_type == "float":
                params[param_name] = trial.suggest_float(
                    param_name,
                    param_config["low"],
                    param_config["high"],
                    step=param_config.get("step"),
                    log=param_config.get("log", False),
                )
            elif param_type == "categorical":
                params[param_name] = trial.suggest_categorical(
                    param_name, param_config["choices"]
                )

        # Log the trial parameters
        logger.info(f"Trial {trial.number}: {params}")

        # Update progress callback if available
        if self.progress_callback:
            self.progress_callback(
                trial.number,
                self.n_trials,
                {
                    "status": "running",
                    "stage": "trial_started",
                    "trial_number": trial.number,
                    "params": params,
                    "trials_completed": trial.number,
                    "total_trials": self.n_trials,
                },
            )

        # Run an experiment with these parameters
        try:
            start_time = time.time()
            result = self._run_experiment(params)
            duration = time.time() - start_time

            # Store details about the trial
            trial_info = {
                "trial_number": trial.number,
                "params": params,
                "result": result,
                "score": result.get("score", 0),
                "duration": duration,
                "timestamp": datetime.now(UTC).isoformat(),
            }
            self.trials_history.append(trial_info)

            # Update callback with results
            if self.progress_callback:
                self.progress_callback(
                    trial.number,
                    self.n_trials,
                    {
                        "status": "completed",
                        "stage": "trial_completed",
                        "trial_number": trial.number,
                        "params": params,
                        "score": result.get("score", 0),
                        "trials_completed": trial.number + 1,
                        "total_trials": self.n_trials,
                    },
                )

            logger.info(
                f"Trial {trial.number} completed: {params}, score: {result['score']:.4f}"
            )

            return result["score"]
        except Exception as e:
            logger.exception(f"Error in trial {trial.number}: {e!s}")

            # Update callback with error
            if self.progress_callback:
                self.progress_callback(
                    trial.number,
                    self.n_trials,
                    {
                        "status": "error",
                        "stage": "trial_error",
                        "trial_number": trial.number,
                        "params": params,
                        "error": str(e),
                        "trials_completed": trial.number,
                        "total_trials": self.n_trials,
                    },
                )

            return float("-inf")  # Return a very low score for failed trials

    def _run_experiment(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        Run a single experiment with the given parameters.

        Args:
            params: Dictionary of parameters to test

        Returns:
            Results dictionary with metrics and score
        """
        # Extract parameters
        iterations = params.get("iterations", 2)
        questions_per_iteration = params.get("questions_per_iteration", 2)
        search_strategy = params.get("search_strategy", "iterdrag")
        max_results = params.get("max_results", 50)

        # Initialize profiling tools
        speed_profiler = SpeedProfiler()

        # Start profiling
        speed_profiler.start()

        try:
            # Create system configuration
            system_config = {
                "iterations": iterations,
                "questions_per_iteration": questions_per_iteration,
                "search_strategy": search_strategy,
                "search_tool": self.search_tool,
                "max_results": max_results,
                "model_name": self.model_name,
                "provider": self.provider,
            }

            # Evaluate quality using composite benchmark evaluator
            # Use a small number of examples for efficiency
            benchmark_dir = str(Path(self.output_dir) / "benchmark_temp")
            quality_results = self.benchmark_evaluator.evaluate(
                system_config=system_config,
                num_examples=5,  # Small number for optimization efficiency
                output_dir=benchmark_dir,
            )

            # Stop timing
            speed_profiler.stop()
            timing_results = speed_profiler.get_summary()

            # Extract key metrics
            quality_score = quality_results.get("quality_score", 0.0)
            benchmark_results = quality_results.get("benchmark_results", {})

            # Speed score: convert duration to a 0-1 score where faster is better
            # Using a reasonable threshold (e.g., 180 seconds for 5 examples)
            # Below this threshold: high score, above it: declining score
            total_duration = timing_results.get("total_duration", 180)
            speed_score = max(0.0, min(1.0, 1.0 - (total_duration - 60) / 180))

            # Calculate combined score based on weights
            combined_score = (
                self.metric_weights.get("quality", 0.6) * quality_score
                + self.metric_weights.get("speed", 0.4) * speed_score
            )

            # Return streamlined results
            return {
                "quality_score": quality_score,
                "benchmark_results": benchmark_results,
                "speed_score": speed_score,
                "total_duration": total_duration,
                "score": combined_score,
                "success": True,
            }

        except Exception as e:
            # Stop profiling on error
            speed_profiler.stop()

            # Log error
            logger.exception(f"Error in experiment: {e!s}")

            # Return error information
            return {"error": str(e), "score": 0.0, "success": False}

    def _optimization_callback(self, study: optuna.Study, trial: optuna.Trial):
        """
        Callback for the Optuna optimization process.

        Args:
            study: Optuna study object
            trial: Current trial
        """
        # Save intermediate results periodically
        if trial.number % 10 == 0 and trial.number > 0:
            self._save_results()
            self._create_quick_visualizations()

    def _save_results(self):
        """Save the optimization results to disk."""
        # Create a timestamp for filenames
        timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")

        # Save trial history
        from ...security.file_write_verifier import write_json_verified

        history_file = str(
            Path(self.output_dir) / f"{self.study_name}_history.json"
        )

        # Convert numpy values to native Python types for JSON serialization
        clean_history = []
        for trial in self.trials_history:
            clean_trial = {}
            for k, v in trial.items():
                if isinstance(v, dict):
                    clean_trial[k] = {
                        dk: (float(dv) if isinstance(dv, np.number) else dv)
                        for dk, dv in v.items()
                    }
                elif isinstance(v, np.number):
                    clean_trial[k] = float(v)
                else:
                    clean_trial[k] = v
            clean_history.append(clean_trial)

        # Sanitize sensitive data before writing to disk
        sanitized_history = sanitize_data(clean_history)

        write_json_verified(
            history_file,
            sanitized_history,
            "benchmark.allow_file_output",
            context="optimization history",
        )

        # Save current best parameters
        if (
            self.study
            and hasattr(self.study, "best_params")
            and self.study.best_params
        ):
            best_params_file = str(
                Path(self.output_dir) / f"{self.study_name}_best_params.json"
            )

            best_params_data = {
                "best_params": self.study.best_params,
                "best_value": float(self.study.best_value),
                "n_trials": len(self.study.trials),
                "timestamp": timestamp,
                "base_query": self.base_query,
                "model_name": self.model_name,
                "provider": self.provider,
                "search_tool": self.search_tool,
                "metric_weights": self.metric_weights,
                "benchmark_weights": self.benchmark_weights,
            }

            # Sanitize sensitive data before writing to disk
            sanitized_best_params = sanitize_data(best_params_data)

            write_json_verified(
                best_params_file,
                sanitized_best_params,
                "benchmark.allow_file_output",
                context="optimization best params",
            )

        # Save the Optuna study
        if self.study:
            study_file = str(
                Path(self.output_dir) / f"{self.study_name}_study.pkl"
            )
            joblib.dump(self.study, study_file)

        logger.info(f"Results saved to {self.output_dir}")

    def _create_visualizations(self):
        """Create and save comprehensive visualizations of the optimization results."""
        if not PLOTTING_AVAILABLE:
            logger.warning(
                "Matplotlib not available, skipping visualization creation"
            )
            return

        if not self.study or len(self.study.trials) < 2:
            logger.warning("Not enough trials to create visualizations")
            return

        # Create directory for visualizations
        viz_dir = Path(self.output_dir) / "visualizations"
        viz_dir.mkdir(parents=True, exist_ok=True)
        viz_dir = str(viz_dir)

        # Create Optuna visualizations
        self._create_optuna_visualizations(viz_dir)

        # Create custom visualizations
        self._create_custom_visualizations(viz_dir)

        logger.info(f"Visualizations saved to {viz_dir}")

    def _create_quick_visualizations(self):
        """Create a smaller set of visualizations for intermediate progress."""
        if (
            not PLOTTING_AVAILABLE
            or not self.study
            or len(self.study.trials) < 2
        ):
            return

        # Create directory for visualizations
        viz_dir = Path(self.output_dir) / "visualizations"
        viz_dir.mkdir(parents=True, exist_ok=True)
        viz_dir = str(viz_dir)

        # Create optimization history only (faster than full visualization)
        try:
            fig = plot_optimization_history(self.study)
            fig.write_image(
                str(
                    Path(viz_dir)
                    / f"{self.study_name}_optimization_history_current.png"
                )
            )
        except Exception as e:
            logger.exception(f"Error creating optimization history plot: {e!s}")

    def _create_optuna_visualizations(self, viz_dir: str):
        """
        Create and save Optuna's built-in visualizations.

        Args:
            viz_dir: Directory to save visualizations
        """
        timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")

        # 1. Optimization history
        try:
            fig = plot_optimization_history(self.study)
            fig.write_image(
                str(
                    Path(viz_dir)
                    / f"{self.study_name}_optimization_history_{timestamp}.png"
                )
            )
        except Exception as e:
            logger.exception(f"Error creating optimization history plot: {e!s}")

        # 2. Parameter importances
        try:
            fig = plot_param_importances(self.study)
            fig.write_image(
                str(
                    Path(viz_dir)
                    / f"{self.study_name}_param_importances_{timestamp}.png"
                )
            )
        except Exception as e:
            logger.exception(
                f"Error creating parameter importances plot: {e!s}"
            )

        # 3. Slice plot for each parameter
        try:
            for param_name in self.study.best_params.keys():
                fig = plot_slice(self.study, [param_name])
                fig.write_image(
                    str(
                        Path(viz_dir)
                        / f"{self.study_name}_slice_{param_name}_{timestamp}.png"
                    )
                )
        except Exception as e:
            logger.exception(f"Error creating slice plots: {e!s}")

        # 4. Contour plots for important parameter pairs
        try:
            # Get all parameter names
            param_names = list(self.study.best_params.keys())

            # Create contour plots for each pair
            for i in range(len(param_names)):
                for j in range(i + 1, len(param_names)):
                    try:
                        fig = plot_contour(
                            self.study, params=[param_names[i], param_names[j]]
                        )
                        fig.write_image(
                            str(
                                Path(viz_dir)
                                / f"{self.study_name}_contour_{param_names[i]}_{param_names[j]}_{timestamp}.png"
                            )
                        )
                    except Exception as e:
                        logger.warning(
                            f"Error creating contour plot for {param_names[i]} vs {param_names[j]}: {e!s}"
                        )
        except Exception as e:
            logger.exception(f"Error creating contour plots: {e!s}")

    def _create_custom_visualizations(self, viz_dir: str):
        """
        Create custom visualizations based on trial history.

        Args:
            viz_dir: Directory to save visualizations
        """
        if not self.trials_history:
            return

        timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")

        # Create quality vs speed plot
        self._create_quality_vs_speed_plot(viz_dir, timestamp)

        # Create parameter evolution plots
        self._create_parameter_evolution_plots(viz_dir, timestamp)

        # Create trial duration vs score plot
        self._create_duration_vs_score_plot(viz_dir, timestamp)

    def _create_quality_vs_speed_plot(self, viz_dir: str, timestamp: str):
        """Create a plot showing quality vs. speed trade-off."""
        if not self.trials_history:
            return

        # Extract data from successful trials
        successful_trials = [
            t
            for t in self.trials_history
            if t.get("result", {}).get("success", False)
        ]

        if not successful_trials:
            logger.warning("No successful trials for visualization")
            return

        try:
            plt.figure(figsize=(10, 8))

            # Extract metrics
            quality_scores = []
            speed_scores = []
            labels = []
            iterations_values = []
            questions_values = []

            for trial in successful_trials:
                result = trial["result"]
                quality = result.get("quality_score", 0)
                speed = result.get("speed_score", 0)
                iterations = trial["params"].get("iterations", 0)
                questions = trial["params"].get("questions_per_iteration", 0)

                quality_scores.append(quality)
                speed_scores.append(speed)
                labels.append(f"Trial {trial['trial_number']}")
                iterations_values.append(iterations)
                questions_values.append(questions)

            # Create scatter plot with size based on iterations*questions
            sizes = [
                i * q * 5
                for i, q in zip(
                    iterations_values, questions_values, strict=False
                )
            ]
            scatter = plt.scatter(
                quality_scores,
                speed_scores,
                s=sizes,
                alpha=0.7,
                c=range(len(quality_scores)),
                cmap="viridis",
            )

            # Highlight best trial
            best_trial = max(
                successful_trials,
                key=lambda x: x.get("result", {}).get("score", 0),
            )
            best_quality = best_trial["result"].get("quality_score", 0)
            best_speed = best_trial["result"].get("speed_score", 0)
            best_iter = best_trial["params"].get("iterations", 0)
            best_questions = best_trial["params"].get(
                "questions_per_iteration", 0
            )

            plt.scatter(
                [best_quality],
                [best_speed],
                s=200,
                facecolors="none",
                edgecolors="red",
                linewidth=2,
                label=f"Best: {best_iter}×{best_questions}",
            )

            # Add annotations for key points
            for i, (q, s, label) in enumerate(
                zip(quality_scores, speed_scores, labels, strict=False)
            ):
                if i % max(1, len(quality_scores) // 5) == 0:  # Label ~5 points
                    plt.annotate(
                        f"{iterations_values[i]}×{questions_values[i]}",
                        (q, s),
                        xytext=(5, 5),
                        textcoords="offset points",
                    )

            # Add colorbar and labels
            cbar = plt.colorbar(scatter)
            cbar.set_label("Trial Progression")

            # Add benchmark weight information
            weights_str = ", ".join(
                [f"{k}:{v:.1f}" for k, v in self.benchmark_weights.items()]
            )
            plt.title(
                f"Quality vs. Speed Trade-off\nBenchmark Weights: {weights_str}"
            )
            plt.xlabel("Quality Score (Benchmark Accuracy)")
            plt.ylabel("Speed Score")
            plt.grid(True, linestyle="--", alpha=0.7)

            # Add legend explaining size
            legend_elements = [
                Line2D(
                    [0],
                    [0],
                    marker="o",
                    color="w",
                    markerfacecolor="gray",
                    markersize=np.sqrt(n * 5 / np.pi),
                    label=f"{n} Total Questions",
                )
                for n in [5, 10, 15, 20, 25]
            ]
            plt.legend(handles=legend_elements, title="Workload")

            # Save the figure
            plt.tight_layout()
            plt.savefig(
                str(
                    Path(viz_dir)
                    / f"{self.study_name}_quality_vs_speed_{timestamp}.png"
                )
            )
            plt.close()
        except Exception as e:
            logger.exception(f"Error creating quality vs speed plot: {e!s}")

    def _create_parameter_evolution_plots(self, viz_dir: str, timestamp: str):
        """Create plots showing how parameter values evolve over trials."""
        try:
            successful_trials = [
                t
                for t in self.trials_history
                if t.get("result", {}).get("success", False)
            ]

            if not successful_trials or len(successful_trials) < 5:
                return

            # Get key parameters
            main_params = list(successful_trials[0]["params"].keys())

            # For each parameter, plot its values over trials
            for param_name in main_params:
                plt.figure(figsize=(12, 6))

                trial_numbers = []
                param_values = []
                scores = []

                for trial in self.trials_history:
                    if "params" in trial and param_name in trial["params"]:
                        trial_numbers.append(trial["trial_number"])
                        param_values.append(trial["params"][param_name])
                        scores.append(trial.get("score", 0))

                # Create evolution plot
                scatter = plt.scatter(
                    trial_numbers,
                    param_values,
                    c=scores,
                    cmap="plasma",
                    alpha=0.8,
                    s=80,
                )

                # Add best trial marker
                best_trial_idx = scores.index(max(scores))
                plt.scatter(
                    [trial_numbers[best_trial_idx]],
                    [param_values[best_trial_idx]],
                    s=150,
                    facecolors="none",
                    edgecolors="red",
                    linewidth=2,
                    label=f"Best Value: {param_values[best_trial_idx]}",
                )

                # Add colorbar
                cbar = plt.colorbar(scatter)
                cbar.set_label("Score")

                # Set chart properties
                plt.title(f"Evolution of {param_name} Values")
                plt.xlabel("Trial Number")
                plt.ylabel(param_name)
                plt.grid(True, linestyle="--", alpha=0.7)
                plt.legend()

                # For categorical parameters, adjust y-axis
                if isinstance(param_values[0], str):
                    unique_values = sorted(set(param_values))
                    plt.yticks(range(len(unique_values)), unique_values)

                # Save the figure
                plt.tight_layout()
                plt.savefig(
                    str(
                        Path(viz_dir)
                        / f"{self.study_name}_param_evolution_{param_name}_{timestamp}.png"
                    )
                )
                plt.close()
        except Exception as e:
            logger.exception(f"Error creating parameter evolution plots: {e!s}")

    def _create_duration_vs_score_plot(self, viz_dir: str, timestamp: str):
        """Create a plot showing trial duration vs score."""
        try:
            plt.figure(figsize=(10, 6))

            successful_trials = [
                t
                for t in self.trials_history
                if t.get("result", {}).get("success", False)
            ]

            if not successful_trials:
                return

            trial_durations = []
            trial_scores = []
            trial_iterations = []
            trial_questions = []

            for trial in successful_trials:
                duration = trial.get("duration", 0)
                score = trial.get("score", 0)
                iterations = trial.get("params", {}).get("iterations", 1)
                questions = trial.get("params", {}).get(
                    "questions_per_iteration", 1
                )

                trial_durations.append(duration)
                trial_scores.append(score)
                trial_iterations.append(iterations)
                trial_questions.append(questions)

            # Total questions per trial
            total_questions = [
                i * q
                for i, q in zip(trial_iterations, trial_questions, strict=False)
            ]

            # Create scatter plot with size based on total questions
            plt.scatter(
                trial_durations,
                trial_scores,
                s=[
                    q * 5 for q in total_questions
                ],  # Size based on total questions
                alpha=0.7,
                c=range(len(trial_durations)),
                cmap="viridis",
            )

            # Add labels
            plt.xlabel("Trial Duration (seconds)")
            plt.ylabel("Score")
            plt.title("Trial Duration vs. Score")
            plt.grid(True, linestyle="--", alpha=0.7)

            # Add trial number annotations for selected points
            for i, (d, s) in enumerate(
                zip(trial_durations, trial_scores, strict=False)
            ):
                if (
                    i % max(1, len(trial_durations) // 5) == 0
                ):  # Annotate ~5 points
                    plt.annotate(
                        f"{trial_iterations[i]}×{trial_questions[i]}",
                        (d, s),
                        xytext=(5, 5),
                        textcoords="offset points",
                    )

            # Save the figure
            plt.tight_layout()
            plt.savefig(
                str(
                    Path(viz_dir)
                    / f"{self.study_name}_duration_vs_score_{timestamp}.png"
                )
            )
            plt.close()
        except Exception as e:
            logger.exception(f"Error creating duration vs score plot: {e!s}")


def optimize_parameters(
    query: str,
    param_space: Optional[Dict[str, Any]] = None,
    output_dir: str = str(Path("data") / "optimization_results"),
    model_name: Optional[str] = None,
    provider: Optional[str] = None,
    search_tool: Optional[str] = None,
    temperature: float = 0.7,
    n_trials: int = 30,
    timeout: Optional[int] = None,
    n_jobs: int = 1,
    study_name: Optional[str] = None,
    optimization_metrics: Optional[List[str]] = None,
    metric_weights: Optional[Dict[str, float]] = None,
    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
    benchmark_weights: Optional[Dict[str, float]] = None,
) -> Tuple[Dict[str, Any], float]:
    """
    Optimize parameters for Local Deep Research.

    Args:
        query: The research query to use for all experiments
        param_space: Dictionary defining parameter search spaces (optional)
        output_dir: Directory to save optimization results
        model_name: Name of the LLM model to use
        provider: LLM provider
        search_tool: Search engine to use
        temperature: LLM temperature
        n_trials: Number of parameter combinations to try
        timeout: Maximum seconds to run optimization (None for no limit)
        n_jobs: Number of parallel jobs for optimization
        study_name: Name of the Optuna study
        optimization_metrics: List of metrics to optimize (default: ["quality", "speed"])
        metric_weights: Dictionary of weights for each metric (e.g., {"quality": 0.6, "speed": 0.4})
        progress_callback: Optional callback for progress updates
        benchmark_weights: Dictionary mapping benchmark types to weights
            (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
            If None, only SimpleQA is used with weight 1.0

    Returns:
        Tuple of (best_parameters, best_score)
    """
    # Create optimizer
    optimizer = OptunaOptimizer(
        base_query=query,
        output_dir=output_dir,
        model_name=model_name,
        provider=provider,
        search_tool=search_tool,
        temperature=temperature,
        n_trials=n_trials,
        timeout=timeout,
        n_jobs=n_jobs,
        study_name=study_name,
        optimization_metrics=optimization_metrics,
        metric_weights=metric_weights,
        progress_callback=progress_callback,
        benchmark_weights=benchmark_weights,
    )

    # Run optimization
    return optimizer.optimize(param_space)


def optimize_for_speed(
    query: str,
    n_trials: int = 20,
    output_dir: str = str(Path("data") / "optimization_results"),
    model_name: Optional[str] = None,
    provider: Optional[str] = None,
    search_tool: Optional[str] = None,
    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
    benchmark_weights: Optional[Dict[str, float]] = None,
) -> Tuple[Dict[str, Any], float]:
    """
    Optimize parameters with a focus on speed performance.

    Args:
        query: The research query to use for all experiments
        n_trials: Number of parameter combinations to try
        output_dir: Directory to save optimization results
        model_name: Name of the LLM model to use
        provider: LLM provider
        search_tool: Search engine to use
        progress_callback: Optional callback for progress updates
        benchmark_weights: Dictionary mapping benchmark types to weights
            (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
            If None, only SimpleQA is used with weight 1.0

    Returns:
        Tuple of (best_parameters, best_score)
    """
    # Focus on speed with reduced parameter space
    param_space = {
        "iterations": {
            "type": "int",
            "low": 1,
            "high": 3,
            "step": 1,
        },
        "questions_per_iteration": {
            "type": "int",
            "low": 1,
            "high": 3,
            "step": 1,
        },
        "search_strategy": {
            "type": "categorical",
            "choices": ["rapid", "parallel", "source_based"],
        },
    }

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

    return optimize_parameters(
        query=query,
        param_space=param_space,
        output_dir=output_dir,
        model_name=model_name,
        provider=provider,
        search_tool=search_tool,
        n_trials=n_trials,
        metric_weights=metric_weights,
        optimization_metrics=["speed", "quality"],
        progress_callback=progress_callback,
        benchmark_weights=benchmark_weights,
    )


def optimize_for_quality(
    query: str,
    n_trials: int = 30,
    output_dir: str = str(Path("data") / "optimization_results"),
    model_name: Optional[str] = None,
    provider: Optional[str] = None,
    search_tool: Optional[str] = None,
    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
    benchmark_weights: Optional[Dict[str, float]] = None,
) -> Tuple[Dict[str, Any], float]:
    """
    Optimize parameters with a focus on result quality.

    Args:
        query: The research query to use for all experiments
        n_trials: Number of parameter combinations to try
        output_dir: Directory to save optimization results
        model_name: Name of the LLM model to use
        provider: LLM provider
        search_tool: Search engine to use
        progress_callback: Optional callback for progress updates
        benchmark_weights: Dictionary mapping benchmark types to weights
            (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
            If None, only SimpleQA is used with weight 1.0

    Returns:
        Tuple of (best_parameters, best_score)
    """
    # Quality-focused weights
    metric_weights = {"quality": 0.9, "speed": 0.1}

    return optimize_parameters(
        query=query,
        output_dir=output_dir,
        model_name=model_name,
        provider=provider,
        search_tool=search_tool,
        n_trials=n_trials,
        metric_weights=metric_weights,
        optimization_metrics=["quality", "speed"],
        progress_callback=progress_callback,
        benchmark_weights=benchmark_weights,
    )


def optimize_for_efficiency(
    query: str,
    n_trials: int = 25,
    output_dir: str = str(Path("data") / "optimization_results"),
    model_name: Optional[str] = None,
    provider: Optional[str] = None,
    search_tool: Optional[str] = None,
    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
    benchmark_weights: Optional[Dict[str, float]] = None,
) -> Tuple[Dict[str, Any], float]:
    """
    Optimize parameters with a focus on resource efficiency.

    Args:
        query: The research query to use for all experiments
        n_trials: Number of parameter combinations to try
        output_dir: Directory to save optimization results
        model_name: Name of the LLM model to use
        provider: LLM provider
        search_tool: Search engine to use
        progress_callback: Optional callback for progress updates
        benchmark_weights: Dictionary mapping benchmark types to weights
            (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
            If None, only SimpleQA is used with weight 1.0

    Returns:
        Tuple of (best_parameters, best_score)
    """
    # Balance of quality, speed and resource usage
    metric_weights = {"quality": 0.4, "speed": 0.3, "resource": 0.3}

    return optimize_parameters(
        query=query,
        output_dir=output_dir,
        model_name=model_name,
        provider=provider,
        search_tool=search_tool,
        n_trials=n_trials,
        metric_weights=metric_weights,
        optimization_metrics=["quality", "speed", "resource"],
        progress_callback=progress_callback,
        benchmark_weights=benchmark_weights,
    )
