"""
简单的音频波形可视化脚本
生成波形图
"""

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

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()
    
    print(f"Waveform saved to: {output_path}")

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='Audio Waveform Generator')
    parser.add_argument('--audio_path', type=str, default='./RAVDESS/Actor_01/03-01-01-01-01-01-01.wav',
                       help='Path to audio file')
    parser.add_argument('--output_path', type=str, default='./output/waveform.png',
                       help='Path to save waveform image')
    
    args = parser.parse_args()
    
    # 检查音频文件是否存在
    if not os.path.exists(args.audio_path):
        print(f"Error: Audio file not found: {args.audio_path}")
        return
    
    # 生成波形图
    generate_waveform(args.audio_path, args.output_path)

if __name__ == "__main__":
    main()