import noisereduce as nr
import soundfile as sf
import librosa
import numpy as np
import os

def normalize_audio_volume(audio, original_rms=None, target_dbfs=-20):
    """
    对音频进行音量归一化
    
    参数:
    audio: 音频数据
    original_rms: 原始音频的RMS值（可选）
    target_dbfs: 目标音量级别（dBFS）
    
    返回:
    归一化后的音频
    """
    # 计算当前音频的RMS
    current_rms = np.sqrt(np.mean(audio**2))
    
    # 方法1：恢复到原始音量水平
    if original_rms is not None and original_rms > 0:
        scale_factor = original_rms / current_rms
        normalized_audio = audio * scale_factor
        print(f"音量恢复: 缩放因子 {scale_factor:.2f}")
    
    # 方法2：基于目标dBFS的音量归一化
    else:
        # 计算当前音频的dBFS
        current_dbfs = 20 * np.log10(current_rms / (np.max(np.abs(audio)) + 1e-8))
        
        # 计算需要的增益
        gain_db = target_dbfs - current_dbfs
        gain_linear = 10 ** (gain_db / 20)
        
        normalized_audio = audio * gain_linear
        print(f"音量归一化: 增益 {gain_db:.1f} dB, 缩放因子 {gain_linear:.2f}")
    
    # 防止削波（clipping）
    max_val = np.max(np.abs(normalized_audio))
    if max_val > 0.95:  # 如果接近最大值1.0
        normalized_audio = normalized_audio * 0.95 / max_val
        print("已应用防削波处理")
    
    return normalized_audio
def reduce_background_noise(audio_path, output_path=None, use_default_noise=True, 
                           noise_path='assets/audio/noise_sample.wav', noise_start=0.1, noise_end=0.5,
                           normalize_volume=True, target_dbfs=-20):
    """
    使用noisereduce库去除背景噪声
    
    参数:
    audio_path: 音频文件路径
    output_path: 输出文件路径（可选）
    use_default_noise: 是否使用默认噪声样本
    noise_path: 默认噪声样本路径
    noise_start: 噪声样本开始时间（秒）（仅当use_default_noise=False时使用）
    noise_end: 噪声样本结束时间（秒）（仅当use_default_noise=False时使用）
    """
    
    # 加载音频
    y, sr = librosa.load(audio_path, sr=None)
    
    print(f"音频信息: 时长 {len(y)/sr:.2f}秒, 采样率 {sr}Hz")
    
    # 提取或加载噪声样本
    if use_default_noise:
        try:
            noise_path = os.path.join("assets", "audio", "noise_sample.wav")
            noise_clip, noise_sr = librosa.load(noise_path, sr=None)
            # 检查采样率是否匹配
            if noise_sr != sr:
                print(f"警告: 噪声样本采样率({noise_sr}Hz)与音频采样率({sr}Hz)不匹配")
                # 重采样噪声样本以匹配音频采样率
                noise_clip = librosa.resample(noise_clip, orig_sr=noise_sr, target_sr=sr)
                print("已对噪声样本进行重采样以匹配音频采样率")
        except FileNotFoundError as e:
            print(f"错误: {e}")
            print("将使用当前音频提取噪声样本")
            use_default_noise = False
    
    if not use_default_noise:
        # 从当前音频提取噪声样本
        noise_start_sample = int(noise_start * sr)
        noise_end_sample = int(noise_end * sr)
        noise_clip = y[noise_start_sample:noise_end_sample]
        print(f"使用当前音频的 {noise_end-noise_start:.2f}秒 噪声样本进行降噪")
    else:
        print(f"使用默认噪声样本进行降噪，噪声长度: {len(noise_clip)/sr:.2f}秒")
    
    # 应用降噪
    reduced_noise = nr.reduce_noise(
        y=y,
        sr=sr,
        y_noise=noise_clip,
        prop_decrease=0.95,  # 降噪比例
        n_fft=1024,
        win_length=1024,
        hop_length=256,
        n_std_thresh_stationary=1.5,
        stationary=True
    )
    
    # 记录原始音频的音量（RMS）
    original_rms = np.sqrt(np.mean(y**2))
    print(f"原始音频RMS: {original_rms:.4f}")
    # 音量归一化处理
    if normalize_volume:
        normalize_audio_volume(reduced_noise, original_rms, target_dbfs)
    
    # 保存结果
    if output_path:
        sf.write(output_path, reduced_noise, sr)
        print(f"降噪后的音频已保存至: {output_path}")

if __name__ == "__main__":
    reduce_background_noise('test/TalkingQ_XQSN00001004_eba844b5.mp3', 'test/denoised2.mp3')