"""
Configuration comparison for Local Deep Research.

This module provides functions for comparing different parameter configurations
and evaluating their performance across various metrics.
"""

import os
from datetime import datetime, UTC
from pathlib import Path
from typing import Any, Dict, List, Optional

import matplotlib.pyplot as plt
import numpy as np
from loguru import logger
from matplotlib.patches import Circle, RegularPolygon

from local_deep_research.benchmarks.efficiency.resource_monitor import (
    ResourceMonitor,
)
from local_deep_research.benchmarks.efficiency.speed_profiler import (
    SpeedProfiler,
)
from local_deep_research.benchmarks.optimization.metrics import (
    calculate_combined_score,
    calculate_quality_metrics,
    calculate_resource_metrics,
    calculate_speed_metrics,
)
from local_deep_research.config.llm_config import get_llm
from local_deep_research.config.search_config import get_search
from local_deep_research.search_system import AdvancedSearchSystem


def compare_configurations(
    query: str,
    configurations: List[Dict[str, Any]],
    output_dir: str = "comparison_results",
    model_name: Optional[str] = None,
    provider: Optional[str] = None,
    search_tool: Optional[str] = None,
    repetitions: int = 1,
    metric_weights: Optional[Dict[str, float]] = None,
) -> Dict[str, Any]:
    """
    Compare multiple parameter configurations.

    Args:
        query: Research query to use for evaluation
        configurations: List of parameter configurations to compare
        output_dir: Directory to save comparison results
        model_name: Name of the LLM model to use
        provider: LLM provider
        search_tool: Search engine to use
        repetitions: Number of repetitions for each configuration
        metric_weights: Dictionary of weights for each metric type

    Returns:
        Dictionary with comparison results
    """
    os.makedirs(output_dir, exist_ok=True)

    # Default metric weights if not provided
    if metric_weights is None:
        metric_weights = {
            "quality": 0.6,
            "speed": 0.4,
            "resource": 0.0,  # Disabled by default
        }

    # Verify valid configurations
    if not configurations:
        logger.error("No configurations provided for comparison")
        return {"error": "No configurations provided"}

    # Results storage
    results = []

    # Process each configuration
    for i, config in enumerate(configurations):
        logger.info(
            f"Evaluating configuration {i + 1}/{len(configurations)}: {config}"
        )

        # Name for this configuration
        config_name = config.get("name", f"Configuration {i + 1}")

        # Results for all repetitions of this configuration
        config_results = []

        # Run multiple repetitions
        for rep in range(repetitions):
            logger.info(
                f"Starting repetition {rep + 1}/{repetitions} for {config_name}"
            )

            try:
                # Run the configuration
                result = _evaluate_single_configuration(
                    query=query,
                    config=config,
                    model_name=model_name,
                    provider=provider,
                    search_tool=search_tool,
                )

                config_results.append(result)
                logger.info(f"Completed repetition {rep + 1} for {config_name}")

            except Exception as e:
                logger.exception(
                    f"Error in {config_name}, repetition {rep + 1}: {e!s}"
                )
                # Add error info but continue with other configurations
                config_results.append({"error": str(e), "success": False})

        # Calculate aggregate metrics across repetitions
        if config_results:
            # Filter out failed runs
            successful_runs = [
                r for r in config_results if r.get("success", False)
            ]

            if successful_runs:
                # Calculate average metrics
                avg_metrics = _calculate_average_metrics(successful_runs)

                # Calculate overall score
                overall_score = calculate_combined_score(
                    quality_metrics=avg_metrics.get("quality_metrics", {}),
                    speed_metrics=avg_metrics.get("speed_metrics", {}),
                    resource_metrics=avg_metrics.get("resource_metrics", {}),
                    weights=metric_weights,
                )

                result_summary = {
                    "name": config_name,
                    "configuration": config,
                    "success": True,
                    "runs_completed": len(successful_runs),
                    "runs_failed": len(config_results) - len(successful_runs),
                    "avg_metrics": avg_metrics,
                    "overall_score": overall_score,
                    "individual_results": config_results,
                }
            else:
                # All runs failed
                result_summary = {
                    "name": config_name,
                    "configuration": config,
                    "success": False,
                    "runs_completed": 0,
                    "runs_failed": len(config_results),
                    "error": "All runs failed",
                    "individual_results": config_results,
                }

            results.append(result_summary)

    # Sort results by overall score (if available)
    sorted_results = sorted(
        [r for r in results if r.get("success", False)],
        key=lambda x: x.get("overall_score", 0),
        reverse=True,
    )

    # Add failed configurations at the end
    sorted_results.extend([r for r in results if not r.get("success", False)])

    # Create comparison report
    comparison_report = {
        "query": query,
        "configurations_tested": len(configurations),
        "successful_configurations": len(
            [r for r in results if r.get("success", False)]
        ),
        "failed_configurations": len(
            [r for r in results if not r.get("success", False)]
        ),
        "repetitions": repetitions,
        "metric_weights": metric_weights,
        "timestamp": datetime.now(UTC).isoformat(),
        "results": sorted_results,
    }

    # Save results to file
    from ...security.file_write_verifier import write_json_verified

    timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")
    result_file = str(Path(output_dir) / f"comparison_results_{timestamp}.json")

    write_json_verified(
        result_file,
        comparison_report,
        "benchmark.allow_file_output",
        context="comparison results",
    )

    # Generate visualizations
    visualizations_dir = Path(output_dir) / "visualizations"
    visualizations_dir.mkdir(parents=True, exist_ok=True)
    visualizations_dir = str(visualizations_dir)

    _create_comparison_visualizations(
        comparison_report, output_dir=visualizations_dir, timestamp=timestamp
    )

    logger.info(f"Comparison completed. Results saved to {result_file}")

    # Add report path to the result
    comparison_report["report_path"] = result_file

    return comparison_report


def _evaluate_single_configuration(
    query: str,
    config: Dict[str, Any],
    model_name: Optional[str] = None,
    provider: Optional[str] = None,
    search_tool: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Evaluate a single configuration.

    Args:
        query: Research query to evaluate
        config: Configuration parameters
        model_name: Name of the LLM model to use
        provider: LLM provider
        search_tool: Search engine to use

    Returns:
        Dictionary with evaluation results
    """
    # Extract configuration parameters
    config_model_name = config.get("model_name", model_name)
    config_provider = config.get("provider", provider)
    config_search_tool = config.get("search_tool", search_tool)
    config_iterations = config.get("iterations", 2)
    config_questions_per_iteration = config.get("questions_per_iteration", 2)
    config_search_strategy = config.get("search_strategy", "iterdrag")
    config_max_results = config.get("max_results", 50)
    config_max_filtered_results = config.get("max_filtered_results", 20)

    # Initialize profiling tools
    speed_profiler = SpeedProfiler()
    resource_monitor = ResourceMonitor(sampling_interval=0.5)

    # Start profiling
    speed_profiler.start()
    resource_monitor.start()

    try:
        # Get LLM
        with speed_profiler.timer("llm_initialization"):
            llm = get_llm(
                temperature=config.get("temperature", 0.7),
                model_name=config_model_name,
                provider=config_provider,
            )

        # Set up search engine if specified
        with speed_profiler.timer("search_initialization"):
            search = None
            if config_search_tool:
                search = get_search(
                    config_search_tool,
                    llm_instance=llm,
                    max_results=config_max_results,
                    max_filtered_results=config_max_filtered_results,
                )

        # Create search system
        system = AdvancedSearchSystem(llm=llm, search=search)
        system.max_iterations = config_iterations
        system.questions_per_iteration = config_questions_per_iteration
        system.strategy_name = config_search_strategy

        # Run the analysis
        with speed_profiler.timer("analysis"):
            results = system.analyze_topic(query)

        # Stop profiling
        speed_profiler.stop()
        resource_monitor.stop()

        # Calculate metrics
        quality_metrics = calculate_quality_metrics(
            results=results,
            system_info={
                "all_links_of_system": getattr(
                    system, "all_links_of_system", []
                )
            },
        )

        speed_metrics = calculate_speed_metrics(
            timing_info=speed_profiler.get_summary(),
            system_info={
                "iterations": config_iterations,
                "questions_per_iteration": config_questions_per_iteration,
                "results": results,
            },
        )

        resource_metrics = calculate_resource_metrics(
            resource_info=resource_monitor.get_combined_stats(),
            system_info={
                "iterations": config_iterations,
                "questions_per_iteration": config_questions_per_iteration,
                "results": results,
            },
        )

        # Return comprehensive results
        return {
            "query": query,
            "config": config,
            "success": True,
            "findings_count": len(results.get("findings", [])),
            "knowledge_length": len(results.get("current_knowledge", "")),
            "quality_metrics": quality_metrics,
            "speed_metrics": speed_metrics,
            "resource_metrics": resource_metrics,
            "timing_details": speed_profiler.get_timings(),
            "resource_details": resource_monitor.get_combined_stats(),
        }

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

        # Log the error
        logger.exception("Error evaluating configuration")

        # Return error information
        return {
            "query": query,
            "config": config,
            "success": False,
            "error": str(e),
            "timing_details": speed_profiler.get_timings(),
            "resource_details": resource_monitor.get_combined_stats(),
        }


def _calculate_average_metrics(results: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Calculate average metrics across multiple runs.

    Args:
        results: List of individual run results

    Returns:
        Dictionary with averaged metrics
    """
    # Check if there are any successful results
    if not results:
        return {}

    # Initialize average metrics
    avg_metrics = {
        "quality_metrics": {},
        "speed_metrics": {},
        "resource_metrics": {},
    }

    # Quality metrics
    quality_keys = set()
    for result in results:
        quality_metrics = result.get("quality_metrics", {})
        quality_keys.update(quality_metrics.keys())

    for key in quality_keys:
        values = [r.get("quality_metrics", {}).get(key) for r in results]
        values = [v for v in values if v is not None]
        if values:
            avg_metrics["quality_metrics"][key] = sum(values) / len(values)

    # Speed metrics
    speed_keys = set()
    for result in results:
        speed_metrics = result.get("speed_metrics", {})
        speed_keys.update(speed_metrics.keys())

    for key in speed_keys:
        values = [r.get("speed_metrics", {}).get(key) for r in results]
        values = [v for v in values if v is not None]
        if values:
            avg_metrics["speed_metrics"][key] = sum(values) / len(values)

    # Resource metrics
    resource_keys = set()
    for result in results:
        resource_metrics = result.get("resource_metrics", {})
        resource_keys.update(resource_metrics.keys())

    for key in resource_keys:
        values = [r.get("resource_metrics", {}).get(key) for r in results]
        values = [v for v in values if v is not None]
        if values:
            avg_metrics["resource_metrics"][key] = sum(values) / len(values)

    return avg_metrics


def _create_comparison_visualizations(
    comparison_report: Dict[str, Any], output_dir: str, timestamp: str
):
    """
    Create visualizations for the comparison results.

    Args:
        comparison_report: Comparison report dictionary
        output_dir: Directory to save visualizations
        timestamp: Timestamp string for filenames
    """
    # Check if there are successful results
    successful_results = [
        r
        for r in comparison_report.get("results", [])
        if r.get("success", False)
    ]

    if not successful_results:
        logger.warning("No successful configurations to visualize")
        return

    # Extract configuration names
    config_names = [
        r.get("name", f"Config {i + 1}")
        for i, r in enumerate(successful_results)
    ]

    # 1. Overall score comparison
    plt.figure(figsize=(12, 6))
    scores = [r.get("overall_score", 0) for r in successful_results]

    # Create horizontal bar chart
    plt.barh(config_names, scores, color="skyblue")
    plt.xlabel("Overall Score")
    plt.ylabel("Configuration")
    plt.title("Configuration Performance Comparison")
    plt.grid(axis="x", linestyle="--", alpha=0.7)
    plt.tight_layout()
    plt.savefig(
        str(Path(output_dir) / f"overall_score_comparison_{timestamp}.png")
    )
    plt.close()

    # 2. Quality metrics comparison
    quality_metrics = ["overall_quality", "source_count", "lexical_diversity"]
    _create_metric_comparison_chart(
        successful_results,
        config_names,
        quality_metrics,
        "quality_metrics",
        "Quality Metrics Comparison",
        str(Path(output_dir) / f"quality_metrics_comparison_{timestamp}.png"),
    )

    # 3. Speed metrics comparison
    speed_metrics = ["overall_speed", "total_duration", "duration_per_question"]
    _create_metric_comparison_chart(
        successful_results,
        config_names,
        speed_metrics,
        "speed_metrics",
        "Speed Metrics Comparison",
        str(Path(output_dir) / f"speed_metrics_comparison_{timestamp}.png"),
    )

    # 4. Resource metrics comparison
    resource_metrics = [
        "overall_resource",
        "process_memory_max_mb",
        "system_cpu_avg",
    ]
    _create_metric_comparison_chart(
        successful_results,
        config_names,
        resource_metrics,
        "resource_metrics",
        "Resource Usage Comparison",
        str(Path(output_dir) / f"resource_metrics_comparison_{timestamp}.png"),
    )

    # 5. Spider chart for multi-dimensional comparison
    _create_spider_chart(
        successful_results,
        config_names,
        str(Path(output_dir) / f"spider_chart_comparison_{timestamp}.png"),
    )

    # 6. Pareto frontier chart for quality vs. speed
    _create_pareto_chart(
        successful_results,
        str(Path(output_dir) / f"pareto_chart_comparison_{timestamp}.png"),
    )


def _create_metric_comparison_chart(
    results: List[Dict[str, Any]],
    config_names: List[str],
    metric_keys: List[str],
    metric_category: str,
    title: str,
    output_path: str,
):
    """
    Create a chart comparing specific metrics across configurations.

    Args:
        results: List of configuration results
        config_names: Names of configurations
        metric_keys: Keys of metrics to compare
        metric_category: Category of metrics (quality_metrics, speed_metrics, etc.)
        title: Chart title
        output_path: Path to save the chart
    """
    # Create figure with multiple subplots (one per metric)
    fig, axes = plt.subplots(
        len(metric_keys), 1, figsize=(12, 5 * len(metric_keys))
    )

    # Handle case with only one metric
    if len(metric_keys) == 1:
        axes = [axes]

    for i, metric_key in enumerate(metric_keys):
        ax = axes[i]

        # Get metric values
        metric_values = []
        for result in results:
            metrics = result.get("avg_metrics", {}).get(metric_category, {})
            value = metrics.get(metric_key)

            # Handle time values for better visualization
            if "duration" in metric_key and value is not None:
                # Convert to seconds if > 60 seconds, minutes if > 60 minutes
                if value > 3600:
                    value = value / 3600  # Convert to hours
                    metric_key += " (hours)"
                elif value > 60:
                    value = value / 60  # Convert to minutes
                    metric_key += " (minutes)"
                else:
                    metric_key += " (seconds)"

            metric_values.append(value if value is not None else 0)

        # Create horizontal bar chart
        bars = ax.barh(config_names, metric_values, color="lightblue")
        ax.set_xlabel(metric_key.replace("_", " ").title())
        ax.set_title(f"{metric_key.replace('_', ' ').title()}")
        ax.grid(axis="x", linestyle="--", alpha=0.7)

        # Add value labels to bars
        for bar in bars:
            width = bar.get_width()
            label_x_pos = width * 1.01
            ax.text(
                label_x_pos,
                bar.get_y() + bar.get_height() / 2,
                f"{width:.2f}",
                va="center",
            )

    plt.suptitle(title, fontsize=16)
    plt.tight_layout()
    plt.savefig(output_path)
    plt.close()


def _create_spider_chart(
    results: List[Dict[str, Any]], config_names: List[str], output_path: str
):
    """
    Create a spider chart comparing metrics across configurations.

    Args:
        results: List of configuration results
        config_names: Names of configurations
        output_path: Path to save the chart
    """
    # Try to import the radar chart module
    try:
        from matplotlib.path import Path
        from matplotlib.projections import register_projection
        from matplotlib.projections.polar import PolarAxes
        from matplotlib.spines import Spine

        def radar_factory(num_vars, frame="circle"):
            """Create a radar chart with `num_vars` axes."""
            # Calculate evenly-spaced axis angles
            theta = np.linspace(0, 2 * np.pi, num_vars, endpoint=False)

            class RadarAxes(PolarAxes):
                name = "radar"

                def __init__(self, *args, **kwargs):
                    super().__init__(*args, **kwargs)
                    self.set_theta_zero_location("N")

                def fill(self, *args, closed=True, **kwargs):
                    return super().fill(closed=closed, *args, **kwargs)

                def plot(self, *args, **kwargs):
                    return super().plot(*args, **kwargs)

                def set_varlabels(self, labels):
                    self.set_thetagrids(np.degrees(theta), labels)

                def _gen_axes_patch(self):
                    if frame == "circle":
                        return Circle((0.5, 0.5), 0.5)
                    elif frame == "polygon":
                        return RegularPolygon(
                            (0.5, 0.5), num_vars, radius=0.5, edgecolor="k"
                        )
                    else:
                        raise ValueError(
                            "Unknown value for 'frame': %s" % frame
                        )

                def _gen_axes_spines(self):
                    if frame == "circle":
                        return super()._gen_axes_spines()
                    elif frame == "polygon":
                        spine_type = Spine.circular_spine
                        verts = unit_poly_verts(num_vars)
                        vertices = [(0.5, 0.5)] + verts
                        codes = (
                            [Path.MOVETO]
                            + [Path.LINETO] * num_vars
                            + [Path.CLOSEPOLY]
                        )
                        path = Path(vertices, codes)
                        spine = Spine(self, spine_type, path)
                        spine.set_transform(self.transAxes)
                        return {"polar": spine}
                    else:
                        raise ValueError(
                            "Unknown value for 'frame': %s" % frame
                        )

            def unit_poly_verts(num_vars):
                """Return vertices of polygon for radar chart."""
                verts = []
                for i in range(num_vars):
                    angle = theta[i]
                    verts.append(
                        (0.5 * (1 + np.cos(angle)), 0.5 * (1 + np.sin(angle)))
                    )
                return verts

            register_projection(RadarAxes)
            return theta

        # Select metrics for the spider chart
        metrics = [
            {"name": "Quality", "key": "quality_metrics.overall_quality"},
            {"name": "Speed", "key": "speed_metrics.overall_speed"},
            {
                "name": "Sources",
                "key": "quality_metrics.normalized_source_count",
            },
            {
                "name": "Content",
                "key": "quality_metrics.normalized_knowledge_length",
            },
            {
                "name": "Memory",
                "key": "resource_metrics.normalized_memory_usage",
                "invert": True,
            },
        ]

        # Extract metric values
        spoke_labels = [m["name"] for m in metrics]
        num_vars = len(spoke_labels)
        theta = radar_factory(num_vars)

        fig, ax = plt.subplots(
            figsize=(10, 10), subplot_kw=dict(projection="radar")
        )

        # Color map for different configurations
        colors = plt.cm.viridis(np.linspace(0, 1, len(results)))

        for i, result in enumerate(results):
            values = []
            for metric in metrics:
                # Extract metric value using the key path (e.g., "quality_metrics.overall_quality")
                key_parts = metric["key"].split(".")
                value = result.get("avg_metrics", {})
                for part in key_parts:
                    value = value.get(part, 0) if isinstance(value, dict) else 0

                # Invert if needed (for metrics where lower is better)
                if metric.get("invert", False):
                    value = 1.0 - value

                values.append(value)

            # Plot this configuration
            ax.plot(
                theta,
                values,
                color=colors[i],
                linewidth=2,
                label=config_names[i],
            )
            ax.fill(theta, values, color=colors[i], alpha=0.25)

        # Set chart properties
        ax.set_varlabels(spoke_labels)
        plt.legend(loc="best", bbox_to_anchor=(0.5, 0.1))
        plt.title("Multi-Dimensional Configuration Comparison", size=16, y=1.05)
        plt.tight_layout()

        # Save chart
        plt.savefig(output_path)
        plt.close()

    except Exception as e:
        logger.exception("Error creating spider chart")
        # Create a text-based chart as fallback
        plt.figure(figsize=(10, 6))
        plt.text(
            0.5,
            0.5,
            f"Spider chart could not be created: {e!s}",
            horizontalalignment="center",
            verticalalignment="center",
        )
        plt.axis("off")
        plt.savefig(output_path)
        plt.close()


def _create_pareto_chart(results: List[Dict[str, Any]], output_path: str):
    """
    Create a Pareto frontier chart showing quality vs. speed tradeoff.

    Args:
        results: List of configuration results
        output_path: Path to save the chart
    """
    # Extract quality and speed metrics
    quality_scores = []
    speed_scores = []
    names = []

    for result in results:
        metrics = result.get("avg_metrics", {})
        quality = metrics.get("quality_metrics", {}).get("overall_quality", 0)

        # For speed, we use inverse of duration (so higher is better)
        duration = metrics.get("speed_metrics", {}).get("total_duration", 1)
        speed = 1.0 / max(duration, 0.001)  # Avoid division by zero

        quality_scores.append(quality)
        speed_scores.append(speed)
        names.append(result.get("name", "Configuration"))

    # Create scatter plot
    plt.figure(figsize=(10, 8))
    plt.scatter(quality_scores, speed_scores, s=100, alpha=0.7)

    # Add labels for each point
    for i, name in enumerate(names):
        plt.annotate(
            name,
            (quality_scores[i], speed_scores[i]),
            xytext=(5, 5),
            textcoords="offset points",
        )

    # Identify Pareto frontier
    pareto_points = []
    for i, (q, s) in enumerate(zip(quality_scores, speed_scores, strict=False)):
        is_pareto = True
        for q2, s2 in zip(quality_scores, speed_scores, strict=False):
            if q2 > q and s2 > s:  # Dominated
                is_pareto = False
                break
        if is_pareto:
            pareto_points.append(i)

    # Highlight Pareto frontier
    pareto_quality = [quality_scores[i] for i in pareto_points]
    pareto_speed = [speed_scores[i] for i in pareto_points]

    # Sort pareto points for line drawing
    pareto_sorted = sorted(
        zip(pareto_quality, pareto_speed, pareto_points, strict=False)
    )
    pareto_quality = [p[0] for p in pareto_sorted]
    pareto_speed = [p[1] for p in pareto_sorted]
    pareto_indices = [p[2] for p in pareto_sorted]

    # Draw Pareto frontier line
    plt.plot(pareto_quality, pareto_speed, "r--", linewidth=2)

    # Highlight Pareto optimal points
    plt.scatter(
        [quality_scores[i] for i in pareto_indices],
        [speed_scores[i] for i in pareto_indices],
        s=150,
        facecolors="none",
        edgecolors="r",
        linewidth=2,
    )

    # Add labels for Pareto optimal configurations
    for i in pareto_indices:
        plt.annotate(
            names[i],
            (quality_scores[i], speed_scores[i]),
            xytext=(8, 8),
            textcoords="offset points",
            bbox=dict(boxstyle="round,pad=0.5", fc="yellow", alpha=0.7),
        )

    # Set chart properties
    plt.xlabel("Quality Score (higher is better)")
    plt.ylabel("Speed Score (higher is better)")
    plt.title("Quality vs. Speed Tradeoff (Pareto Frontier)", size=14)
    plt.grid(True, linestyle="--", alpha=0.7)

    # Add explanation
    plt.figtext(
        0.5,
        0.01,
        "Points on the red line are Pareto optimal configurations\n"
        "(no other configuration is better in both quality and speed)",
        ha="center",
        fontsize=10,
        bbox=dict(boxstyle="round", fc="white", alpha=0.7),
    )

    plt.tight_layout()
    plt.savefig(output_path)
    plt.close()
