class AudioDropWidget(QWidget):
    def dragEnterEvent(self, event):
        if event.mimeData().hasUrls():
            event.acceptProposedAction()
    
    def dropEvent(self, event):
        urls = event.mimeData().urls()
        # 支持多文件和文件夾拖放

def enhanced_reduce_noise(y, sr, strength=0.5, stationary=True):
    # 第一級(jí)：譜減法降噪
    y_denoised = nr.reduce_noise(
        y=y,
        sr=sr,
        stationary=stationary,
        prop_decrease=min(0.95, strength),
        n_fft=2048
    )
    
    # 第二級(jí)：Butterworth高通濾波
    sos = signal.butter(10, 0.05, 'highpass', fs=sr, output='sos')
    return signal.sosfilt(sos, y_denoised)

class BatchProcessor(QThread):
    progress_updated = pyqtSignal(int, str)
    
    def run(self):
        for i, filepath in enumerate(self.files):
            # ...處理邏輯...
            self.progress_updated.emit(progress, f"處理中: {filename}")

import os
import numpy as np
import matplotlib.pyplot as plt
import librosa
import librosa.display
import noisereduce as nr
import soundfile as sf
from scipy import signal
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                             QLabel, QPushButton, QLineEdit, QFileDialog, QMessageBox, 
                             QProgressBar, QStatusBar, QSlider, QComboBox, QCheckBox)
from PyQt5.QtCore import Qt, QMimeData, QUrl, QThread, pyqtSignal
from PyQt5.QtGui import QFont, QDragEnterEvent, QDropEvent

class AudioDropWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setAcceptDrops(True)
        self.setFixedHeight(100)
        
    def dragEnterEvent(self, event: QDragEnterEvent):
        if event.mimeData().hasUrls():
            event.acceptProposedAction()
    
    def dropEvent(self, event: QDropEvent):
        urls = event.mimeData().urls()
        if urls:
            main_window = self.window()
            if hasattr(main_window, 'handle_dropped_files'):
                main_window.handle_dropped_files(urls)
            else:
                QMessageBox.warning(self, "?? 錯(cuò)誤", "無法處理拖放的文件")

class BatchProcessor(QThread):
    progress_updated = pyqtSignal(int, str)
    finished = pyqtSignal()
    
    def __init__(self, files, output_dir, strength, stationary, voice_gain, parent=None):
        super().__init__(parent)
        self.files = files
        self.output_dir = output_dir
        self.strength = strength
        self.stationary = stationary
        self.voice_gain = voice_gain
        self.canceled = False
    
    def run(self):
        total_files = len(self.files)
        for i, filepath in enumerate(self.files):
            if self.canceled:
                break
                
            try:
                self.progress_updated.emit(int((i+1)/total_files*100), f"正在處理: {os.path.basename(filepath)}")
                
                # 加載音頻
                y, sr = librosa.load(filepath, sr=None)
                
                # 降噪處理
                processed_audio = self.enhanced_reduce_noise(
                    y=y,
                    sr=sr,
                    strength=self.strength,
                    stationary=self.stationary
                )
                
                # 人聲音量增強(qiáng)
                if self.voice_gain > 0:
                    processed_audio = self.enhance_voice_volume(
                        y=processed_audio,
                        gain=self.voice_gain
                    )
                
                # 音量歸一化
                processed_audio = librosa.util.normalize(processed_audio)
                
                # 保存文件
                output_path = os.path.join(self.output_dir, f"processed_{os.path.basename(filepath)}")
                sf.write(output_path, processed_audio, sr)
                
            except Exception as e:
                self.progress_updated.emit(int((i+1)/total_files*100), f"處理失敗: {os.path.basename(filepath)} ({str(e)})")
                continue
                
        self.finished.emit()
    
    def enhanced_reduce_noise(self, y, sr, strength=0.5, stationary=True):
        """增強(qiáng)的降噪算法"""
        try:
            # 第一級(jí)：譜減法降噪
            y_denoised = nr.reduce_noise(
                y=y,
                sr=sr,
                stationary=stationary,
                prop_decrease=min(0.95, strength),
                n_fft=2048,
                win_length=2048,
                hop_length=512
            )
            
            # 第二級(jí)：Butterworth濾波器去除高頻噪聲
            sos = signal.butter(10, 0.05, 'highpass', fs=sr, output='sos')
            y_filtered = signal.sosfilt(sos, y_denoised)
            
            return y_filtered
        except Exception as e:
            print(f"降噪過程中出錯(cuò): {str(e)}")
            return y
    
    def enhance_voice_volume(self, y, gain=0.0):
        """簡單的人聲音量增強(qiáng) (0.0-1.0對(duì)應(yīng)0%-100%增益)"""
        try:
            if gain == 0.0:  # 0%增益時(shí)不處理
                return y
            # 增益公式：output = input * (1 + gain)
            return y * (1.0 + gain)
        except Exception as e:
            print(f"人聲增強(qiáng)過程中出錯(cuò): {str(e)}")
            return y

class AudioDenoiserUI(QMainWindow):
    def __init__(self):
        super().__init__()
        
        # 初始化UI
        self.setWindowTitle("?? Python音頻降噪軟件")
        self.setGeometry(100, 100, 600, 600)
        
        # 主窗口部件
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        self.main_layout = QVBoxLayout(self.central_widget)
        
        # 標(biāo)題
        self.title_label = QLabel("?? 音頻降噪與人聲增強(qiáng)工具")
        self.title_label.setFont(QFont('Microsoft YaHei', 16, QFont.Bold))
        self.title_label.setAlignment(Qt.AlignCenter)
        self.main_layout.addWidget(self.title_label)
        
        # 拖放區(qū)域
        self.drop_area = AudioDropWidget(self)
        self.drop_layout = QVBoxLayout(self.drop_area)
        self.drop_label = QLabel("?? 拖放音頻文件或文件夾到這里\n或點(diǎn)擊下方按鈕選擇文件")
        self.drop_label.setFont(QFont('Microsoft YaHei'))
        self.drop_label.setAlignment(Qt.AlignCenter)
        self.drop_label.setStyleSheet("border: 2px dashed #aaa; padding: 20px;")
        self.drop_layout.addWidget(self.drop_label)
        self.main_layout.addWidget(self.drop_area)
        
        # 文件選擇按鈕
        self.file_button_layout = QHBoxLayout()
        self.load_button = QPushButton("?? 選擇音頻文件")
        self.load_button.setFont(QFont('Microsoft YaHei'))
        self.load_button.clicked.connect(self.open_file_dialog)
        self.file_button_layout.addWidget(self.load_button)
        
        self.load_folder_button = QPushButton("?? 選擇文件夾")
        self.load_folder_button.setFont(QFont('Microsoft YaHei'))
        self.load_folder_button.clicked.connect(self.open_folder_dialog)
        self.file_button_layout.addWidget(self.load_folder_button)
        
        self.main_layout.addLayout(self.file_button_layout)
        
        # 文件信息
        self.file_label = QLabel("? 未選擇文件")
        self.file_label.setFont(QFont('Microsoft YaHei'))
        self.file_label.setAlignment(Qt.AlignCenter)
        self.main_layout.addWidget(self.file_label)
        
        # 批量處理選項(xiàng)
        self.batch_checkbox = QCheckBox("啟用批量處理模式")
        self.batch_checkbox.setFont(QFont('Microsoft YaHei'))
        self.batch_checkbox.stateChanged.connect(self.toggle_batch_mode)
        self.main_layout.addWidget(self.batch_checkbox)
        
        # 批量文件列表
        self.batch_list_label = QLabel("待處理文件列表 (0):")
        self.batch_list_label.setFont(QFont('Microsoft YaHei'))
        self.batch_list_label.setVisible(False)
        self.main_layout.addWidget(self.batch_list_label)
        
        self.batch_files_text = QLabel("")
        self.batch_files_text.setFont(QFont('Microsoft YaHei', 8))
        self.batch_files_text.setWordWrap(True)
        self.batch_files_text.setVisible(False)
        self.main_layout.addWidget(self.batch_files_text)
        
        # 降噪設(shè)置
        self.denoise_layout = QVBoxLayout()
        
        # 降噪強(qiáng)度滑塊
        self.strength_layout = QHBoxLayout()
        self.strength_label = QLabel("?? 降噪強(qiáng)度 (0.1-1.0):")
        self.strength_label.setFont(QFont('Microsoft YaHei'))
        self.strength_slider = QSlider(Qt.Horizontal)
        self.strength_slider.setRange(1, 10)
        self.strength_slider.setValue(5)
        self.strength_value_label = QLabel("5")
        self.strength_value_label.setFont(QFont('Microsoft YaHei'))
        self.strength_slider.valueChanged.connect(self.update_strength_value)
        self.strength_layout.addWidget(self.strength_label)
        self.strength_layout.addWidget(self.strength_slider)
        self.strength_layout.addWidget(self.strength_value_label)
        self.denoise_layout.addLayout(self.strength_layout)
        
        # 降噪模式
        self.mode_layout = QHBoxLayout()
        self.mode_label = QLabel("降噪模式:")
        self.mode_label.setFont(QFont('Microsoft YaHei'))
        self.mode_combo = QComboBox()
        self.mode_combo.setFont(QFont('Microsoft YaHei'))
        self.mode_combo.addItems(["穩(wěn)態(tài)噪聲", "非穩(wěn)態(tài)噪聲"])
        self.mode_layout.addWidget(self.mode_label)
        self.mode_layout.addWidget(self.mode_combo)
        self.denoise_layout.addLayout(self.mode_layout)
        
        self.main_layout.addLayout(self.denoise_layout)
        
        # 人聲音量增益滑塊
        self.voice_layout = QHBoxLayout()
        self.voice_label = QLabel("?? 人聲音量增益 (0-100%):")
        self.voice_label.setFont(QFont('Microsoft YaHei'))
        self.voice_slider = QSlider(Qt.Horizontal)
        self.voice_slider.setRange(0, 100)  # 0%到100%
        self.voice_slider.setValue(0)       # 默認(rèn)0%(無增益)
        self.voice_value_label = QLabel("0%")
        self.voice_value_label.setFont(QFont('Microsoft YaHei'))
        self.voice_slider.valueChanged.connect(self.update_voice_value)
        self.voice_layout.addWidget(self.voice_label)
        self.voice_layout.addWidget(self.voice_slider)
        self.voice_layout.addWidget(self.voice_value_label)
        self.main_layout.addLayout(self.voice_layout)
        
        # 處理按鈕
        self.process_button = QPushButton("? 處理音頻")
        self.process_button.setFont(QFont('Microsoft YaHei', 12))
        self.process_button.clicked.connect(self.process_audio)
        self.main_layout.addWidget(self.process_button)
        
        # 保存按鈕
        self.save_button = QPushButton("?? 保存處理結(jié)果")
        self.save_button.setFont(QFont('Microsoft YaHei', 12))
        self.save_button.clicked.connect(self.save_audio)
        self.save_button.setEnabled(False)
        self.main_layout.addWidget(self.save_button)
        
        # 批量處理按鈕
        self.batch_process_button = QPushButton("? 批量處理所有文件")
        self.batch_process_button.setFont(QFont('Microsoft YaHei', 12))
        self.batch_process_button.clicked.connect(self.process_batch)
        self.batch_process_button.setVisible(False)
        self.main_layout.addWidget(self.batch_process_button)
        
        # 進(jìn)度條
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        self.main_layout.addWidget(self.progress_bar)
        
        # 狀態(tài)欄
        self.status_bar = QStatusBar()
        self.status_bar.setFont(QFont('Microsoft YaHei'))
        self.setStatusBar(self.status_bar)
        self.status_bar.showMessage("?? 準(zhǔn)備就緒")
        
        # 音頻數(shù)據(jù)
        self.audio_data = None
        self.sample_rate = None
        self.processed_audio = None
        self.filename = None
        self.batch_mode = False
        self.batch_files = []
        
        # 樣式
        self.setStyleSheet("""
            QPushButton {
                padding: 8px;
                font-size: 14px;
            }
            QLabel {
                font-size: 14px;
            }
            QLineEdit {
                font-size: 14px;
            }
            QComboBox {
                font-size: 14px;
            }
        """)
    
    def update_strength_value(self, value):
        self.strength_value_label.setText(str(value))
    
    def update_voice_value(self, value):
        self.voice_value_label.setText(f"{value}%")
    
    def toggle_batch_mode(self, state):
        self.batch_mode = state == Qt.Checked
        self.batch_list_label.setVisible(self.batch_mode)
        self.batch_files_text.setVisible(self.batch_mode)
        self.batch_process_button.setVisible(self.batch_mode)
        
        if not self.batch_mode:
            self.batch_files = []
            self.update_batch_list()
    
    def open_file_dialog(self):
        if self.batch_mode:
            filepaths, _ = QFileDialog.getOpenFileNames(
                self, "選擇音頻文件", "", 
                "音頻文件 (*.wav *.mp3 *.ogg *.flac);;所有文件 (*.*)"
            )
            if filepaths:
                self.add_batch_files(filepaths)
        else:
            filepath, _ = QFileDialog.getOpenFileName(
                self, "選擇音頻文件", "", 
                "音頻文件 (*.wav *.mp3 *.ogg *.flac);;所有文件 (*.*)"
            )
            if filepath:
                self.load_audio(filepath)
    
    def open_folder_dialog(self):
        folder = QFileDialog.getExistingDirectory(self, "選擇包含音頻文件的文件夾")
        if folder:
            if self.batch_mode:
                self.scan_folder_for_audio(folder)
            else:
                # 在非批量模式下，只加載文件夾中的第一個(gè)音頻文件
                audio_files = self.find_audio_files(folder)
                if audio_files:
                    self.load_audio(audio_files[0])
                else:
                    QMessageBox.warning(self, "?? 警告", "文件夾中沒有找到支持的音頻文件")
    
    def scan_folder_for_audio(self, folder):
        audio_files = self.find_audio_files(folder)
        if audio_files:
            self.add_batch_files(audio_files)
        else:
            QMessageBox.warning(self, "?? 警告", "文件夾中沒有找到支持的音頻文件")
    
    def find_audio_files(self, folder):
        supported_formats = ('.wav', '.mp3', '.ogg', '.flac')
        audio_files = []
        for root, _, files in os.walk(folder):
            for file in files:
                if file.lower().endswith(supported_formats):
                    audio_files.append(os.path.join(root, file))
        return audio_files
    
    def handle_dropped_files(self, urls):
        filepaths = [url.toLocalFile() for url in urls if url.isLocalFile()]
        
        if not filepaths:
            return
            
        if self.batch_mode:
            # 在批量模式下，處理所有拖放的文件和文件夾
            all_files = []
            for path in filepaths:
                if os.path.isfile(path):
                    if path.lower().endswith(('.wav', '.mp3', '.ogg', '.flac')):
                        all_files.append(path)
                elif os.path.isdir(path):
                    all_files.extend(self.find_audio_files(path))
            
            if all_files:
                self.add_batch_files(all_files)
            else:
                QMessageBox.warning(self, "?? 警告", "沒有找到支持的音頻文件")
        else:
            # 在非批量模式下，只處理第一個(gè)音頻文件
            for path in filepaths:
                if os.path.isfile(path) and path.lower().endswith(('.wav', '.mp3', '.ogg', '.flac')):
                    self.load_audio(path)
                    break
                elif os.path.isdir(path):
                    audio_files = self.find_audio_files(path)
                    if audio_files:
                        self.load_audio(audio_files[0])
                        break
            else:
                QMessageBox.warning(self, "?? 警告", "請(qǐng)拖放音頻文件 (.wav, .mp3, .ogg, .flac)")
    
    def add_batch_files(self, files):
        new_files = [f for f in files if f not in self.batch_files]
        if new_files:
            self.batch_files.extend(new_files)
            self.update_batch_list()
        else:
            QMessageBox.information(self, "提示", "所有文件已在處理列表中")
    
    def update_batch_list(self):
        count = len(self.batch_files)
        self.batch_list_label.setText(f"待處理文件列表 ({count}):")
        
        if count <= 5:
            file_list = "\n".join(os.path.basename(f) for f in self.batch_files)
        else:
            file_list = "\n".join(os.path.basename(f) for f in self.batch_files[:5]) + f"\n...和其他 {count-5} 個(gè)文件"
        
        self.batch_files_text.setText(file_list)
    
    def load_audio(self, filepath):
        try:
            self.progress_bar.setVisible(True)
            self.progress_bar.setRange(0, 0)
            self.status_bar.showMessage("?? 正在加載音頻...")
            QApplication.processEvents()
            
            self.audio_data, self.sample_rate = librosa.load(filepath, sr=None)
            self.filename = os.path.basename(filepath)
            self.file_label.setText(f"? 已選擇: {self.filename}")
            self.save_button.setEnabled(False)
            
            self.status_bar.showMessage("?? 音頻加載完成")
        except Exception as e:
            QMessageBox.critical(self, "? 錯(cuò)誤", f"加載音頻文件失敗: {str(e)}")
            self.status_bar.showMessage("?? 加載失敗")
        finally:
            self.progress_bar.setVisible(False)
    
    def enhanced_reduce_noise(self, y, sr, strength=0.5, stationary=True):
        """增強(qiáng)的降噪算法"""
        try:
            # 第一級(jí)：譜減法降噪
            y_denoised = nr.reduce_noise(
                y=y,
                sr=sr,
                stationary=stationary,
                prop_decrease=min(0.95, strength),  # 限制最大降噪強(qiáng)度
                n_fft=2048,
                win_length=2048,
                hop_length=512
            )
            
            # 第二級(jí)：Butterworth濾波器去除高頻噪聲
            sos = signal.butter(10, 0.05, 'highpass', fs=sr, output='sos')
            y_filtered = signal.sosfilt(sos, y_denoised)
            
            return y_filtered
        except Exception as e:
            print(f"降噪過程中出錯(cuò): {str(e)}")
            return y
    
    def enhance_voice_volume(self, y, gain=0.0):
        """簡單的人聲音量增強(qiáng) (0.0-1.0對(duì)應(yīng)0%-100%增益)"""
        try:
            if gain == 0.0:  # 0%增益時(shí)不處理
                return y
            # 增益公式：output = input * (1 + gain)
            return y * (1.0 + gain)
        except Exception as e:
            print(f"人聲增強(qiáng)過程中出錯(cuò): {str(e)}")
            return y
    
    def process_audio(self):
        if self.batch_mode:
            self.process_batch()
            return
            
        if self.audio_data is None:
            QMessageBox.warning(self, "?? 警告", "請(qǐng)先選擇音頻文件")
            return
        
        try:
            # 獲取參數(shù)
            strength = self.strength_slider.value() / 10.0  # 轉(zhuǎn)換為0.1-1.0范圍
            stationary = self.mode_combo.currentIndex() == 0
            voice_gain = self.voice_slider.value() / 100.0  # 轉(zhuǎn)換為0.0-1.0范圍
            
            self.progress_bar.setVisible(True)
            self.status_bar.showMessage("?? 正在處理音頻...")
            QApplication.processEvents()
            
            # 1. 降噪處理
            self.processed_audio = self.enhanced_reduce_noise(
                y=self.audio_data,
                sr=self.sample_rate,
                strength=strength,
                stationary=stationary
            )
            
            # 2. 人聲音量增強(qiáng) (只有g(shù)ain>0時(shí)才處理)
            if voice_gain > 0:
                self.processed_audio = self.enhance_voice_volume(
                    y=self.processed_audio,
                    gain=voice_gain
                )
            
            # 3. 音量歸一化
            self.processed_audio = librosa.util.normalize(self.processed_audio)
            
            # 顯示波形對(duì)比
            self.show_waveform_comparison()
            
            self.save_button.setEnabled(True)
            self.status_bar.showMessage("?? 音頻處理完成")
            QMessageBox.information(self, "? 完成", "音頻處理完成！")
        except Exception as e:
            QMessageBox.critical(self, "? 錯(cuò)誤", f"處理音頻時(shí)出錯(cuò): {str(e)}")
            self.status_bar.showMessage("?? 處理失敗")
        finally:
            self.progress_bar.setVisible(False)
    
    def process_batch(self):
        if not self.batch_files:
            QMessageBox.warning(self, "?? 警告", "沒有可處理的文件")
            return
            
        # 選擇輸出目錄
        output_dir = QFileDialog.getExistingDirectory(self, "選擇輸出文件夾")
        if not output_dir:
            return
            
        # 獲取參數(shù)
        strength = self.strength_slider.value() / 10.0
        stationary = self.mode_combo.currentIndex() == 0
        voice_gain = self.voice_slider.value() / 100.0
        
        # 創(chuàng)建批處理器
        self.batch_processor = BatchProcessor(
            files=self.batch_files,
            output_dir=output_dir,
            strength=strength,
            stationary=stationary,
            voice_gain=voice_gain
        )
        
        # 連接信號(hào)
        self.batch_processor.progress_updated.connect(self.update_batch_progress)
        self.batch_processor.finished.connect(self.batch_complete)
        
        # 禁用UI控件
        self.set_ui_enabled(False)
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(0)
        self.status_bar.showMessage("?? 開始批量處理...")
        
        # 啟動(dòng)處理線程
        self.batch_processor.start()
    
    def update_batch_progress(self, progress, message):
        self.progress_bar.setValue(progress)
        self.status_bar.showMessage(message)
    
    def batch_complete(self):
        self.set_ui_enabled(True)
        self.progress_bar.setVisible(False)
        self.status_bar.showMessage("?? 批量處理完成")
        QMessageBox.information(self, "? 完成", "批量處理完成！")
    
    def set_ui_enabled(self, enabled):
        self.load_button.setEnabled(enabled)
        self.load_folder_button.setEnabled(enabled)
        self.process_button.setEnabled(enabled)
        self.save_button.setEnabled(enabled)
        self.batch_process_button.setEnabled(enabled)
        self.strength_slider.setEnabled(enabled)
        self.voice_slider.setEnabled(enabled)
        self.mode_combo.setEnabled(enabled)
        self.batch_checkbox.setEnabled(enabled)

    def show_waveform_comparison(self):
        """顯示處理前后的波形對(duì)比圖"""
        plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
        plt.rcParams['axes.unicode_minus'] = False
        
        plt.figure(figsize=(12, 6))
        
        # 原始音頻波形
        plt.subplot(2, 1, 1)
        librosa.display.waveshow(self.audio_data, sr=self.sample_rate, color='b')
        plt.title("原始音頻波形")
        plt.xlabel("時(shí)間 (秒)")
        plt.ylabel("振幅")
        
        # 處理后音頻波形
        plt.subplot(2, 1, 2)
        librosa.display.waveshow(self.processed_audio, sr=self.sample_rate, color='g')
        plt.title("處理后音頻波形 (降噪+人聲增強(qiáng))")
        plt.xlabel("時(shí)間 (秒)")
        plt.ylabel("振幅")
        
        plt.tight_layout()
        plt.show()
    
    def save_audio(self):
        if self.processed_audio is None:
            QMessageBox.warning(self, "?? 警告", "沒有可保存的處理結(jié)果")
            return
        
        save_path, _ = QFileDialog.getSaveFileName(
            self, "保存處理后的音頻", 
            f"processed_{self.filename}", 
            "WAV文件 (*.wav);;所有文件 (*.*)"
        )
        
        if save_path:
            try:
                self.progress_bar.setVisible(True)
                self.status_bar.showMessage("?? 正在保存音頻...")
                QApplication.processEvents()
                
                sf.write(save_path, self.processed_audio, self.sample_rate)
                
                self.status_bar.showMessage("?? 音頻保存成功")
                QMessageBox.information(self, "? 成功", f"音頻已保存到: {save_path}")
            except Exception as e:
                QMessageBox.critical(self, "? 錯(cuò)誤", f"保存音頻失敗: {str(e)}")
                self.status_bar.showMessage("?? 保存失敗")
            finally:
                self.progress_bar.setVisible(False)

if __name__ == "__main__":
    app = QApplication([])
    # 設(shè)置全局字體
    app.setFont(QFont('Microsoft YaHei'))
    window = AudioDenoiserUI()
    window.show()
    app.exec_()