plot_results.py

"""
Plot Results Script for RL Training Visualization (Task 7)
Generates publication-quality figures from training and testing logs

Outputs:
  1. reward_curve.png     - Episode reward over training
  2. success_rate.png    - Success rate over training  
  3. mode_usage_pie.png  - Mode distribution pie chart
  4. completion_time.png - Completion time histogram/bar chart

Usage:
    python plot_results.py --log-dir ./logs/ppo_map_a_xxx/
    python plot_results.py --test-results ./test_results/test_tri_mode_composite_xxx/
"""

import numpy as np
import csv
import os
import sys
import json
from datetime import datetime

_CACHE_ROOT = os.path.join(os.path.abspath(os.path.dirname(__file__)), '.cache')
os.makedirs(os.path.join(_CACHE_ROOT, 'matplotlib'), exist_ok=True)
os.environ.setdefault('XDG_CACHE_HOME', _CACHE_ROOT)
os.environ.setdefault('MPLCONFIGDIR', os.path.join(_CACHE_ROOT, 'matplotlib'))

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

sys.path.append(os.path.abspath(os.path.dirname(__file__)))

from map_manager import MapManager

# Set publication-quality plot settings
plt.rcParams.update({
    'figure.dpi': 150,
    'savefig.dpi': 300,
    'font.size': 11,
    'axes.titlesize': 14,
    'axes.labelsize': 12,
    'legend.fontsize': 10,
    'xtick.labelsize': 10,
    'ytick.labelsize': 10,
    'figure.figsize': (10, 6),
    'axes.grid': True,
    'grid.alpha': 0.3
})


def load_action_stats(log_dir):
    """
    Load action statistics CSV file (Task 4 output)
    
    Args:
        log_dir: Directory containing action_stats.csv
    
    Returns:
        list: List of episode stat dictionaries
    """
    csv_path = os.path.join(log_dir, 'action_stats.csv')
    
    if not os.path.exists(csv_path):
        print(f"Warning: {csv_path} not found")
        return []
    
    stats = []
    with open(csv_path, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            # Convert numeric fields
            row['episode'] = int(row['episode'])
            row['success'] = row['success'] == 'True' or row['success'] == 'True'
            row['total_reward'] = float(row['total_reward'])
            row['steps'] = int(row['steps'])
            row['time(s)'] = float(row['time(s)'])
            row['final_distance'] = float(row['final_distance'])
            row['afm_count'] = int(row['afm_count'])
            row['apt_count'] = int(row['apt_count'])
            row['azr_count'] = int(row['azr_count'])
            row['afm_pct'] = float(row['afm_pct'])
            row['apt_pct'] = float(row['apt_pct'])
            row['azr_pct'] = float(row['azr_pct'])
            row['mode_switches'] = int(row['mode_switches'])
            if 'failure_reason' in row:
                row['failure_reason'] = row['failure_reason']
            if 'apt_candidate_steps' in row:
                row['apt_candidate_steps'] = int(row['apt_candidate_steps'])
            if 'azr_candidate_steps' in row:
                row['azr_candidate_steps'] = int(row['azr_candidate_steps'])
            if 'blocked_steps' in row:
                row['blocked_steps'] = int(row['blocked_steps'])
            
            stats.append(row)
    
    return stats


def load_eval_results(log_dir):
    """
    Load evaluation results from EvalCallback (Task 2 output)
    
    Args:
        log_dir: Directory containing results.csv
    
    Returns:
        list: List of evaluation result dictionaries
    """
    csv_path = os.path.join(log_dir, 'results.csv')
    
    if not os.path.exists(csv_path):
        print(f"Warning: {csv_path} not found")
        return []
    
    results = []
    with open(csv_path, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            for key in row:
                try:
                    row[key] = float(row[key])
                except ValueError:
                    pass
            results.append(row)
    
    return results


def load_episode_details(test_dir):
    """Load episode_details.csv generated by test.py."""
    csv_path = os.path.join(test_dir, 'episode_details.csv')

    if not os.path.exists(csv_path):
        print(f"Warning: {csv_path} not found")
        return []

    rows = []
    with open(csv_path, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            row['episode'] = int(row['episode'])
            row['success'] = row['success'] == 'True'
            row['reward'] = float(row['reward'])
            row['steps'] = int(row['steps'])
            row['time_s'] = float(row['time_s'])
            row['final_distance_m'] = float(row['final_distance_m'])
            row['afm_count'] = int(row['afm_count'])
            row['apt_count'] = int(row['apt_count'])
            row['azr_count'] = int(row['azr_count'])
            row['mode_switches'] = int(row['mode_switches'])
            row['total_actions'] = int(row['total_actions'])
            if 'apt_candidate_steps' in row:
                row['apt_candidate_steps'] = int(row['apt_candidate_steps'])
            if 'azr_candidate_steps' in row:
                row['azr_candidate_steps'] = int(row['azr_candidate_steps'])
            if 'blocked_steps' in row:
                row['blocked_steps'] = int(row['blocked_steps'])
            rows.append(row)
    return rows


def load_step_trace(test_dir):
    """Load per-step traces generated by test.py."""
    csv_path = os.path.join(test_dir, 'step_trace.csv')

    if not os.path.exists(csv_path):
        print(f"Warning: {csv_path} not found")
        return []

    rows = []
    with open(csv_path, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            row['episode'] = int(row['episode'])
            row['step'] = int(row['step'])
            row['mode'] = int(row['mode'])
            row['reward'] = float(row['reward'])
            row['x'] = float(row['x'])
            row['y'] = float(row['y'])
            row['heading_rad'] = float(row['heading_rad'])
            row['distance_to_goal'] = float(row['distance_to_goal'])
            row['forward_clearance'] = float(row['forward_clearance'])
            row['left_clearance'] = float(row['left_clearance'])
            row['right_clearance'] = float(row['right_clearance'])
            row['path_progress'] = float(row['path_progress'])
            row['path_heading_change'] = float(row['path_heading_change'])
            row['is_apt_candidate'] = row['is_apt_candidate'] == 'True'
            row['is_azr_candidate'] = row['is_azr_candidate'] == 'True'
            row['suggested_mode'] = int(row['suggested_mode'])
            row['step_longitudinal'] = float(row['step_longitudinal'])
            row['step_lateral'] = float(row['step_lateral'])
            row['step_heading_change'] = float(row['step_heading_change'])
            rows.append(row)
    return rows


def load_test_summary(test_dir):
    """Load test_summary.json from a test-results directory."""
    json_path = os.path.join(test_dir, 'test_summary.json')
    if not os.path.exists(json_path):
        print(f"Warning: {json_path} not found")
        return None
    with open(json_path, 'r') as f:
        return json.load(f)


def plot_reward_curve(stats, output_path, title='Training Reward Curve'):
    """
    Plot 1: Episode reward over training (Task 7.1)
    
    Args:
        stats: List of episode statistics
        output_path: Directory to save figure
        title: Plot title
    """
    if len(stats) == 0:
        print("No data to plot reward curve")
        return
    
    fig, ax = plt.subplots(figsize=(12, 6))
    
    episodes = [s['episode'] for s in stats]
    rewards = [s['total_reward'] for s in stats]
    
    # Calculate rolling average for smoothing
    window = min(20, len(rewards))
    if window > 1:
        rewards_smooth = np.convolve(rewards, np.ones(window)/window, mode='valid')
        episodes_smooth = episodes[window-1:]
        
        ax.plot(episodes_smooth, rewards_smooth, 'b-', linewidth=2.5, 
               label=f'Rolling Mean (window={window})', alpha=0.8)
    
    # Plot raw data with low alpha
    ax.scatter(episodes, rewards, c=rewards, cmap='viridis', 
              s=15, alpha=0.4, label='Episode Reward')
    
    ax.set_xlabel('Episode Number', fontweight='bold')
    ax.set_ylabel('Total Reward', fontweight='bold')
    ax.set_title(title, fontweight='bold', pad=15)
    ax.legend(loc='best')
    
    # Add colorbar
    cbar = plt.colorbar(ax.collections[0] if len(ax.collections) > 0 else None, ax=ax)
    cbar.set_label('Reward Value', rotation=270, labelpad=15)
    
    # Statistics annotation
    mean_r = np.mean(rewards)
    std_r = np.std(rewards)
    max_r = np.max(rewards)
    min_r = np.min(rewards)
    
    textstr = f'Mean: {mean_r:.2f}\nStd: {std_r:.2f}\nMax: {max_r:.2f}\nMin: {min_r:.2f}'
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.8)
    ax.text(0.02, 0.98, textstr, transform=ax.transAxes, fontsize=10,
           verticalalignment='top', bbox=props)
    
    plt.tight_layout()
    
    save_path = os.path.join(output_path, 'reward_curve.png')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"✓ Saved: {save_path}")


def plot_success_rate(stats, output_path, title='Success Rate Over Training'):
    """
    Plot 2: Success rate curve (Task 7.2)
    
    Shows running success rate over training episodes
    """
    if len(stats) == 0:
        print("No data to plot success rate")
        return
    
    fig, ax = plt.subplots(figsize=(12, 6))
    
    episodes = [s['episode'] for s in stats]
    successes = [1 if s['success'] else 0 for s in stats]
    
    # Calculate cumulative success rate with window
    window = min(50, len(successes))
    success_rate = []
    for i in range(len(successes)):
        start_idx = max(0, i - window + 1)
        rate = np.mean(successes[start_idx:i+1]) * 100
        success_rate.append(rate)
    
    ax.fill_between(episodes, success_rate, alpha=0.3, color='green')
    ax.plot(episodes, success_rate, 'g-', linewidth=2, label='Success Rate (%)')
    
    # Mark final success rate
    final_rate = success_rate[-1]
    ax.axhline(y=final_rate, color='red', linestyle='--', linewidth=1.5, 
               alpha=0.7, label=f'Final: {final_rate:.1f}%')
    
    ax.set_xlabel('Episode Number', fontweight='bold')
    ax.set_ylabel('Success Rate (%)', fontweight='bold')
    ax.set_title(title, fontweight='bold', pad=15)
    ax.legend(loc='lower right')
    ax.set_ylim(0, 105)
    
    # Annotation
    total_successes = sum(successes)
    total_episodes = len(successes)
    textstr = f'Total Episodes: {total_episodes}\nSuccesses: {total_successes}\nFinal Rate: {final_rate:.1f}%'
    props = dict(boxstyle='round', facecolor='lightgreen', alpha=0.8)
    ax.text(0.02, 0.98, textstr, transform=ax.transAxes, fontsize=10,
           verticalalignment='top', bbox=props)
    
    plt.tight_layout()
    
    save_path = os.path.join(output_path, 'success_rate.png')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"✓ Saved: {save_path}")


def plot_mode_usage_pie(stats, output_path, title='Mode Usage Distribution'):
    """
    Plot 3: Mode usage pie chart (Task 7.3)
    
    Shows AFM/APT/AZR usage percentages across all episodes
    """
    if len(stats) == 0:
        print("No data to plot mode usage")
        return
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
    
    # Aggregate mode counts across all episodes
    total_afm = sum(s.get('afm_count', 0) for s in stats)
    total_apt = sum(s.get('apt_count', 0) for s in stats)
    total_azr = sum(s.get('azr_count', 0) for s in stats)
    total = total_afm + total_apt + total_azr or 1
    
    # Pie chart data
    sizes = [total_afm, total_apt, total_azr]
    labels = ['AFM\n(Path Tracking)', 'APT\n(Translation)', 'AZR\n(Rotation)']
    colors = ['#3498db', '#e74c3c', '#2ecc71']
    explode = (0.05, 0.05, 0.05)  # Slightly separate all slices
    
    def make_autopct(values):
        """Custom autopct function to show both percentage and count"""
        def my_autopct(pct):
            absolute = int(round(pct/100.*sum(values)))
            return f'{pct:.1f}%\n({absolute:,})'
        return my_autopct
    
    # Main pie chart
    wedges, texts, autotexts = ax1.pie(
        sizes, explode=explode, labels=labels, colors=colors,
        autopct=make_autopct(sizes), shadow=True, startangle=90,
        textprops={'fontsize': 11, 'fontweight': 'bold'}
    )
    
    ax1.set_title('Overall Mode Usage\n(All Episodes)', fontweight='bold', pad=20)
    
    # Equal aspect ratio ensures that pie is drawn as a circle
    ax1.axis('equal')
    
    # Bar chart showing mode usage per episode
    episodes = [s['episode'] for s in stats]
    afm_pcts = [s.get('afm_pct', 0) for s in stats]
    apt_pcts = [s.get('apt_pct', 0) for s in stats]
    azr_pcts = [s.get('azr_pct', 0) for s in stats]
    
    x = np.arange(len(episodes))
    width = 0.25
    
    bars1 = ax2.bar(x - width, afm_pcts, width, label='AFM', color='#3498db', alpha=0.8)
    bars2 = ax2.bar(x, apt_pcts, width, label='APT', color='#e74c3c', alpha=0.8)
    bars3 = ax2.bar(x + width, azr_pcts, width, label='AZR', color='#2ecc71', alpha=0.8)
    
    ax2.set_xlabel('Episode Number', fontweight='bold')
    ax2.set_ylabel('Usage Percentage (%)', fontweight='bold')
    ax2.set_title('Mode Usage Per Episode', fontweight='bold', pad=15)
    ax2.legend(loc='upper right')
    ax2.set_xticks(x[::max(1, len(episodes)//10)])  # Show subset of ticks
    
    plt.tight_layout()
    
    save_path = os.path.join(output_path, 'mode_usage_pie.png')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"✓ Saved: {save_path}")
    
    # Print summary
    print(f"\n📊 Mode Usage Summary:")
    print(f"   AFM (Path Track): {100*total_afm/total:5.1f}% ({total_afm:,} actions)")
    print(f"   APT (Translate):  {100*total_apt/total:5.1f}% ({total_apt:,} actions)")
    print(f"   AZR (Rotate):     {100*total_azr/total:5.1f}% ({total_azr:,} actions)")


def plot_completion_time(stats, output_path, title='Completion Time Analysis'):
    """
    Plot 4: Completion time histogram/bar chart (Task 7.4)
    
    Shows distribution of completion times for successful vs failed episodes
    """
    if len(stats) == 0:
        print("No data to plot completion time")
        return
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
    
    # Separate successful and failed episodes
    success_times = [s['time(s)'] for s in stats if s['success']]
    fail_times = [s['time(s)'] for s in stats if not s['success']]
    
    # Histogram for all episodes
    all_times = [s['time(s)'] for s in stats]
    
    bins = np.linspace(0, max(all_times) + 5, 20)
    
    ax1.hist(success_times, bins=bins, alpha=0.7, color='green', 
            label=f'Success (n={len(success_times)})', edgecolor='black')
    ax1.hist(fail_times, bins=bins, alpha=0.7, color='red',
            label=f'Failed (n={len(fail_times)})', edgecolor='black')
    
    ax1.set_xlabel('Completion Time (seconds)', fontweight='bold')
    ax1.set_ylabel('Frequency', fontweight='bold')
    ax1.set_title('Distribution of Episode Durations', fontweight='bold', pad=15)
    ax1.legend(loc='upper right')
    
    # Statistics annotation
    if len(success_times) > 0:
        textstr = (f'Successful Episodes:\n'
                  f'  Count: {len(success_times)}\n'
                  f'  Mean Time: {np.mean(success_times):.2f}s\n'
                  f'  Std: {np.std(success_times):.2f}s\n'
                  f'  Min: {min(success_times):.2f}s\n'
                  f'  Max: {max(success_times):.2f}s')
        props = dict(boxstyle='round', facecolor='lightgreen', alpha=0.8)
        ax1.text(0.98, 0.98, textstr, transform=ax1.transAxes, fontsize=9,
                verticalalignment='top', horizontalalignment='right', bbox=props)
    
    # Box plot comparing success vs failure
    data_to_plot = []
    labels = []
    
    if len(success_times) > 0:
        data_to_plot.append(success_times)
        labels.append(f'Success\n(n={len(success_times)})')
    
    if len(fail_times) > 0:
        data_to_plot.append(fail_times)
        labels.append(f'Failed\n(n={len(fail_times)})')
    
    if len(data_to_plot) > 0:
        bp = ax2.boxplot(data_to_plot, tick_labels=labels, patch_artist=True)
        colors = ['lightgreen', 'lightcoral']
        for patch, color in zip(bp['boxes'], colors):
            patch.set_facecolor(color)
            patch.set_alpha(0.7)
    
    ax2.set_ylabel('Time (seconds)', fontweight='bold')
    ax2.set_title('Success vs Failure: Time Comparison', fontweight='bold', pad=15)
    ax2.grid(True, axis='y', alpha=0.3)
    
    plt.tight_layout()
    
    save_path = os.path.join(output_path, 'completion_time.png')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"✓ Saved: {save_path}")


def plot_failure_reasons(episode_details, output_path, title='Failure Reason Breakdown'):
    """Plot failure reason counts from test episodes."""
    if len(episode_details) == 0:
        print("No data to plot failure reasons")
        return

    reasons = {}
    for row in episode_details:
        reason = row.get('failure_reason', 'unknown')
        reasons[reason] = reasons.get(reason, 0) + 1

    labels = list(reasons.keys())
    values = [reasons[label] for label in labels]

    fig, ax = plt.subplots(figsize=(10, 5))
    bars = ax.bar(labels, values, color=['#27ae60' if label == 'success' else '#e74c3c' for label in labels])
    ax.set_xlabel('Reason', fontweight='bold')
    ax.set_ylabel('Episode Count', fontweight='bold')
    ax.set_title(title, fontweight='bold', pad=15)
    ax.set_ylim(0, max(values) + 1)

    for bar, value in zip(bars, values):
        ax.text(bar.get_x() + bar.get_width() / 2, value + 0.05, str(value),
                ha='center', va='bottom', fontsize=10, fontweight='bold')

    plt.tight_layout()
    save_path = os.path.join(output_path, 'failure_reasons.png')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"✓ Saved: {save_path}")


def plot_test_trajectories(step_trace, episode_details, summary, output_path,
                           title='Episode Trajectories'):
    """Overlay all test trajectories on the map."""
    if len(step_trace) == 0:
        print("No data to plot trajectories")
        return

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

    map_type = 'tri_mode_composite'
    if summary is not None:
        map_type = summary.get('test_configuration', {}).get('map_type', 'tri_mode_composite')

    try:
        env_map = MapManager().create_map(map_type)
        env_map.draw_track(ax)
    except Exception as exc:
        print(f"Warning: failed to draw map background: {exc}")

    success_by_episode = {row['episode']: row['success'] for row in episode_details}
    grouped = {}
    for row in step_trace:
        grouped.setdefault(row['episode'], []).append(row)

    for episode, rows in grouped.items():
        xs = [row['x'] for row in rows]
        ys = [row['y'] for row in rows]
        success = success_by_episode.get(episode, False)
        color = '#27ae60' if success else '#c0392b'
        alpha = 0.85 if success else 0.25
        linewidth = 2.0 if success else 1.1
        ax.plot(xs, ys, color=color, alpha=alpha, linewidth=linewidth)

    ax.set_title(title, fontweight='bold', pad=15)
    ax.set_xlabel('X (m)', fontweight='bold')
    ax.set_ylabel('Y (m)', fontweight='bold')
    ax.grid(True, alpha=0.3)

    save_path = os.path.join(output_path, 'test_trajectories.png')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"✓ Saved: {save_path}")


def plot_mode_timeline(step_trace, output_path, title='Mode Timeline Heatmap', max_episodes=20):
    """Visualize per-step mode choices as a heatmap."""
    if len(step_trace) == 0:
        print("No data to plot mode timeline")
        return

    episodes = sorted(set(row['episode'] for row in step_trace))[:max_episodes]
    max_steps = max(row['step'] for row in step_trace if row['episode'] in episodes)

    matrix = np.full((len(episodes), max_steps), np.nan)
    episode_to_idx = {ep: idx for idx, ep in enumerate(episodes)}
    for row in step_trace:
        if row['episode'] not in episode_to_idx:
            continue
        matrix[episode_to_idx[row['episode']], row['step'] - 1] = row['mode']

    masked = np.ma.masked_invalid(matrix)
    cmap = ListedColormap(['#3498db', '#e74c3c', '#2ecc71'])

    fig, ax = plt.subplots(figsize=(12, 6))
    im = ax.imshow(masked, aspect='auto', interpolation='nearest', cmap=cmap, vmin=0, vmax=2)
    cbar = plt.colorbar(im, ax=ax, ticks=[0, 1, 2])
    cbar.ax.set_yticklabels(['AFM', 'APT', 'AZR'])

    ax.set_xlabel('Step', fontweight='bold')
    ax.set_ylabel('Episode', fontweight='bold')
    ax.set_yticks(np.arange(len(episodes)))
    ax.set_yticklabels([str(ep) for ep in episodes])
    ax.set_title(title, fontweight='bold', pad=15)

    plt.tight_layout()
    save_path = os.path.join(output_path, 'mode_timeline.png')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"✓ Saved: {save_path}")


def generate_test_plots(test_dir, output_dir='./figures/'):
    """Generate figures from a test-results directory."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_path = os.path.join(output_dir, f'test_plots_{timestamp}')
    os.makedirs(output_path, exist_ok=True)

    print("\n" + "="*80)
    print(" GENERATING TEST RESULT PLOTS")
    print("="*80)
    print(f"\nInput directory:  {test_dir}")
    print(f"Output directory: {output_path}\n")

    summary = load_test_summary(test_dir)
    episode_details = load_episode_details(test_dir)
    action_stats = load_action_stats(test_dir)
    step_trace = load_step_trace(test_dir)

    if len(episode_details) == 0:
        print("Error: No test episode data found!")
        return None

    stats_for_plots = action_stats if len(action_stats) > 0 else [
        {
            'episode': row['episode'],
            'success': row['success'],
            'total_reward': row['reward'],
            'time(s)': row['time_s'],
            'afm_count': row['afm_count'],
            'apt_count': row['apt_count'],
            'azr_count': row['azr_count'],
            'afm_pct': 100 * row['afm_count'] / max(1, row['total_actions']),
            'apt_pct': 100 * row['apt_count'] / max(1, row['total_actions']),
            'azr_pct': 100 * row['azr_count'] / max(1, row['total_actions']),
            'mode_switches': row['mode_switches']
        }
        for row in episode_details
    ]

    plot_mode_usage_pie(stats_for_plots, output_path, title='Test Mode Usage Distribution')
    plot_completion_time(stats_for_plots, output_path, title='Test Completion Time Analysis')
    plot_failure_reasons(episode_details, output_path)
    plot_test_trajectories(step_trace, episode_details, summary, output_path)
    plot_mode_timeline(step_trace, output_path)

    summary_file = os.path.join(output_path, 'plot_summary.txt')
    with open(summary_file, 'w') as f:
        f.write("="*80 + "\n")
        f.write("TEST PLOT GENERATION SUMMARY\n")
        f.write("="*80 + "\n\n")
        f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Source: {test_dir}\n")
        f.write(f"Output: {output_path}\n\n")
        f.write("Generated Figures:\n")
        f.write("  1. mode_usage_pie.png   - Mode usage distribution\n")
        f.write("  2. completion_time.png  - Episode duration analysis\n")
        f.write("  3. failure_reasons.png  - Failure reason counts\n")
        f.write("  4. test_trajectories.png - All episode trajectories\n")
        f.write("  5. mode_timeline.png    - Per-step mode heatmap\n")

    print(f"\n{'-'*80}")
    print(" TEST PLOT GENERATION COMPLETED!")
    print(f"{'-'*80}")
    print(f"\nGenerated files:")
    print(f"  -> {output_path}/mode_usage_pie.png")
    print(f"  -> {output_path}/completion_time.png")
    print(f"  -> {output_path}/failure_reasons.png")
    print(f"  -> {output_path}/test_trajectories.png")
    print(f"  -> {output_path}/mode_timeline.png")
    print(f"  -> {output_path}/plot_summary.txt\n")

    return output_path


def generate_all_plots(log_dir, output_dir='./figures/'):
    """
    Generate all visualization plots (Task 7 main function)
    
    Args:
        log_dir: Directory containing training logs
        output_dir: Directory to save generated figures
    """
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_path = os.path.join(output_dir, f'plots_{timestamp}')
    os.makedirs(output_path, exist_ok=True)
    
    print("\n" + "="*80)
    print(" GENERATING VISUALIZATION PLOTS")
    print("="*80)
    print(f"\nInput directory:  {log_dir}")
    print(f"Output directory: {output_path}\n")
    
    # Load data
    stats = load_action_stats(log_dir)
    
    if len(stats) == 0:
        print("Error: No training data found!")
        return None
    
    print(f"Loaded {len(stats)} episodes from action_stats.csv\n")
    
    # Generate plots
    print("Generating plots...")
    
    plot_reward_curve(stats, output_path)
    plot_success_rate(stats, output_path)
    plot_mode_usage_pie(stats, output_path)
    plot_completion_time(stats, output_path)
    
    # Generate summary report
    summary_file = os.path.join(output_path, 'plot_summary.txt')
    with open(summary_file, 'w') as f:
        f.write("="*80 + "\n")
        f.write("PLOT GENERATION SUMMARY\n")
        f.write("="*80 + "\n\n")
        f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Source: {log_dir}\n")
        f.write(f"Output: {output_path}\n\n")
        f.write(f"Total Episodes Analyzed: {len(stats)}\n\n")
        f.write("Generated Figures:\n")
        f.write("  1. reward_curve.png      - Episode reward over training\n")
        f.write("  2. success_rate.png     - Success rate progression\n")
        f.write("  3. mode_usage_pie.png   - Mode distribution (pie + bar)\n")
        f.write("  4. completion_time.png - Duration analysis (histogram + boxplot)\n")
    
    print(f"\n{'─'*80}")
    print(f" PLOT GENERATION COMPLETED!")
    print(f"{'─'*80}")
    print(f"\nGenerated files:")
    print(f"  → {output_path}/reward_curve.png")
    print(f"  → {output_path}/success_rate.png")
    print(f"  → {output_path}/mode_usage_pie.png")
    print(f"  → {output_path}/completion_time.png")
    print(f"  → {output_path}/plot_summary.txt\n")
    
    return output_path


def main():
    """Command-line interface"""
    import argparse
    
    parser = argparse.ArgumentParser(
        description='Generate publication-quality plots from training logs',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Generate all plots from training logs
  python plot_results.py --log-dir ./logs/ppo_map_a_xxx/

  # Custom output directory
  python plot_results.py --log-dir ./logs/xxx/ --output-dir ./paper_figures/

  # From test results
  python plot_results.py --test-results ./test_results/test_xxx/
        """
    )
    
    parser.add_argument('--log-dir', type=str, default=None,
                       help='Directory containing training logs (with action_stats.csv)')
    parser.add_argument('--test-results', type=str, default=None,
                       help='Directory containing test results (with test_summary.json)')
    parser.add_argument('--output-dir', type=str, default='./figures/',
                       help='Output directory for generated figures')
    
    args = parser.parse_args()
    
    if args.log_dir:
        generate_all_plots(args.log_dir, args.output_dir)
    elif args.test_results:
        generate_test_plots(args.test_results, args.output_dir)
    else:
        parser.print_help()
        print("\nExample usage:")
        print("  python plot_results.py --log-dir ./logs/ppo_map_a_20260414_xxx/")


if __name__ == "__main__":
    main()