"""
批量处理音频文件生成波形图的脚本
"""

import os
import matplotlib.pyplot as plt
import librosa
import librosa.display
import argparse
import glob

def generate_waveform(audio_path, output_path):
    """
    生成音频波形图并保存
    
    Args:
        audio_path: 音频文件路径
        output_path: 输出图片保存路径
    
    Returns:
        None
    """
    # 创建输出目录
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    # 加载音频
    y, sr = librosa.load(audio_path, sr=None)
    
    # 创建图形
    plt.figure(figsize=(12, 4))
    
    # 绘制波形
    librosa.display.waveshow(y, sr=sr)
    plt.title('Waveform')
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    
    # 保存图片
    plt.tight_layout()
    plt.savefig(output_path)
    plt.close()

def process_audio_directory(audio_dir, output_dir, limit=None, file_ext='wav'):
    """
    批量处理音频文件生成波形图
    
    Args:
        audio_dir: 音频文件目录
        output_dir: 输出图片保存目录
        limit: 处理文件数量限制
        file_ext: 音频文件扩展名
    
    Returns:
        int: 处理的文件数量
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 获取所有音频文件
    audio_files = glob.glob(os.path.join(audio_dir, f"**/*.{file_ext}"), recursive=True)
    
    if limit and len(audio_files) > limit:
        audio_files = audio_files[:limit]
    
    processed_count = 0
    for audio_file in audio_files:
        # 生成输出文件路径
        rel_path = os.path.relpath(audio_file, audio_dir)
        output_path = os.path.join(output_dir, f"{os.path.splitext(rel_path)[0]}.png")
        
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        try:
            # 生成波形图
            generate_waveform(audio_file, output_path)
            processed_count += 1
            print(f"Processed {processed_count}/{len(audio_files)}: {rel_path}")
        except Exception as e:
            print(f"Error processing {audio_file}: {e}")
    
    return processed_count

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='Batch Audio Waveform Generator')
    parser.add_argument('--audio_dir', type=str, default='./RAVDESS/Actor_01',
                       help='Directory containing audio files')
    parser.add_argument('--output_dir', type=str, default='./output/waveforms',
                       help='Directory to save waveform images')
    parser.add_argument('--limit', type=int, default=None,
                       help='Limit the number of files to process')
    parser.add_argument('--file_ext', type=str, default='wav',
                       help='Audio file extension to process')
    
    args = parser.parse_args()
    
    # 检查目录是否存在
    if not os.path.exists(args.audio_dir):
        print(f"Error: Audio directory not found: {args.audio_dir}")
        return
    
    # 批量处理音频文件
    start_time = __import__('time').time()
    count = process_audio_directory(args.audio_dir, args.output_dir, args.limit, args.file_ext)
    end_time = __import__('time').time()
    
    print(f"\nProcessed {count} files in {end_time - start_time:.2f} seconds")
    print(f"Waveform images saved to: {args.output_dir}")

if __name__ == "__main__":
    main()