plot.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.gridspec import GridSpec
from scipy.signal import savgol_filter
import pandas as pd
import os
from double_lane_change import SharedControl_DLC
from lane_change import SharedControl_LC
import argparse

class Plotter:
    def __init__(self):
        self.metric_names = [
            'lateral_error',      # Path tracking
            'heading_error',      # Vehicle stability
            'human_control',      # Control effort from human
        ]
        self.setup_plotting_style()
        
    def setup_plotting_style(self):
        plt.style.use('seaborn-v0_8-paper')
        mpl.rcParams['font.family'] = 'serif'
        mpl.rcParams['font.serif'] = ['Times New Roman']
        mpl.rcParams['font.size'] = 14
        mpl.rcParams['axes.labelsize'] = 16
        mpl.rcParams['axes.titlesize'] = 16
        mpl.rcParams['xtick.labelsize'] = 16
        mpl.rcParams['ytick.labelsize'] = 16
        mpl.rcParams['legend.fontsize'] = 14
        mpl.rcParams['figure.titlesize'] = 16
        mpl.rcParams['axes.grid'] = True
        mpl.rcParams['grid.alpha'] = 0.3
        mpl.rcParams['lines.linewidth'] = 2.0
        mpl.rcParams['savefig.dpi'] = 600
        mpl.rcParams['savefig.bbox'] = 'tight'
        mpl.rcParams['savefig.pad_inches'] = 0.1
    
    def calculate_metrics(self, states, controls_D, controls_A, y_ref, psi_ref, t):
        metrics = {}

        # note: have to add x_error to it 
        
        lateral_error = states[:, 1] - y_ref
        heading_error = states[:, 2] - psi_ref
        
        metrics['time_series'] = {
            'lateral_error': lateral_error,
            'heading_error': heading_error,
            'controls_D': controls_D,
            'time': t
        }

        dt = t[1] - t[0]
        metrics['lateral_error_total'] = np.sum(lateral_error**2) * dt
        metrics['heading_error_total'] = np.sum(heading_error**2) * dt
        metrics['human_control_total'] = np.sum(controls_D**2) * dt
        
        metrics['lateral_error_max'] = np.max(np.abs(lateral_error))
        metrics['heading_error_max'] = np.max(np.abs(heading_error))
        metrics['human_control_max'] = np.max(np.abs(controls_D))
        
        return metrics
    
    def normalize_metrics(self, all_metrics, transition_types):
        normalized_metrics = {trans_type: {} for trans_type in transition_types}
        
        max_values = {}
        for metric in self.metric_names:
            max_values[f'{metric}_total'] = max(
                metrics[f'{metric}_total'] 
                for metrics in all_metrics.values()
            )

        for trans_type in transition_types:
            metrics = all_metrics[trans_type]
            total_scores = {}
            normalized_scores = {}
            
            for metric in self.metric_names:
                total_key = f'{metric}_total'
                total_scores[metric] = metrics[total_key]
                normalized_scores[metric] = metrics[total_key] / max_values[total_key]

            total_score = np.sum(list(total_scores.values()))
            normalized_score = np.sum(list(normalized_scores.values()))
            
            normalized_metrics[trans_type] = {
                'total_scores': total_scores,
                'individual_scores': normalized_scores,
                'total_score': total_score,
                'time_series': metrics['time_series'],
                **{k: v for k, v in metrics.items() if k not in ['individual_scores', 'total_score', 'total_scores']}
            }
        
        return normalized_metrics
    
    def apply_savgol_filter(self, data, window_length=101, polyorder=3):
        return savgol_filter(data, window_length=window_length, polyorder=polyorder)
    
    def plot_scores_with_variance(self, all_metrics, transition_types):
        fig, ax = plt.subplots(figsize=(10, 6), dpi=300)
        
        scores_data = {trans_type: [] for trans_type in transition_types}
        for trans_type in transition_types:
            for q_metrics in all_metrics.values():
                scores_data[trans_type].append(q_metrics[trans_type]['total_score'])
        
        x = np.arange(len(transition_types))
        means = [np.mean(scores_data[t]) for t in transition_types]
        stds = [np.std(scores_data[t]) for t in transition_types]
        
        colors = ['#7CB9E8', '#7CB9E8', '#7CB9E8', '#7CB9E8', '#7CB9E8', '#F0B27A']
        
        bars = ax.bar(x, means, yerr=stds, 
                    capsize=4, 
                    alpha=0.8,
                    width=0.6,
                    color=colors[:len(transition_types)],
                    edgecolor='grey',
                    linewidth=0.8,
                    error_kw={'ecolor': 'grey',
                            'capthick': 1,
                            'elinewidth': 1})
        
        ax.set_ylabel('Total Error', fontsize=20)
        ax.tick_params(axis='y', labelsize=17)
        
        ax.set_xticks(x)
        ax.set_xticklabels([t.capitalize() for t in transition_types], 
                        rotation=30, fontsize=20)
        
        ax.yaxis.grid(True, linestyle='--', alpha=0.3, color='grey')
        ax.set_axisbelow(True)
        
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)

        for spine in ['left', 'bottom']:
            ax.spines[spine].set_linewidth(0.8)
            ax.spines[spine].set_color('grey')
        
        plt.tight_layout()
        
        return fig
    
    def create_metrics_table_for_q(self, q_id, metrics, transition_types):
        """
        Create a metrics table for a specific Q matrix
        
        Args:
            q_id: The Q-matrix identifier (e.g., 'Q0')
            metrics: Metrics for this Q-matrix
            transition_types: List of transition types
        
        Returns:
            DataFrame: Formatted table with metrics for each transition type
        """
        # Create a table with total metrics
        data_total = []
        for trans_type in transition_types:
            trans_metrics = metrics[trans_type]
            row = [trans_type.capitalize()]
            
            # Add raw metrics
            for metric in self.metric_names:
                row.append(trans_metrics[f'{metric}_total'])
            
            # Add max metrics
            for metric in self.metric_names:
                row.append(trans_metrics[f'{metric}_max'])
            
            # Add normalized score
            row.append(trans_metrics['total_score'])
            
            data_total.append(row)
        
        # Create column names
        columns = ['Transition Type']
        
        # Total metric columns
        for metric in self.metric_names:
            columns.append(f'{metric.replace("_", " ").title()} (Total)')
        
        # Max metric columns
        for metric in self.metric_names:
            columns.append(f'{metric.replace("_", " ").title()} (Max)')
        
        # Normalized score column
        columns.append('Total Score')
        
        # Create DataFrame and format
        df = pd.DataFrame(data_total, columns=columns)
        df = df.round(4)
        
        return df
    
    def save_metrics_to_csv(self, all_metrics, transition_types, Q_matrices, key):
        """
        Save metrics to CSV files
        
        Args:
            all_metrics: Dictionary with metrics for all Q matrices
            transition_types: List of transition types
            Q_matrices: List of Q matrices
            key: Scenario key ('lane_change' or 'double_lane_change')
        """
        # Create directory for metrics
        metrics_dir = f"plots/{key}/metrics"
        os.makedirs(metrics_dir, exist_ok=True)
        
        # Save aggregated metrics
        data = []
        for q_id in all_metrics.keys():
            for trans_type in transition_types:
                metrics = all_metrics[q_id][trans_type]
                row = {
                    'Q_Matrix': q_id,
                    'Transition Type': trans_type.capitalize(),
                    'Total Score': metrics['total_score']
                }
                for metric in self.metric_names:
                    row[f'{metric.replace("_", " ").title()} Score'] = \
                        metrics['individual_scores'][metric]
                    
                    # Add raw metrics
                    row[f'{metric.replace("_", " ").title()} Total'] = \
                        metrics[f'{metric}_total']
                    row[f'{metric.replace("_", " ").title()} Max'] = \
                        metrics[f'{metric}_max']
                
                data.append(row)
        
        df = pd.DataFrame(data)
        df.to_csv(f"{metrics_dir}/all_metrics.csv", index=False)
        
        # Save individual Q matrix metrics
        for q_id in all_metrics.keys():
            df = self.create_metrics_table_for_q(q_id, all_metrics[q_id], transition_types)
            df.to_csv(f"{metrics_dir}/{q_id}_metrics.csv", index=False)
        
        return df
    
    def print_summary_statistics(self, all_metrics, transition_types):
        """
        Print summary statistics for all Q matrices and transition types
        
        Args:
            all_metrics: Dictionary with metrics for all Q matrices
            transition_types: List of transition types
        """
        print("\n" + "="*100)
        print(f"PERFORMANCE STATISTICS ACROSS {len(all_metrics)} Q MATRICES")
        print("="*100)
        
        # Print summary by transition type
        for trans_type in transition_types:
            scores = [metrics[trans_type]['total_score'] 
                     for metrics in all_metrics.values()]
            
            print(f"\n{trans_type.capitalize()} Transition:")
            print(f"  Mean Score: {np.mean(scores):.4f}")
            print(f"  Std Dev:    {np.std(scores):.4f}")
            print(f"  Min Score:  {np.min(scores):.4f}")
            print(f"  Max Score:  {np.max(scores):.4f}")
    
    def print_individual_q_statistics(self, q_id, metrics, transition_types):
        """
        Print statistics for an individual Q matrix
        
        Args:
            q_id: The Q-matrix identifier (e.g., 'Q0')
            metrics: Metrics for this Q-matrix
            transition_types: List of transition types
        """
        print("\n" + "-"*100)
        print(f"PERFORMANCE STATISTICS FOR {q_id}")
        print("-"*100)
        
        # Create a formatted table with metrics
        df = self.create_metrics_table_for_q(q_id, metrics, transition_types)
        
        # Print the table
        print(f"\nMetrics by Transition Type:")
        print(df.to_string(index=False))
        
        # Calculate best and worst transition types
        scores = {trans_type: metrics[trans_type]['total_score'] for trans_type in transition_types}
        best_trans = min(scores, key=scores.get)
        worst_trans = max(scores, key=scores.get)
        
        print(f"\nBest transition type: {best_trans.capitalize()} (Score: {scores[best_trans]:.4f})")
        print(f"Worst transition type: {worst_trans.capitalize()} (Score: {scores[worst_trans]:.4f})")
        
        # Print metrics by category
        print("\nMetric Breakdown by Category:")
        for metric in self.metric_names:
            print(f"\n{metric.replace('_', ' ').title()}:")
            for trans_type in transition_types:
                print(f"  {trans_type.capitalize()}: Total = {metrics[trans_type][f'{metric}_total']:.4f}, " +
                      f"Max = {metrics[trans_type][f'{metric}_max']:.4f}")
    
    def plot_combined_for_q_matrix(self, q_id, results, transition_types, type):
        """
        Create a combined plot for each transition type in a Q-matrix showing trajectory, control input, and authority
        
        Args:
            q_id: The Q-matrix identifier (e.g., 'Q0')
            results: Dictionary of results for this Q-matrix
            transition_types: List of transition types
            type: 'LC' for lane change or 'DLC' for double lane change
        """
        if type == 'LC':
            key = 'lane_change'
            title_prefix = 'Lane Change'
        elif type == 'DLC':
            key = 'double_lane_change'
            title_prefix = 'double_lane_change'
        
        # Create output directory structure
        plot_dir = f"plots/{key}/{q_id}/transitions"
        os.makedirs(plot_dir, exist_ok=True)
        
        for trans_type in transition_types:
            states, controls_D, controls_A, alphas_D, alphas_A, t, y_ref, _, _, _ = results[trans_type]
            controls_D, controls_A = controls_D[:, 0], controls_A[:, 0]
            
            fig = plt.figure(figsize=(8, 10))
            gs = GridSpec(3, 1, figure=fig, height_ratios=[1, 1, 1])
            
            # Plot 1: Trajectory
            ax1 = fig.add_subplot(gs[0])
            ax1.plot(t, y_ref, 'k--', linewidth=1.5, label='Reference')
            ax1.plot(t, states[:, 1], 'b-', linewidth=2.0, label='Actual')
            ax1.set_ylabel('Lateral Position (m)')
            ax1.set_title(f'{title_prefix} Trajectory - {trans_type.capitalize()}')
            ax1.legend(loc='best', frameon=True, fancybox=True, framealpha=0.7)
            ax1.set_xticklabels([])
            
            # Plot 2: Control inputs
            ax2 = fig.add_subplot(gs[1])
            # steering control only
            controls_D = np.ravel(controls_D)
            controls_A = np.ravel(controls_A)
            controls_D_smooth = self.apply_savgol_filter(controls_D)
            controls_A_smooth = self.apply_savgol_filter(controls_A)
            
            ax2.plot(t[:-1], controls_D_smooth, 'b-', label='$u_D$')
            ax2.plot(t[:-1], controls_A_smooth, 'r-', label='$u_A$')
            ax2.set_ylabel('Steering Angle (rad)')
            ax2.legend(loc='best', frameon=True, fancybox=True, framealpha=0.7)
            ax2.set_xticklabels([])
            
            # Plot 3: Authority transition
            ax3 = fig.add_subplot(gs[2])
            ax3.plot(t[:-1], alphas_D, 'b-', label='Driver')
            ax3.plot(t[:-1], alphas_A, 'r-', label='ADAS')
            ax3.set_ylabel('Control Authority')
            ax3.set_xlabel('Time (s)')
            ax3.legend(loc='best', frameon=True, fancybox=True, framealpha=0.7)
            
            plt.tight_layout()
            plt.savefig(f"{plot_dir}/{trans_type}.png")
            plt.close(fig)
    
    def compare_all_trajectories_for_q(self, q_id, results, transition_types, type):
        """
        Create a plot comparing all trajectories for a single Q-matrix
        
        Args:
            q_id: The Q-matrix identifier (e.g., 'Q0')
            results: Dictionary of results for this Q-matrix
            transition_types: List of transition types
            type: 'LC' for lane change or 'DLC' for double lane change
        """
        if type == 'LC':
            key = 'lane_change'
            title_prefix = 'Lane Change'
        elif type == 'DLC':
            key = 'double_lane_change'
            title_prefix = 'double_lane_change'
        
        # Create output directory
        plot_dir = f"plots/{key}/{q_id}/comparisons"
        os.makedirs(plot_dir, exist_ok=True)
        
        fig = plt.figure(figsize=(10, 6))
        
        _, _, _, _, _, t, y_ref, _, _, _ = results[transition_types[0]]
        plt.plot(t, y_ref, 'k--', linewidth=1.5, label='Reference')
        
        colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']
        
        for i, trans_type in enumerate(transition_types):
            states, _, _, _, _, t, _, _, _, _ = results[trans_type]
            plt.plot(t, states[:, 1], color=colors[i % len(colors)], alpha=0.8,
                    linewidth=2.0, label=f'{trans_type.capitalize()}')
        
        plt.xlabel('Time (s)')
        plt.ylabel('Lateral Position (m)')
        plt.title(f'Trajectory Comparison - {q_id}')
        plt.legend(loc='best', frameon=True, fancybox=True, framealpha=0.7)
        plt.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig(f"{plot_dir}/trajectory_comparison.png")
        plt.close(fig)
    
    def compare_all_controls_for_q(self, q_id, results, transition_types, type):
        """
        Create plots comparing control inputs for a single Q-matrix
        
        Args:
            q_id: The Q-matrix identifier (e.g., 'Q0')
            results: Dictionary of results for this Q-matrix
            transition_types: List of transition types
            type: 'LC' for lane change or 'DLC' for double lane change
        """
        if type == 'LC':
            key = 'lane_change'
            title_prefix = 'Lane Change'
        elif type == 'DLC':
            key = 'double_lane_change'
            title_prefix = 'Double Lane Change'
        
        # Create output directory
        plot_dir = f"plots/{key}/{q_id}/comparisons"
        os.makedirs(plot_dir, exist_ok=True)
        
        # Combined plot showing both driver and ADAS for each transition type
        fig_combined = plt.figure(figsize=(12, 8))
        
        rows = len(transition_types)//2 + len(transition_types)%2
        cols = 2 if len(transition_types) > 1 else 1
        
        for i, trans_type in enumerate(transition_types):
            plt.subplot(rows, cols, i+1)
            
            _, controls_D, controls_A, _, _, t, _, _, _, _ = results[trans_type]
            controls_D, controls_A = controls_D[:, 0], controls_A[:, 0]
            controls_D = np.ravel(controls_D)
            controls_A = np.ravel(controls_A)
            controls_D_smooth = self.apply_savgol_filter(controls_D)
            controls_A_smooth = self.apply_savgol_filter(controls_A)
            
            plt.plot(t[:-1], controls_D_smooth, 'b-', linewidth=1.5, label='$u_D$', alpha=0.8)
            plt.plot(t[:-1], controls_A_smooth, 'r-', linewidth=1.5, label='$u_A$', alpha=0.8)
            
            plt.title(f'{trans_type.capitalize()}')
            plt.xlabel('Time (s)')
            plt.ylabel('Steering angle (rad)')
            plt.legend()
            plt.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.suptitle(f'Control Input Comparison - {q_id}', fontsize=16, y=1.02)
        plt.savefig(f"{plot_dir}/control_comparison.png")
        plt.close(fig_combined)
    
    def plot_both_scenarios_comparison(self, lc_metrics, dlc_metrics, transition_types):
        """
        Create a bar chart comparing both scenarios (Lane Change and Double Lane Change) for each transition type
        
        Args:
            lc_metrics: Normalized metrics for lane change scenario
            dlc_metrics: Normalized metrics for double lane change scenario
            transition_types: List of transition types to compare
        """
        fig, ax = plt.subplots(figsize=(12, 7), dpi=300)
        
        x = np.arange(len(transition_types))
        width = 0.35
        
        lc_means = []
        lc_stds = []
        dlc_means = []
        dlc_stds = []
        
        for trans_type in transition_types:
            lc_scores = [q_metrics[trans_type]['total_score'] for q_metrics in lc_metrics.values()]
            dlc_scores = [q_metrics[trans_type]['total_score'] for q_metrics in dlc_metrics.values()]
            
            lc_means.append(np.mean(lc_scores))
            lc_stds.append(np.std(lc_scores))
            dlc_means.append(np.mean(dlc_scores))
            dlc_stds.append(np.std(dlc_scores))
        
        lc_bars = ax.bar(x - width/2, lc_means, width, yerr=lc_stds, 
                         capsize=4, alpha=0.9, color='#7CB9E8', edgecolor='grey',
                         linewidth=0.9, error_kw={'ecolor': 'black', 'capthick': 1, 'elinewidth': 1},
                         label='Lane Change')
        
        dlc_bars = ax.bar(x + width/2, dlc_means, width, yerr=dlc_stds, 
                        capsize=4, alpha=0.95, color='#F0B27A', edgecolor='grey',
                        linewidth=0.9, hatch='\\\\', error_kw={'ecolor': 'black', 'capthick': 1, 'elinewidth': 1},
                        label='Double Lane Change')
        
        ax.set_ylabel('Total Error', fontsize=20)
        ax.tick_params(axis='y', labelsize=17)
        
        ax.set_xticks(x)
        ax.set_xticklabels([t.capitalize() for t in transition_types], rotation=30, fontsize=20)
        
        ax.legend(fontsize=18, frameon=True, fancybox=True, framealpha=0.9, loc="upper right")
        
        ax.yaxis.grid(True, linestyle='--', alpha=0.3, color='grey')
        ax.xaxis.grid(False)
        ax.set_axisbelow(True)
        
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        
        for spine in ['left', 'bottom']:
            ax.spines[spine].set_linewidth(0.8)
            ax.spines[spine].set_color('grey')
        
        plt.tight_layout()
        
        return fig
    
    def compare_all_metrics(self, transition_types, controllers, Q_human_values, type):
        """
        Compare metrics across all Q matrices and transition types
        
        Args:
            transition_types: List of transition types to compare
            controllers: Dictionary of controllers (LC and/or DLC)
            Q_human_values: List of Q matrices
            type: 'LC' for lane change or 'DLC' for double lane change
        
        Returns:
            Dictionary: All normalized metrics
        """
        if type == 'LC':    
            key = 'lane_change'
        elif type == 'DLC':   
            key = 'double_lane_change'

        # Create main output directories
        os.makedirs(f"plots/{key}", exist_ok=True)
        
        all_results = {f'Q{i}': {} for i in range(len(Q_human_values))}
        raw_metrics = {f'Q{i}': {} for i in range(len(Q_human_values))}
        
        print(f"\n{'-'*80}")
        print(f"RUNNING SIMULATION FOR {key.upper()} SCENARIO")
        print(f"{'-'*80}")
        
        for i, Q in enumerate(Q_human_values):
            print(f"\nSimulating with Q{i}...")
            
            if type == 'LC':
                controllers['LC'].QD_max = Q
                results = controllers['LC'].do_lane_change(transition_types)
            else:
                controllers['DLC'].QD_max = Q
                results = controllers['DLC'].do_double_lane_change(transition_types)
            
            all_results[f'Q{i}'] = results
            
            # Process and save metrics for this Q matrix
            q_id = f'Q{i}'
            
            # Create directories for this Q matrix
            os.makedirs(f"plots/{key}/{q_id}", exist_ok=True)
            os.makedirs(f"plots/{key}/{q_id}/transitions", exist_ok=True)
            os.makedirs(f"plots/{key}/{q_id}/comparisons", exist_ok=True)
            
            # Generate plots for this Q matrix
            self.plot_combined_for_q_matrix(q_id, results, transition_types, type)
            self.compare_all_trajectories_for_q(q_id, results, transition_types, type)
            self.compare_all_controls_for_q(q_id, results, transition_types, type)
            
            # Calculate metrics for this Q matrix
            for trans_type in transition_types:
                states, controls_D, controls_A, _, _, t, y_ref, psi_ref, _, _ = results[trans_type]
                controls_D, controls_A = controls_D[:, 0], controls_A[:, 0]
                metrics = self.calculate_metrics(states, controls_D, controls_A, y_ref, psi_ref, t)
                raw_metrics[f'Q{i}'][trans_type] = metrics

        # Normalize and aggregate metrics
        all_metrics = {q_id: self.normalize_metrics(raw_metrics[q_id], transition_types) 
                      for q_id in all_results.keys()}
        
        # Save overall results plot
        fig = self.plot_scores_with_variance(all_metrics, transition_types)
        fig.savefig(f"plots/{key}/overall_comparison.png", bbox_inches='tight')

        # Save metrics to CSV
        self.save_metrics_to_csv(all_metrics, transition_types, Q_human_values, key)
        
        # Print summary statistics for all Q matrices
        self.print_summary_statistics(all_metrics, transition_types)
        
        # Print individual Q matrix statistics
        for q_id, metrics in all_metrics.items():
            self.print_individual_q_statistics(q_id, metrics, transition_types)
        
        return all_metrics
    
    def extract_plots(self, transition_types, scenario, log, Q_human_values=None):
        """
        Main function to extract plots and analyze results
        
        Args:
            transition_types: List of transition types to compare
            scenario: 'LC' for lane change, 'DLC' for double lane change, or 'both'
            log: Boolean indicating whether to enable logging
            Q_human_values: Optional list of Q matrices to use
        """
        # Default Q matrices if none provided
        if Q_human_values is None:
            Q_human_values = [
                np.diag([0, 5.21, 11.64, 0]),
                np.diag([0, 4.36, 8.07, 0]),
                np.diag([0, 4.30, 10.90, 0]),
                np.diag([0, 5.30, 9.70, 0]),
            ]
        
        # Create base directories
        os.makedirs("plots", exist_ok=True)
        if scenario in ['LC', 'both']:
            os.makedirs("plots/lane_change", exist_ok=True)
        if scenario in ['DLC', 'both']:
            os.makedirs("plots/double_lane_change", exist_ok=True)
        if scenario == 'both':
            os.makedirs("plots/combined", exist_ok=True)
        
        # Initialize controllers
        controllers = {}
        if scenario in ['LC', 'both']:
            controllers['LC'] = SharedControl_LC(log)
        if scenario in ['DLC', 'both']:
            controllers['DLC'] = SharedControl_DLC(log)

        # Run simulations and generate plots
        lc_metrics = None
        dlc_metrics = None
        
        if scenario == 'LC':
            lc_metrics = self.compare_all_metrics(transition_types, controllers, Q_human_values, "LC")
        elif scenario == 'DLC':
            dlc_metrics = self.compare_all_metrics(transition_types, controllers, Q_human_values, "DLC")
        elif scenario == 'both':
            print("\n" + "="*100)
            print("RUNNING LANE CHANGE SCENARIO")
            print("="*100)
            lc_metrics = self.compare_all_metrics(transition_types, controllers, Q_human_values, "LC")
            
            print("\n" + "="*100)
            print("RUNNING DOUBLE LANE CHANGE SCENARIO")
            print("="*100)
            dlc_metrics = self.compare_all_metrics(transition_types, controllers, Q_human_values, "DLC")
            
            # Create combined comparison plot
            print("\n" + "="*100)
            print("GENERATING SCENARIO COMPARISON PLOTS")
            combined_fig = self.plot_both_scenarios_comparison(lc_metrics, dlc_metrics, transition_types)
            combined_fig.savefig(f"plots/combined/scenario_comparison.png", bbox_inches='tight')
            print("DONE!")
            print("="*100)
            
        else:
            raise ValueError("Invalid scenario. Choose from 'LC' or 'DLC'")
        
        print("\n" + "="*100)
        print("SIMULATION COMPLETE")
        print("="*100)
        print(f"\nResults saved in 'plots/' directory")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Run and analyze shared control simulations with enhanced plotting.')
    parser.add_argument('--scenario', choices=['LC', 'DLC', 'both'], default='DLC',
                        help='Choose the driving scenario: LC (Lane Change), DLC (Double Lane Change), or both')
    parser.add_argument('--transitions', nargs='+', default=['step', 'linear', 'cooperative', 'sigmoid', 'lagged', 'adaptive'],
                        help='List of transition types to evaluate: step, linear, cooperative, sigmoid, exponential, blend, lagged, adaptive')
    parser.add_argument('--log', action='store_true', help='Enable logging for control state')
    args = parser.parse_args()

    plotter = Plotter()
    plotter.extract_plots(args.transitions, args.scenario, args.log)