ComputerAssignments/CA24_multi_agent_systems/src/utils/visualization.py

import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
from loguru import logger

class VisualizationManager:

    def __init__(self, output_dir: str = "results/plots"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)


        plt.style.use("seaborn-v0_8")
        sns.set_palette("husl")

        logger.info(
            f"VisualizationManager initialized with output directory: {self.output_dir}"
        )

    def plot_agent_performance_comparison(
        self,
        results: Dict[str, List[Dict]],
        save_file: str = "agent_performance_comparison.png",
    ) -> str:


        data = []
        for agent_name, agent_results in results.items():
            for result in agent_results:
                data.append(
                    {
                        "Agent": agent_name,
                        "Scenario": result.get("scenario_name", "Unknown"),
                        "Success": result.get("success", False),
                        "Execution Time": result.get("execution_time", 0),
                        "Steps": result.get("steps_taken", 0),
                        "Cost": result.get("total_cost", 0),
                    }
                )

        df = pd.DataFrame(data)


        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        fig.suptitle("Agent Performance Comparison", fontsize=16, fontweight="bold")


        success_rates = df.groupby("Agent")["Success"].mean()
        axes[0, 0].bar(success_rates.index, success_rates.values, color="skyblue")
        axes[0, 0].set_title("Success Rate by Agent")
        axes[0, 0].set_ylabel("Success Rate")
        axes[0, 0].set_ylim(0, 1)


        avg_times = df.groupby("Agent")["Execution Time"].mean()
        axes[0, 1].bar(avg_times.index, avg_times.values, color="lightcoral")
        axes[0, 1].set_title("Average Execution Time by Agent")
        axes[0, 1].set_ylabel("Time (seconds)")


        avg_steps = df.groupby("Agent")["Steps"].mean()
        axes[1, 0].bar(avg_steps.index, avg_steps.values, color="lightgreen")
        axes[1, 0].set_title("Average Steps by Agent")
        axes[1, 0].set_ylabel("Number of Steps")


        avg_costs = df.groupby("Agent")["Cost"].mean()
        axes[1, 1].bar(avg_costs.index, avg_costs.values, color="gold")
        axes[1, 1].set_title("Average Cost by Agent")
        axes[1, 1].set_ylabel("Cost ($)")


        for ax in axes.flat:
            ax.tick_params(axis="x", rotation=45)

        plt.tight_layout()


        filepath = self.output_dir / save_file
        plt.savefig(filepath, dpi=300, bbox_inches="tight")
        plt.close()

        logger.info(f"Performance comparison plot saved to {filepath}")
        return str(filepath)

    def plot_scenario_difficulty_analysis(
        self,
        results: Dict[str, List[Dict]],
        save_file: str = "scenario_difficulty_analysis.png",
    ) -> str:


        scenario_data = {}
        for agent_name, agent_results in results.items():
            for result in agent_results:
                scenario_name = result.get("scenario_name", "Unknown")
                if scenario_name not in scenario_data:
                    scenario_data[scenario_name] = []
                scenario_data[scenario_name].append(result.get("success", False))


        scenario_success_rates = {}
        for scenario, successes in scenario_data.items():
            scenario_success_rates[scenario] = sum(successes) / len(successes)


        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))


        scenarios = list(scenario_success_rates.keys())
        success_rates = list(scenario_success_rates.values())

        colors = [
            "red" if rate < 0.5 else "orange" if rate < 0.8 else "green"
            for rate in success_rates
        ]
        bars = ax1.bar(scenarios, success_rates, color=colors)
        ax1.set_title("Scenario Difficulty Analysis")
        ax1.set_ylabel("Success Rate")
        ax1.set_ylim(0, 1)
        ax1.tick_params(axis="x", rotation=45)


        for bar, rate in zip(bars, success_rates):
            ax1.text(
                bar.get_x() + bar.get_width() / 2,
                bar.get_height() + 0.01,
                f"{rate:.2f}",
                ha="center",
                va="bottom",
            )


        difficulty_counts = {"Easy": 0, "Medium": 0, "Hard": 0}
        for rate in success_rates:
            if rate >= 0.8:
                difficulty_counts["Easy"] += 1
            elif rate >= 0.5:
                difficulty_counts["Medium"] += 1
            else:
                difficulty_counts["Hard"] += 1

        ax2.pie(
            difficulty_counts.values(),
            labels=difficulty_counts.keys(),
            autopct="%1.1f%%",
            colors=["green", "orange", "red"],
        )
        ax2.set_title("Difficulty Distribution")

        plt.tight_layout()


        filepath = self.output_dir / save_file
        plt.savefig(filepath, dpi=300, bbox_inches="tight")
        plt.close()

        logger.info(f"Scenario difficulty analysis plot saved to {filepath}")
        return str(filepath)

    def plot_agent_evolution_timeline(
        self,
        results: Dict[str, List[Dict]],
        save_file: str = "agent_evolution_timeline.png",
    ) -> str:


        fig = make_subplots(
            rows=2,
            cols=2,
            subplot_titles=(
                "Success Rate Evolution",
                "Execution Time Evolution",
                "Steps Evolution",
                "Cost Evolution",
            ),
            specs=[
                [{"secondary_y": False}, {"secondary_y": False}],
                [{"secondary_y": False}, {"secondary_y": False}],
            ],
        )


        agent_order = [
            "ReAct (v1.0)",
            "Reflexion (v2.0)",
            "HippoRAG (v3.0)",
            "Nexus (v4.0)",
        ]


        metrics = {
            "success_rate": [],
            "avg_execution_time": [],
            "avg_steps": [],
            "avg_cost": [],
        }

        for agent_name in agent_order:
            if agent_name in results:
                agent_results = results[agent_name]


                success_rate = sum(
                    r.get("success", False) for r in agent_results
                ) / len(agent_results)
                avg_time = sum(r.get("execution_time", 0) for r in agent_results) / len(
                    agent_results
                )
                avg_steps = sum(r.get("steps_taken", 0) for r in agent_results) / len(
                    agent_results
                )
                avg_cost = sum(r.get("total_cost", 0) for r in agent_results) / len(
                    agent_results
                )

                metrics["success_rate"].append(success_rate)
                metrics["avg_execution_time"].append(avg_time)
                metrics["avg_steps"].append(avg_steps)
                metrics["avg_cost"].append(avg_cost)


        fig.add_trace(
            go.Scatter(
                x=agent_order,
                y=metrics["success_rate"],
                mode="lines+markers",
                name="Success Rate",
                line=dict(color="green", width=3),
            ),
            row=1,
            col=1,
        )

        fig.add_trace(
            go.Scatter(
                x=agent_order,
                y=metrics["avg_execution_time"],
                mode="lines+markers",
                name="Execution Time",
                line=dict(color="red", width=3),
            ),
            row=1,
            col=2,
        )

        fig.add_trace(
            go.Scatter(
                x=agent_order,
                y=metrics["avg_steps"],
                mode="lines+markers",
                name="Steps",
                line=dict(color="blue", width=3),
            ),
            row=2,
            col=1,
        )

        fig.add_trace(
            go.Scatter(
                x=agent_order,
                y=metrics["avg_cost"],
                mode="lines+markers",
                name="Cost",
                line=dict(color="orange", width=3),
            ),
            row=2,
            col=2,
        )


        fig.update_layout(
            title_text="Agent Evolution Timeline", showlegend=False, height=800
        )


        fig.update_xaxes(title_text="Agent Version", row=1, col=1)
        fig.update_xaxes(title_text="Agent Version", row=1, col=2)
        fig.update_xaxes(title_text="Agent Version", row=2, col=1)
        fig.update_xaxes(title_text="Agent Version", row=2, col=2)

        fig.update_yaxes(title_text="Success Rate", row=1, col=1)
        fig.update_yaxes(title_text="Time (seconds)", row=1, col=2)
        fig.update_yaxes(title_text="Number of Steps", row=2, col=1)
        fig.update_yaxes(title_text="Cost ($)", row=2, col=2)


        filepath = self.output_dir / save_file
        fig.write_html(str(filepath).replace(".png", ".html"))
        fig.write_image(str(filepath), width=1200, height=800)

        logger.info(f"Agent evolution timeline plot saved to {filepath}")
        return str(filepath)

    def plot_capability_radar_chart(
        self,
        results: Dict[str, List[Dict]],
        save_file: str = "capability_radar_chart.png",
    ) -> str:


        capabilities = {
            "Basic Reasoning": {"weight": 1.0, "key": "success"},
            "Tool Integration": {"weight": 0.8, "key": "tools_used"},
            "Error Recovery": {"weight": 0.9, "key": "error_recovery"},
            "Learning": {"weight": 0.7, "key": "learning"},
            "Memory": {"weight": 0.6, "key": "memory"},
            "Scalability": {"weight": 0.8, "key": "scalability"},
        }


        agent_scores = {}

        for agent_name, agent_results in results.items():
            scores = {}


            total_tests = len(agent_results)
            successful_tests = sum(1 for r in agent_results if r.get("success", False))


            scores["Basic Reasoning"] = (
                successful_tests / total_tests if total_tests > 0 else 0
            )


            scores["Tool Integration"] = min(
                1.0, successful_tests / max(total_tests * 0.5, 1)
            )


            if "Reflexion" in agent_name:
                scores["Error Recovery"] = 0.8
                scores["Learning"] = 0.9
            else:
                scores["Error Recovery"] = 0.3
                scores["Learning"] = 0.2


            if "HippoRAG" in agent_name:
                scores["Memory"] = 0.9
            elif "Nexus" in agent_name:
                scores["Memory"] = 0.7
            else:
                scores["Memory"] = 0.2


            if "Nexus" in agent_name:
                scores["Scalability"] = 0.9
            else:
                scores["Scalability"] = 0.4

            agent_scores[agent_name] = scores


        fig = go.Figure()

        capability_names = list(capabilities.keys())

        for agent_name, scores in agent_scores.items():
            values = [scores.get(cap, 0) for cap in capability_names]
            values += values[:1]

            fig.add_trace(
                go.Scatterpolar(
                    r=values,
                    theta=capability_names + [capability_names[0]],
                    fill="toself",
                    name=agent_name,
                    line=dict(width=2),
                )
            )

        fig.update_layout(
            polar=dict(radialaxis=dict(visible=True, range=[0, 1])),
            showlegend=True,
            title="Agent Capability Comparison",
            height=600,
        )


        filepath = self.output_dir / save_file
        fig.write_html(str(filepath).replace(".png", ".html"))
        fig.write_image(str(filepath), width=800, height=600)

        logger.info(f"Capability radar chart saved to {filepath}")
        return str(filepath)

    def plot_efficiency_heatmap(
        self, results: Dict[str, List[Dict]], save_file: str = "efficiency_heatmap.png"
    ) -> str:


        agent_names = list(results.keys())
        scenario_names = set()

        for agent_results in results.values():
            for result in agent_results:
                scenario_names.add(result.get("scenario_name", "Unknown"))

        scenario_names = sorted(list(scenario_names))


        efficiency_matrix = np.zeros((len(agent_names), len(scenario_names)))

        for i, agent_name in enumerate(agent_names):
            agent_results = results[agent_name]
            for result in agent_results:
                scenario_name = result.get("scenario_name", "Unknown")
                if scenario_name in scenario_names:
                    j = scenario_names.index(scenario_name)


                    success = 1.0 if result.get("success", False) else 0.0
                    time_score = max(
                        0, 1.0 - result.get("execution_time", 0) / 300.0
                    )
                    cost_score = max(
                        0, 1.0 - result.get("total_cost", 0) / 5000.0
                    )

                    efficiency = success * 0.5 + time_score * 0.3 + cost_score * 0.2
                    efficiency_matrix[i, j] = efficiency


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

        sns.heatmap(
            efficiency_matrix,
            xticklabels=scenario_names,
            yticklabels=agent_names,
            annot=True,
            fmt=".2f",
            cmap="RdYlGn",
            vmin=0,
            vmax=1,
            ax=ax,
        )

        ax.set_title("Agent Efficiency Heatmap", fontsize=16, fontweight="bold")
        ax.set_xlabel("Scenarios")
        ax.set_ylabel("Agents")


        plt.xticks(rotation=45, ha="right")
        plt.yticks(rotation=0)

        plt.tight_layout()


        filepath = self.output_dir / save_file
        plt.savefig(filepath, dpi=300, bbox_inches="tight")
        plt.close()

        logger.info(f"Efficiency heatmap saved to {filepath}")
        return str(filepath)

    def generate_all_visualizations(self, results: Dict[str, List[Dict]]) -> List[str]:

        logger.info("Generating all visualization plots...")

        plots = []


        plots.append(self.plot_agent_performance_comparison(results))
        plots.append(self.plot_scenario_difficulty_analysis(results))
        plots.append(self.plot_agent_evolution_timeline(results))
        plots.append(self.plot_capability_radar_chart(results))
        plots.append(self.plot_efficiency_heatmap(results))

        logger.info(f"Generated {len(plots)} visualization plots")
        return plots