plotManager.py

import os
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import csv

class PlotManager:
    def __init__(self, results_dir='results', bin_threshold=100, bin_size=10):
        self.results_dir = results_dir
        self.bin_threshold = bin_threshold
        self.bin_size = bin_size
        os.makedirs(self.results_dir, exist_ok=True)

    def plot_metric(self, values, ylabel, title, filename, smooth=False, bin=None, agent_label = "Agent"):
        os.makedirs(os.path.dirname(filename), exist_ok=True)
        episodes = list(range(len(values)))

        # Auto-enable binning if episode count > threshold
        if bin is None:
            bin = len(values) > self.bin_threshold

        if smooth:
            values = pd.Series(values).rolling(window=10, min_periods=1).mean().tolist()

        if bin:
            binned_episodes = list(range(0, len(values), self.bin_size))
            binned_values = [np.mean(values[i:i+self.bin_size]) for i in binned_episodes]
            episodes = binned_episodes
            values = binned_values

        plt.figure(figsize=(12, 6))
        plt.plot(episodes, values, label= f"{agent_label} (Smoothed)" if smooth else f"{agent_label}", marker='o' if bin else None)

        # Highlight best episode
        if ylabel.lower() == f"delay {agent_label}":
            best_idx = np.argmin(values)
            best_label = f"Lowest {agent_label}"
        else:
            best_idx = np.argmax(values)
            best_label = f"Highest ({agent_label})"
        best_episode = episodes[best_idx]
        best_value = values[best_idx]
        plt.scatter(best_episode, best_value, color='red', zorder=5, label=f'{best_label}: {best_value:.2f}')
        plt.annotate(f"{best_value:.2f}", (best_episode, best_value), textcoords="offset points", xytext=(0,10), ha='center')

        plt.xlabel(f"Episodes ({agent_label})")
        plt.ylabel(ylabel)
        plt.title(title)
        plt.grid(alpha=0.3)
        plt.legend()
        plt.xticks(np.linspace(0, max(episodes), num=10, dtype=int))
        plt.tight_layout()
        plt.savefig(filename)
        plt.close()
        print(f"📊 Saved enhanced plot: {filename}")



    def generate_csv_plots(self, episode, scenario, variant, suffix, agent_label = "Agent"):
        csv_dir = os.path.join(self.results_dir, 'plots', 'csv', f'{agent_label}_{variant}_{episode}_episodes')
        csv_files = [
            f for f in os.listdir(csv_dir)
            if f.startswith(f"episode_stats_{variant}_scenario_{scenario}_{episode}_episodes{suffix}") and f.endswith('.csv')
        ]
        for file in csv_files:
            df = pd.read_csv(os.path.join(csv_dir, file))
            base = os.path.splitext(file)[0]
            base_path = os.path.join(csv_dir, f"{variant}_scenario_{scenario}_{episode}_episodes{suffix}")

            self.generate_all_plots(df, base_path)

    def _get_unit_label(self, value_col):
        unit_labels = {
            "Avg_Delay": " (ms)",
            "QoS_Success": " (packets)",
            "Served": " (packets)",
            "Arrived": " (packets)",
            "Dropped": " (packets)",
            "Success_Rate": " (%)",
            "Failure_Rate": " (%)"
        }
        return unit_labels.get(value_col, "")

    def _aggregate_and_plot(self, df, value_col, ylabel, filename):
        episode_count = df['Episode'].nunique()
        unit = self._get_unit_label(value_col)
        full_ylabel = ylabel + unit

        if episode_count > self.bin_threshold:
            df['Episode_Bin'] = (df['Episode'] // self.bin_size) * self.bin_size
            pivot = df.pivot_table(index='Episode_Bin', columns='Type', values=value_col, aggfunc='mean')
            xlabel = f"Episode Bins (every {self.bin_size})"
        else:
            pivot = df.pivot_table(index='Episode', columns='Type', values=value_col, aggfunc='mean')
            xlabel = "Episode"

        plot_type = 'bar' if len(pivot) <= 50 else 'line'

        plt.figure(figsize=(12, 6))
        ax = plt.gca()
        if plot_type == 'line':
            pivot.plot(kind='line', marker='o', ax=ax)
        else:
            pivot.plot(kind='bar', ax=ax)

        ax.set_ylabel(full_ylabel)
        ax.set_xlabel(xlabel)
        ax.set_title(f"{ylabel} over {xlabel}")
        ax.grid(True)
        ax.legend(title="Traffic Type")

        if len(pivot) > 30:
            ax.set_xticks(ax.get_xticks()[::2])
            plt.xticks(rotation=45)

        plt.tight_layout()
        os.makedirs(os.path.dirname(filename), exist_ok=True)
        plt.savefig(filename)
        plt.close()
        print(f"📈 Saved CSV plot: {filename}")


    def save_csv_with_header(self, path, header):
        """
        Safely creates the directory for the file and writes the header row to a CSV.

        Args:
            path (str): Full path to the CSV file to be created.
            header (list[str]): List of column names for the CSV header.
        """
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(header)

    def plot_traffic_stats(self, df, base_path):
        os.makedirs(base_path, exist_ok=True)
        metrics = ["Arrived", "Served", "QoS_Success", "Dropped"]
        for metric in metrics:
            filename = os.path.join(base_path, f"{metric.lower()}_per_type.png")
            self._aggregate_and_plot(df, metric, metric, filename)
            

    def plot_success_failure_rates(self, df, base_path):
        os.makedirs(base_path, exist_ok=True)
        for metric in ["Success_Rate", "Failure_Rate"]:
            filename = os.path.join(base_path, f"{metric.lower()}_per_type.png")
            self._aggregate_and_plot(df, metric, metric, filename)
            

    def plot_combined_metrics(self, df, base_path, metrics=None):
        if metrics is None:
            metrics = ["Avg_Delay", "QoS_Success", "Served", "Dropped"]

        unit_labels = {
            "Avg_Delay": " (ms)",
            "QoS_Success": " (packets)",
            "Served": " (packets)",
            "Arrived": " (packets)",
            "Dropped": " (packets)",
            "Success_Rate": " (%)",
            "Failure_Rate": " (%)"
        }

        fig, axs = plt.subplots(len(metrics), 1, figsize=(12, 5 * len(metrics)))

        for i, metric in enumerate(metrics):
            ax = axs[i] if len(metrics) > 1 else axs
            episode_count = df['Episode'].nunique()

            if episode_count > self.bin_threshold:
                df['Episode_Bin'] = (df['Episode'] // self.bin_size) * self.bin_size
                pivot = df.pivot_table(index='Episode_Bin', columns='Type', values=metric, aggfunc='mean')
                xlabel = f"Episode Bins (every {self.bin_size})"
            else:
                pivot = df.pivot_table(index='Episode', columns='Type', values=metric, aggfunc='mean')
                xlabel = "Episode"

            unit = unit_labels.get(metric, "")
            pivot.plot(marker='o', ax=ax)
            ax.set_ylabel(metric + unit)
            ax.set_xlabel(xlabel)
            ax.set_title(f"{metric} over {xlabel}")
            ax.grid(True)
            ax.legend(title="Traffic Type")

        plt.tight_layout()
        os.makedirs(base_path, exist_ok=True)
        save_path = os.path.join(base_path, "combined_metrics.png")
        plt.savefig(save_path)
        plt.close()
        print(f"📊 Combined plot saved to: {save_path}")

    def generate_all_plots(self, df, base_path):
        self.plot_traffic_stats(df, base_path)
        self.plot_success_failure_rates(df, base_path)
        self.plot_combined_metrics(df, base_path)