import os
import pandas as pd
import matplotlib.pyplot as plt


def plot_success_rates(logs_dir):
    """
    Plot success rates from all CSV files found recursively in the logs directory.
    
    Args:
        logs_dir: Directory containing subdirectories with CSV files
    """
    success_rates = {}
    
    # Recursively find all CSV files
    for root, dirs, files in os.walk(logs_dir):
        for file in files:
            if file.endswith('.csv'):
                csv_path = os.path.join(root, file)
                
                # Get relative path from logs_dir
                rel_path = os.path.relpath(csv_path, logs_dir)
                
                try:
                    # Read the CSV file
                    df = pd.read_csv(csv_path, sep="\t")
                    
                    # Calculate success rate (assuming max_reward > 0 means success)
                    success_rate = (df["max_reward"] > 0.03).mean() * 100
                    
                    success_rates[rel_path] = success_rate
                    print(f"{rel_path}: {success_rate:.2f}%")
                    
                except Exception as e:
                    print(f"Error processing {rel_path}: {e}")
                    continue
    
    if not success_rates:
        print("No valid CSV files found!")
        return
    
    # Sort by relative path
    sorted_items = sorted(success_rates.items())
    
    folders, rates = zip(*sorted_items)
    
    # Create bar plot
    plt.figure(figsize=(12, 6))
    bars = plt.bar(range(len(folders)), rates, color='steelblue', edgecolor='black')
    
    # Add value labels on top of bars
    for i, bar in enumerate(bars):
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height,
                f'{height:.1f}%',
                ha='center', va='bottom', fontsize=8)
    
    plt.xticks(range(len(folders)), folders, rotation=45, ha='right', fontsize=8)
    plt.xlabel('Log File', fontsize=12)
    plt.ylabel('Success Rate (%)', fontsize=12)
    plt.title('Success Rate by Log File', fontsize=14, fontweight='bold')
    plt.ylim(0, 100)
    plt.grid(axis='y', alpha=0.3)
    plt.tight_layout()
    
    # Save the plot
    output_path = os.path.join(logs_dir, 'success_rates.png')
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    print(f"\nPlot saved to: {output_path}")
    
    plt.show()


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1:
        logs_directory = sys.argv[1]
    else:
        # Default to logs folder in the project
        logs_directory = "logs"
    
    if not os.path.exists(logs_directory):
        print(f"Error: Directory '{logs_directory}' not found!")
        sys.exit(1)
    
    plot_success_rates(logs_directory)
