"""
Visualization utilities for optimization results.

This module provides functions for generating visual representations
of benchmark and optimization results.
"""

from loguru import logger
from typing import Any, Dict, List, Optional

import numpy as np


# Check if matplotlib is available
try:
    import matplotlib.pyplot as plt
    from matplotlib.figure import Figure

    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False
    logger.warning(
        "Matplotlib not available. Visualization functions will be limited."
    )


def plot_optimization_history(
    trial_values: List[float],
    best_values: List[float],
    output_file: Optional[str] = None,
    title: str = "Optimization History",
) -> Optional[Figure]:
    """
    Plot the optimization history.

    Args:
        trial_values: List of objective values for each trial
        best_values: List of best values observed up to each trial
        output_file: Path to save the plot (if None, returns figure without saving)
        title: Plot title

    Returns:
        Matplotlib figure or None if matplotlib is not available
    """
    if not MATPLOTLIB_AVAILABLE:
        logger.warning("Matplotlib not available. Cannot create plot.")
        return None

    fig, ax = plt.subplots(figsize=(10, 6))
    trials = list(range(1, len(trial_values) + 1))

    # Plot trial values and best values
    ax.plot(trials, trial_values, "o-", alpha=0.5, label="Trial Value")
    ax.plot(trials, best_values, "r-", label="Best Value")

    # Add labels and title
    ax.set_xlabel("Trial Number")
    ax.set_ylabel("Objective Value")
    ax.set_title(title)
    ax.grid(True, linestyle="--", alpha=0.7)
    ax.legend()

    # Save or return
    if output_file:
        fig.tight_layout()
        fig.savefig(output_file, dpi=300, bbox_inches="tight")
        logger.info(f"Saved optimization history plot to {output_file}")

    return fig


def plot_parameter_importance(
    parameter_names: List[str],
    importance_values: List[float],
    output_file: Optional[str] = None,
    title: str = "Parameter Importance",
) -> Optional[Figure]:
    """
    Plot parameter importance.

    Args:
        parameter_names: List of parameter names
        importance_values: List of importance values
        output_file: Path to save the plot (if None, returns figure without saving)
        title: Plot title

    Returns:
        Matplotlib figure or None if matplotlib is not available
    """
    if not MATPLOTLIB_AVAILABLE:
        logger.warning("Matplotlib not available. Cannot create plot.")
        return None

    # Sort by importance
    sorted_indices = np.argsort(importance_values)
    sorted_names = [parameter_names[i] for i in sorted_indices]
    sorted_values = [importance_values[i] for i in sorted_indices]

    fig, ax = plt.subplots(figsize=(10, 6))
    y_pos = range(len(sorted_names))

    # Create horizontal bar chart
    ax.barh(y_pos, sorted_values, align="center")
    ax.set_yticks(y_pos)
    ax.set_yticklabels(sorted_names)
    ax.invert_yaxis()  # Labels read top-to-bottom

    # Add labels and title
    ax.set_xlabel("Importance")
    ax.set_title(title)
    ax.grid(True, linestyle="--", alpha=0.3, axis="x")

    # Save or return
    if output_file:
        fig.tight_layout()
        fig.savefig(output_file, dpi=300, bbox_inches="tight")
        logger.info(f"Saved parameter importance plot to {output_file}")

    return fig


def plot_quality_vs_speed(
    quality_scores: List[float],
    speed_scores: List[float],
    _parameter_values: Optional[List[Dict[str, Any]]] = None,
    output_file: Optional[str] = None,
    title: str = "Quality vs. Speed Trade-off",
) -> Optional[Figure]:
    """
    Plot quality vs. speed trade-off.

    Args:
        quality_scores: List of quality scores
        speed_scores: List of speed scores
        parameter_values: Optional list of parameter dictionaries for each point
        output_file: Path to save the plot (if None, returns figure without saving)
        title: Plot title

    Returns:
        Matplotlib figure or None if matplotlib is not available
    """
    if not MATPLOTLIB_AVAILABLE:
        logger.warning("Matplotlib not available. Cannot create plot.")
        return None

    fig, ax = plt.subplots(figsize=(10, 8))

    # Create scatter plot
    scatter = ax.scatter(
        speed_scores,
        quality_scores,
        c=np.arange(len(quality_scores)),
        cmap="viridis",
        alpha=0.7,
        s=100,
    )

    # Add colorbar to show trial number
    cbar = plt.colorbar(scatter)
    cbar.set_label("Trial Number")

    # Add labels and title
    ax.set_xlabel("Speed Score (higher = faster)")
    ax.set_ylabel("Quality Score (higher = better)")
    ax.set_title(title)
    ax.grid(True, linestyle="--", alpha=0.5)

    # Add reference lines
    ax.axhline(
        y=0.7,
        color="r",
        linestyle="--",
        alpha=0.3,
        label="Good Quality Threshold",
    )
    ax.axvline(
        x=0.7,
        color="g",
        linestyle="--",
        alpha=0.3,
        label="Good Speed Threshold",
    )

    # Mark Pareto frontier
    if len(quality_scores) > 2:
        try:
            # Identify Pareto frontier points
            pareto_points = []
            for i in range(len(quality_scores)):
                is_pareto = True
                for j in range(len(quality_scores)):
                    if i != j:
                        if (
                            quality_scores[j] >= quality_scores[i]
                            and speed_scores[j] >= speed_scores[i]
                        ):
                            if (
                                quality_scores[j] > quality_scores[i]
                                or speed_scores[j] > speed_scores[i]
                            ):
                                is_pareto = False
                                break
                if is_pareto:
                    pareto_points.append((speed_scores[i], quality_scores[i]))

            # Sort pareto points by speed score
            pareto_points.sort()
            if pareto_points:
                pareto_x, pareto_y = zip(*pareto_points, strict=False)
                ax.plot(pareto_x, pareto_y, "k--", label="Pareto Frontier")
                ax.scatter(pareto_x, pareto_y, c="red", s=50, alpha=0.8)
        except Exception:
            logger.warning("Error calculating Pareto frontier")

    ax.legend()

    # Save or return
    if output_file:
        fig.tight_layout()
        fig.savefig(output_file, dpi=300, bbox_inches="tight")
        logger.info(f"Saved quality vs. speed plot to {output_file}")

    return fig
