#!/usr/bin/env python3
"""
Rollout Trace Visualization Tool

This script analyzes and visualizes agent rollout system traces logged by the AgentOrchestrator.
It provides comprehensive analysis of rollout performance, timing, and success metrics.

Usage:
    python rollout_trace_viz.py <csv_file_path> [--output-dir <output_directory>]

Example:
    python rollout_trace_viz.py logs/agentarium_rollout_traces_20231201_120000.csv --output-dir ./plots
"""

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import argparse
import os
import json
from datetime import datetime
from typing import Dict, List, Optional, Any
import warnings
warnings.filterwarnings('ignore')

class RolloutTraceAnalyzer:
    """Analyzer for agent rollout system traces."""
    
    def __init__(self, csv_filepath: str):
        """
        Initialize the analyzer with a CSV file of rollout traces.
        
        Args:
            csv_filepath: Path to the CSV file containing rollout traces
        """
        self.csv_filepath = csv_filepath
        self.df = None
        self.load_data()
        
    def load_data(self):
        """Load and preprocess the rollout trace data."""
        try:
            self.df = pd.read_csv(self.csv_filepath)
            print(f"Loaded {len(self.df)} rollout traces from {self.csv_filepath}")
            
            # Convert timestamps to datetime
            if 'start_time' in self.df.columns:
                self.df['start_datetime'] = pd.to_datetime(self.df['start_time'], unit='s')
            if 'end_time' in self.df.columns:
                self.df['end_datetime'] = pd.to_datetime(self.df['end_time'], unit='s')
                
            # Parse metadata JSON if present
            if 'metadata' in self.df.columns:
                self.df['metadata_parsed'] = self.df['metadata'].apply(self._parse_json_safely)
                
            # Fill NaN values for numeric columns
            numeric_columns = ['duration', 'num_turns', 'input_length', 'output_length', 
                             'input_token_count', 'output_token_count', 'final_reward']
            for col in numeric_columns:
                if col in self.df.columns:
                    self.df[col] = self.df[col].fillna(0)
                    
            print(f"Data preprocessing complete. Columns: {list(self.df.columns)}")
            
        except Exception as e:
            print(f"Error loading data: {e}")
            raise
            
    def _parse_json_safely(self, json_str: str) -> Dict[str, Any]:
        """Safely parse JSON string, return empty dict if parsing fails."""
        if pd.isna(json_str) or json_str == "":
            return {}
        try:
            return json.loads(json_str)
        except:
            return {}
    
    def print_summary_stats(self):
        """Print summary statistics of the rollout traces."""
        print("\n" + "="*60)
        print("ROLLOUT TRACE SUMMARY STATISTICS")
        print("="*60)
        
        if self.df is None or len(self.df) == 0:
            print("No data available for analysis.")
            return
            
        # Basic counts
        total_traces = len(self.df)
        completed_traces = len(self.df[self.df['status'] == 'completed'])
        failed_traces = len(self.df[self.df['status'] == 'failed'])
        running_traces = len(self.df[self.df['status'] == 'running'])
        
        print(f"Total Traces: {total_traces}")
        print(f"Completed: {completed_traces} ({completed_traces/total_traces*100:.1f}%)")
        print(f"Failed: {failed_traces} ({failed_traces/total_traces*100:.1f}%)")
        print(f"Running: {running_traces} ({running_traces/total_traces*100:.1f}%)")
        
        # Timing statistics (for completed traces)
        completed_df = self.df[self.df['status'] == 'completed']
        if len(completed_df) > 0:
            print(f"\nTIMING STATISTICS (Completed Traces)")
            print(f"Average Duration: {completed_df['duration'].mean():.2f} seconds")
            print(f"Median Duration: {completed_df['duration'].median():.2f} seconds")
            print(f"Min Duration: {completed_df['duration'].min():.2f} seconds")
            print(f"Max Duration: {completed_df['duration'].max():.2f} seconds")
            print(f"Total Runtime: {completed_df['duration'].sum():.2f} seconds")
            
            # Turn statistics
            if 'num_turns' in completed_df.columns:
                print(f"\nTURN STATISTICS")
                print(f"Average Turns: {completed_df['num_turns'].mean():.1f}")
                print(f"Median Turns: {completed_df['num_turns'].median():.1f}")
                print(f"Max Turns: {completed_df['num_turns'].max()}")
                
            # Length statistics
            if 'input_length' in completed_df.columns and 'output_length' in completed_df.columns:
                print(f"\nLENGTH STATISTICS")
                print(f"Average Input Length: {completed_df['input_length'].mean():.0f} chars")
                print(f"Average Output Length: {completed_df['output_length'].mean():.0f} chars")
                print(f"Total Input Length: {completed_df['input_length'].sum():,} chars")
                print(f"Total Output Length: {completed_df['output_length'].sum():,} chars")
                
            # Token statistics
            if 'input_token_count' in completed_df.columns and 'output_token_count' in completed_df.columns:
                print(f"\nTOKEN STATISTICS")
                print(f"Average Input Tokens: {completed_df['input_token_count'].mean():.0f}")
                print(f"Average Output Tokens: {completed_df['output_token_count'].mean():.0f}")
                print(f"Total Input Tokens: {completed_df['input_token_count'].sum():,}")
                print(f"Total Output Tokens: {completed_df['output_token_count'].sum():,}")
                
            # Reward statistics
            if 'final_reward' in completed_df.columns:
                reward_df = completed_df[completed_df['final_reward'].notna()]
                if len(reward_df) > 0:
                    print(f"\nREWARD STATISTICS")
                    print(f"Average Reward: {reward_df['final_reward'].mean():.3f}")
                    print(f"Median Reward: {reward_df['final_reward'].median():.3f}")
                    print(f"Min Reward: {reward_df['final_reward'].min():.3f}")
                    print(f"Max Reward: {reward_df['final_reward'].max():.3f}")
                    
        # Agent type breakdown
        if 'agent_type' in self.df.columns:
            print(f"\nAGENT TYPE BREAKDOWN")
            agent_counts = self.df['agent_type'].value_counts()
            for agent_type, count in agent_counts.items():
                print(f"  {agent_type}: {count} traces")
                
        # Model breakdown
        if 'model' in self.df.columns:
            print(f"\nMODEL BREAKDOWN")
            model_counts = self.df['model'].value_counts()
            for model, count in model_counts.items():
                print(f"  {model}: {count} traces")
    
    def create_visualizations(self, output_dir: str = "plots"):
        """Create comprehensive visualizations of the rollout traces."""
        if self.df is None or len(self.df) == 0:
            print("No data available for visualization.")
            return
            
        # Create output directory
        os.makedirs(output_dir, exist_ok=True)
        print(f"\nCreating visualizations in {output_dir}/")
        
        # Set style
        plt.style.use('default')
        sns.set_palette("husl")
        
        # 1. Status Distribution Pie Chart
        self._plot_status_distribution(output_dir)
        
        # 2. Duration Distribution Histogram
        self._plot_duration_distribution(output_dir)
        
        # 3. Performance Over Time
        self._plot_performance_over_time(output_dir)
        
        # 4. Turn Count Analysis
        self._plot_turn_analysis(output_dir)
        
        # 5. Input/Output Length Analysis
        self._plot_length_analysis(output_dir)
        
        # 6. Reward Distribution
        self._plot_reward_analysis(output_dir)
        
        # 7. Agent Type Performance
        self._plot_agent_type_performance(output_dir)
        
        # 8. Model Performance Comparison
        self._plot_model_performance(output_dir)
        
        # 9. Correlation Heatmap
        self._plot_correlation_heatmap(output_dir)
        
        # 10. Comprehensive Dashboard
        self._create_dashboard(output_dir)
        
        print(f"All visualizations saved to {output_dir}/")
    
    def _plot_status_distribution(self, output_dir: str):
        """Plot the distribution of rollout statuses."""
        plt.figure(figsize=(8, 6))
        status_counts = self.df['status'].value_counts()
        colors = ['#2ecc71', '#e74c3c', '#f39c12', '#3498db']
        
        plt.pie(status_counts.values, labels=status_counts.index, autopct='%1.1f%%', 
                colors=colors[:len(status_counts)], startangle=90)
        plt.title('Rollout Status Distribution', fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'status_distribution.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    def _plot_duration_distribution(self, output_dir: str):
        """Plot the distribution of rollout durations."""
        completed_df = self.df[self.df['status'] == 'completed']
        if len(completed_df) == 0:
            return
            
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
        
        # Histogram
        ax1.hist(completed_df['duration'], bins=30, alpha=0.7, color='skyblue', edgecolor='black')
        ax1.set_xlabel('Duration (seconds)')
        ax1.set_ylabel('Frequency')
        ax1.set_title('Duration Distribution')
        ax1.grid(True, alpha=0.3)
        
        # Box plot
        ax2.boxplot(completed_df['duration'], vert=True, patch_artist=True, 
                   boxprops=dict(facecolor='lightblue', alpha=0.7))
        ax2.set_ylabel('Duration (seconds)')
        ax2.set_title('Duration Box Plot')
        ax2.grid(True, alpha=0.3)
        
        plt.suptitle('Rollout Duration Analysis', fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'duration_distribution.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    def _plot_performance_over_time(self, output_dir: str):
        """Plot performance metrics over time."""
        if 'start_datetime' not in self.df.columns:
            return
            
        completed_df = self.df[self.df['status'] == 'completed'].copy()
        if len(completed_df) == 0:
            return
            
        completed_df = completed_df.sort_values('start_datetime')
        
        fig, axes = plt.subplots(2, 2, figsize=(16, 12))
        
        # Duration over time
        axes[0, 0].plot(completed_df['start_datetime'], completed_df['duration'], 
                       alpha=0.6, marker='o', markersize=3)
        axes[0, 0].set_title('Duration Over Time')
        axes[0, 0].set_ylabel('Duration (seconds)')
        axes[0, 0].tick_params(axis='x', rotation=45)
        axes[0, 0].grid(True, alpha=0.3)
        
        # Cumulative completed traces
        completed_df['cumulative_count'] = range(1, len(completed_df) + 1)
        axes[0, 1].plot(completed_df['start_datetime'], completed_df['cumulative_count'], 
                       color='green', linewidth=2)
        axes[0, 1].set_title('Cumulative Completed Traces')
        axes[0, 1].set_ylabel('Count')
        axes[0, 1].tick_params(axis='x', rotation=45)
        axes[0, 1].grid(True, alpha=0.3)
        
        # Turns over time
        if 'num_turns' in completed_df.columns:
            axes[1, 0].plot(completed_df['start_datetime'], completed_df['num_turns'], 
                           alpha=0.6, marker='s', markersize=3, color='orange')
            axes[1, 0].set_title('Number of Turns Over Time')
            axes[1, 0].set_ylabel('Turns')
            axes[1, 0].tick_params(axis='x', rotation=45)
            axes[1, 0].grid(True, alpha=0.3)
        
        # Reward over time
        if 'final_reward' in completed_df.columns:
            reward_df = completed_df[completed_df['final_reward'].notna()]
            if len(reward_df) > 0:
                axes[1, 1].plot(reward_df['start_datetime'], reward_df['final_reward'], 
                               alpha=0.6, marker='^', markersize=3, color='red')
                axes[1, 1].set_title('Reward Over Time')
                axes[1, 1].set_ylabel('Final Reward')
                axes[1, 1].tick_params(axis='x', rotation=45)
                axes[1, 1].grid(True, alpha=0.3)
        
        plt.suptitle('Performance Metrics Over Time', fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'performance_over_time.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    def _plot_turn_analysis(self, output_dir: str):
        """Plot analysis of turn counts."""
        if 'num_turns' not in self.df.columns:
            return
            
        completed_df = self.df[self.df['status'] == 'completed']
        if len(completed_df) == 0:
            return
            
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
        
        # Turn count distribution
        turn_counts = completed_df['num_turns'].value_counts().sort_index()
        ax1.bar(turn_counts.index, turn_counts.values, alpha=0.7, color='lightcoral')
        ax1.set_xlabel('Number of Turns')
        ax1.set_ylabel('Frequency')
        ax1.set_title('Turn Count Distribution')
        ax1.grid(True, alpha=0.3)
        
        # Duration vs Turns scatter plot
        ax2.scatter(completed_df['num_turns'], completed_df['duration'], alpha=0.6, s=50)
        ax2.set_xlabel('Number of Turns')
        ax2.set_ylabel('Duration (seconds)')
        ax2.set_title('Duration vs Number of Turns')
        ax2.grid(True, alpha=0.3)
        
        # Add correlation coefficient
        if len(completed_df) > 1:
            correlation = np.corrcoef(completed_df['num_turns'], completed_df['duration'])[0, 1]
            ax2.text(0.05, 0.95, f'Correlation: {correlation:.3f}', transform=ax2.transAxes, 
                    bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        
        plt.suptitle('Turn Analysis', fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'turn_analysis.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    def _plot_length_analysis(self, output_dir: str):
        """Plot analysis of input/output lengths."""
        if 'input_length' not in self.df.columns or 'output_length' not in self.df.columns:
            return
            
        completed_df = self.df[self.df['status'] == 'completed']
        if len(completed_df) == 0:
            return
            
        fig, axes = plt.subplots(2, 2, figsize=(16, 12))
        
        # Input length distribution
        axes[0, 0].hist(completed_df['input_length'], bins=30, alpha=0.7, color='lightblue')
        axes[0, 0].set_xlabel('Input Length (characters)')
        axes[0, 0].set_ylabel('Frequency')
        axes[0, 0].set_title('Input Length Distribution')
        axes[0, 0].grid(True, alpha=0.3)
        
        # Output length distribution
        axes[0, 1].hist(completed_df['output_length'], bins=30, alpha=0.7, color='lightgreen')
        axes[0, 1].set_xlabel('Output Length (characters)')
        axes[0, 1].set_ylabel('Frequency')
        axes[0, 1].set_title('Output Length Distribution')
        axes[0, 1].grid(True, alpha=0.3)
        
        # Input vs Output scatter
        axes[1, 0].scatter(completed_df['input_length'], completed_df['output_length'], alpha=0.6)
        axes[1, 0].set_xlabel('Input Length (characters)')
        axes[1, 0].set_ylabel('Output Length (characters)')
        axes[1, 0].set_title('Input vs Output Length')
        axes[1, 0].grid(True, alpha=0.3)
        
        # Length vs Duration
        total_length = completed_df['input_length'] + completed_df['output_length']
        axes[1, 1].scatter(total_length, completed_df['duration'], alpha=0.6, color='purple')
        axes[1, 1].set_xlabel('Total Length (characters)')
        axes[1, 1].set_ylabel('Duration (seconds)')
        axes[1, 1].set_title('Total Length vs Duration')
        axes[1, 1].grid(True, alpha=0.3)
        
        plt.suptitle('Input/Output Length Analysis', fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'length_analysis.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    def _plot_reward_analysis(self, output_dir: str):
        """Plot reward distribution and analysis."""
        if 'final_reward' not in self.df.columns:
            return
            
        completed_df = self.df[self.df['status'] == 'completed']
        reward_df = completed_df[completed_df['final_reward'].notna()]
        
        if len(reward_df) == 0:
            return
            
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
        
        # Reward distribution
        ax1.hist(reward_df['final_reward'], bins=30, alpha=0.7, color='gold')
        ax1.set_xlabel('Final Reward')
        ax1.set_ylabel('Frequency')
        ax1.set_title('Reward Distribution')
        ax1.grid(True, alpha=0.3)
        
        # Reward vs Duration
        ax2.scatter(reward_df['final_reward'], reward_df['duration'], alpha=0.6, color='crimson')
        ax2.set_xlabel('Final Reward')
        ax2.set_ylabel('Duration (seconds)')
        ax2.set_title('Reward vs Duration')
        ax2.grid(True, alpha=0.3)
        
        plt.suptitle('Reward Analysis', fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'reward_analysis.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    def _plot_agent_type_performance(self, output_dir: str):
        """Plot performance comparison across agent types."""
        if 'agent_type' not in self.df.columns:
            return
            
        completed_df = self.df[self.df['status'] == 'completed']
        if len(completed_df) == 0:
            return
            
        agent_types = completed_df['agent_type'].unique()
        if len(agent_types) < 2:
            return  # Skip if only one agent type
            
        fig, axes = plt.subplots(2, 2, figsize=(16, 12))
        
        # Duration by agent type
        agent_durations = [completed_df[completed_df['agent_type'] == agent]['duration'].values 
                          for agent in agent_types]
        axes[0, 0].boxplot(agent_durations, labels=agent_types, patch_artist=True)
        axes[0, 0].set_ylabel('Duration (seconds)')
        axes[0, 0].set_title('Duration by Agent Type')
        axes[0, 0].grid(True, alpha=0.3)
        
        # Success rate by agent type
        success_rates = []
        for agent in agent_types:
            agent_df = self.df[self.df['agent_type'] == agent]
            success_rate = len(agent_df[agent_df['status'] == 'completed']) / len(agent_df) * 100
            success_rates.append(success_rate)
        
        axes[0, 1].bar(agent_types, success_rates, alpha=0.7, color='lightseagreen')
        axes[0, 1].set_ylabel('Success Rate (%)')
        axes[0, 1].set_title('Success Rate by Agent Type')
        axes[0, 1].grid(True, alpha=0.3)
        
        # Turn count by agent type
        if 'num_turns' in completed_df.columns:
            agent_turns = [completed_df[completed_df['agent_type'] == agent]['num_turns'].values 
                          for agent in agent_types]
            axes[1, 0].boxplot(agent_turns, labels=agent_types, patch_artist=True)
            axes[1, 0].set_ylabel('Number of Turns')
            axes[1, 0].set_title('Turns by Agent Type')
            axes[1, 0].grid(True, alpha=0.3)
        
        # Output length by agent type
        if 'output_length' in completed_df.columns:
            agent_outputs = [completed_df[completed_df['agent_type'] == agent]['output_length'].values 
                           for agent in agent_types]
            axes[1, 1].boxplot(agent_outputs, labels=agent_types, patch_artist=True)
            axes[1, 1].set_ylabel('Output Length (characters)')
            axes[1, 1].set_title('Output Length by Agent Type')
            axes[1, 1].grid(True, alpha=0.3)
        
        plt.suptitle('Agent Type Performance Comparison', fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'agent_type_performance.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    def _plot_model_performance(self, output_dir: str):
        """Plot performance comparison across models."""
        if 'model' not in self.df.columns:
            return
            
        completed_df = self.df[self.df['status'] == 'completed']
        if len(completed_df) == 0:
            return
            
        models = completed_df['model'].unique()
        if len(models) < 2:
            return  # Skip if only one model
            
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
        
        # Duration by model
        model_durations = [completed_df[completed_df['model'] == model]['duration'].values 
                          for model in models]
        ax1.boxplot(model_durations, labels=models, patch_artist=True)
        ax1.set_ylabel('Duration (seconds)')
        ax1.set_title('Duration by Model')
        ax1.tick_params(axis='x', rotation=45)
        ax1.grid(True, alpha=0.3)
        
        # Success rate by model
        success_rates = []
        for model in models:
            model_df = self.df[self.df['model'] == model]
            success_rate = len(model_df[model_df['status'] == 'completed']) / len(model_df) * 100
            success_rates.append(success_rate)
        
        ax2.bar(models, success_rates, alpha=0.7, color='mediumpurple')
        ax2.set_ylabel('Success Rate (%)')
        ax2.set_title('Success Rate by Model')
        ax2.tick_params(axis='x', rotation=45)
        ax2.grid(True, alpha=0.3)
        
        plt.suptitle('Model Performance Comparison', fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'model_performance.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    def _plot_correlation_heatmap(self, output_dir: str):
        """Plot correlation heatmap of numeric metrics."""
        completed_df = self.df[self.df['status'] == 'completed']
        if len(completed_df) == 0:
            return
            
        # Select numeric columns for correlation
        numeric_columns = ['duration', 'num_turns', 'input_length', 'output_length', 
                          'input_token_count', 'output_token_count', 'final_reward']
        available_columns = [col for col in numeric_columns if col in completed_df.columns]
        
        if len(available_columns) < 2:
            return
            
        correlation_df = completed_df[available_columns].corr()
        
        plt.figure(figsize=(10, 8))
        sns.heatmap(correlation_df, annot=True, cmap='coolwarm', center=0, 
                   square=True, linewidths=0.5)
        plt.title('Correlation Heatmap of Rollout Metrics', fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'correlation_heatmap.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    def _create_dashboard(self, output_dir: str):
        """Create a comprehensive dashboard with key metrics."""
        completed_df = self.df[self.df['status'] == 'completed']
        
        fig = plt.figure(figsize=(20, 16))
        gs = fig.add_gridspec(4, 4, hspace=0.3, wspace=0.3)
        
        # Status pie chart (top-left)
        ax1 = fig.add_subplot(gs[0, 0])
        status_counts = self.df['status'].value_counts()
        colors = ['#2ecc71', '#e74c3c', '#f39c12', '#3498db']
        ax1.pie(status_counts.values, labels=status_counts.index, autopct='%1.1f%%', 
               colors=colors[:len(status_counts)], startangle=90)
        ax1.set_title('Status Distribution', fontweight='bold')
        
        # Duration histogram (top-center-left)
        ax2 = fig.add_subplot(gs[0, 1])
        if len(completed_df) > 0:
            ax2.hist(completed_df['duration'], bins=20, alpha=0.7, color='skyblue')
        ax2.set_xlabel('Duration (s)')
        ax2.set_ylabel('Count')
        ax2.set_title('Duration Distribution', fontweight='bold')
        ax2.grid(True, alpha=0.3)
        
        # Performance over time (top-center-right)
        ax3 = fig.add_subplot(gs[0, 2])
        if len(completed_df) > 0 and 'start_datetime' in completed_df.columns:
            completed_sorted = completed_df.sort_values('start_datetime')
            completed_sorted['cumulative'] = range(1, len(completed_sorted) + 1)
            ax3.plot(completed_sorted['start_datetime'], completed_sorted['cumulative'], 
                    color='green', linewidth=2)
        ax3.set_ylabel('Completed Count')
        ax3.set_title('Cumulative Progress', fontweight='bold')
        ax3.tick_params(axis='x', rotation=45)
        ax3.grid(True, alpha=0.3)
        
        # Key metrics text (top-right)
        ax4 = fig.add_subplot(gs[0, 3])
        ax4.axis('off')
        if len(completed_df) > 0:
            metrics_text = f"""
KEY METRICS

Total Traces: {len(self.df)}
Completed: {len(completed_df)}
Success Rate: {len(completed_df)/len(self.df)*100:.1f}%

Avg Duration: {completed_df['duration'].mean():.2f}s
Median Duration: {completed_df['duration'].median():.2f}s
Total Runtime: {completed_df['duration'].sum():.1f}s
            """
            if 'num_turns' in completed_df.columns:
                metrics_text += f"\nAvg Turns: {completed_df['num_turns'].mean():.1f}"
            if 'final_reward' in completed_df.columns:
                reward_df = completed_df[completed_df['final_reward'].notna()]
                if len(reward_df) > 0:
                    metrics_text += f"\nAvg Reward: {reward_df['final_reward'].mean():.3f}"
        else:
            metrics_text = "No completed traces available"
        
        ax4.text(0.05, 0.95, metrics_text, transform=ax4.transAxes, fontsize=12,
                verticalalignment='top', bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.8))
        
        # Additional visualizations for remaining subplots
        if len(completed_df) > 0:
            # Turn distribution (row 2, col 1)
            if 'num_turns' in completed_df.columns:
                ax5 = fig.add_subplot(gs[1, 0])
                turn_counts = completed_df['num_turns'].value_counts().sort_index()
                ax5.bar(turn_counts.index, turn_counts.values, alpha=0.7, color='orange')
                ax5.set_xlabel('Turns')
                ax5.set_ylabel('Count')
                ax5.set_title('Turn Distribution', fontweight='bold')
                ax5.grid(True, alpha=0.3)
            
            # Length scatter (row 2, col 2)
            if 'input_length' in completed_df.columns and 'output_length' in completed_df.columns:
                ax6 = fig.add_subplot(gs[1, 1])
                ax6.scatter(completed_df['input_length'], completed_df['output_length'], alpha=0.6)
                ax6.set_xlabel('Input Length')
                ax6.set_ylabel('Output Length')
                ax6.set_title('Input vs Output Length', fontweight='bold')
                ax6.grid(True, alpha=0.3)
            
            # Duration vs turns (row 2, col 3)
            if 'num_turns' in completed_df.columns:
                ax7 = fig.add_subplot(gs[1, 2])
                ax7.scatter(completed_df['num_turns'], completed_df['duration'], alpha=0.6, color='red')
                ax7.set_xlabel('Turns')
                ax7.set_ylabel('Duration (s)')
                ax7.set_title('Duration vs Turns', fontweight='bold')
                ax7.grid(True, alpha=0.3)
            
            # Reward distribution (row 2, col 4)
            if 'final_reward' in completed_df.columns:
                reward_df = completed_df[completed_df['final_reward'].notna()]
                if len(reward_df) > 0:
                    ax8 = fig.add_subplot(gs[1, 3])
                    ax8.hist(reward_df['final_reward'], bins=20, alpha=0.7, color='gold')
                    ax8.set_xlabel('Final Reward')
                    ax8.set_ylabel('Count')
                    ax8.set_title('Reward Distribution', fontweight='bold')
                    ax8.grid(True, alpha=0.3)
        
        plt.suptitle('Rollout Trace Analysis Dashboard', fontsize=20, fontweight='bold', y=0.95)
        plt.savefig(os.path.join(output_dir, 'dashboard.png'), dpi=300, bbox_inches='tight')
        plt.close()


def main():
    parser = argparse.ArgumentParser(description='Analyze and visualize agent rollout traces')
    parser.add_argument('csv_file', help='Path to the CSV file containing rollout traces')
    parser.add_argument('--output-dir', '-o', default='plots', 
                       help='Output directory for plots (default: plots)')
    parser.add_argument('--summary-only', '-s', action='store_true',
                       help='Only print summary statistics, skip visualizations')
    
    args = parser.parse_args()
    
    if not os.path.exists(args.csv_file):
        print(f"Error: CSV file '{args.csv_file}' not found.")
        return 1
    
    try:
        # Initialize analyzer
        analyzer = RolloutTraceAnalyzer(args.csv_file)
        
        # Print summary statistics
        analyzer.print_summary_stats()
        
        # Create visualizations unless summary-only is specified
        if not args.summary_only:
            analyzer.create_visualizations(args.output_dir)
            print(f"\nAnalysis complete! Check {args.output_dir}/ for visualizations.")
        
        return 0
        
    except Exception as e:
        print(f"Error during analysis: {e}")
        return 1


if __name__ == "__main__":
    exit(main())

