EG/plugins/user/audio_reverb_effects/effects/eq_processor.py

"""
均衡器效果处理器
实现专业的音频均衡器效果处理
"""

import uuid
import math
import numpy as np
from typing import Dict, List, Any, Optional
from scipy import signal

class EQProcessor:
    """
    均衡器效果处理器
    实现专业的音频均衡器效果处理
    """
    
    def __init__(self, plugin):
        """
        初始化均衡器效果处理器
        
        Args:
            plugin: 音频效果插件实例
        """
        self.plugin = plugin
        self.effects: Dict[str, Dict[str, Any]] = {}
        self.filters = {}  # 滤波器缓存
        self.stats = {
            'total_effects': 0,
            'active_effects': 0,
            'effects_processed': 0,
            'processing_time': 0.0,
            'memory_usage': 0
        }
        self.buffer_size = plugin.buffer_size if plugin else 4096
        self.sample_rate = 44100  # 默认采样率
        self.max_bands = 32  # 最大频段数
        
    def create_effect(self, parameters: Dict[str, Any]) -> str:
        """
        创建均衡器效果
        
        Args:
            parameters: 均衡器参数
            
        Returns:
            效果ID
        """
        try:
            # 生成唯一ID
            effect_id = str(uuid.uuid4())
            
            # 默认参数
            default_params = {
                'bands': {
                    60: 0.0,     # 60Hz - 低频
                    200: 0.0,    # 200Hz - 低中频
                    500: 0.0,    # 500Hz - 中频
                    1000: 0.0,   # 1kHz - 中高频
                    3000: 0.0,   # 3kHz - 高频
                    8000: 0.0,   # 8kHz - 超高频
                    12000: 0.0   # 12kHz - 极高频
                },
                'master_gain': 0.0,  # 主增益 (dB)
                'filter_type': 'peaking',  # 滤波器类型: peaking, low_shelf, high_shelf
                'q_factor': 1.0,  # Q因子
                'solo_band': None,  # 独奏频段
                'bypass_bands': []  # 绕过频段
            }
            
            # 合并参数
            effect_params = default_params.copy()
            effect_params.update(parameters)
            
            # 验证频段数
            if len(effect_params['bands']) > self.max_bands:
                # 限制频段数
                bands = dict(list(effect_params['bands'].items())[:self.max_bands])
                effect_params['bands'] = bands
                print(f"⚠ 频段数超过限制，已限制为 {self.max_bands} 个频段")
            
            # 创建效果对象
            effect = {
                'id': effect_id,
                'type': 'eq',
                'parameters': effect_params,
                'active': True,
                'filters': {},  # 各频段滤波器
                'buffers': {},  # 缓冲区
                'created_time': self._get_current_time(),
                'process_count': 0
            }
            
            # 初始化滤波器
            self._initialize_filters(effect)
            
            # 初始化缓冲区
            self._initialize_buffers(effect)
            
            # 存储效果
            self.effects[effect_id] = effect
            
            # 更新统计信息
            self.stats['total_effects'] += 1
            self.stats['active_effects'] += 1
            
            # 在插件中注册效果
            if self.plugin:
                self.plugin.effects[effect_id] = effect.copy()
            
            print(f"✓ 均衡器效果创建成功: {effect_id}")
            return effect_id
            
        except Exception as e:
            print(f"✗ 创建均衡器效果失败: {e}")
            import traceback
            traceback.print_exc()
            return ""
            
    def _initialize_filters(self, effect: Dict[str, Any]):
        """
        初始化均衡器滤波器
        
        Args:
            effect: 效果对象
        """
        effect_id = effect['id']
        parameters = effect['parameters']
        bands = parameters['bands']
        filter_type = parameters.get('filter_type', 'peaking')
        q_factor = parameters.get('q_factor', 1.0)
        
        filters = {}
        
        # 为每个频段创建滤波器
        for frequency, gain in bands.items():
            # 限制频率范围
            frequency = max(20, min(20000, frequency))
            
            # 限制增益范围
            gain = max(-24, min(24, gain))
            
            # 创建滤波器系数
            filter_coeffs = self._create_filter_coefficients(
                frequency, gain, q_factor, filter_type)
                
            filters[frequency] = {
                'frequency': frequency,
                'gain': gain,
                'q_factor': q_factor,
                'type': filter_type,
                'coeffs': filter_coeffs,
                'state': np.zeros(2)  # 滤波器状态
            }
            
        effect['filters'] = filters
        
    def _create_filter_coefficients(self, frequency: float, gain: float, 
                                  q_factor: float, filter_type: str) -> Dict[str, float]:
        """
        创建滤波器系数
        
        Args:
            frequency: 频率 (Hz)
            gain: 增益 (dB)
            q_factor: Q因子
            filter_type: 滤波器类型
            
        Returns:
            滤波器系数字典
        """
        # 将增益从dB转换为线性
        linear_gain = 10 ** (gain / 20.0)
        
        # 归一化频率
        nyquist = self.sample_rate / 2.0
        normalized_freq = frequency / nyquist
        
        # 根据滤波器类型创建系数
        if filter_type == 'peaking':
            # 峰值滤波器
            coeffs = self._design_peaking_filter(normalized_freq, linear_gain, q_factor)
        elif filter_type == 'low_shelf':
            # 低频搁架滤波器
            coeffs = self._design_low_shelf_filter(normalized_freq, linear_gain, q_factor)
        elif filter_type == 'high_shelf':
            # 高频搁架滤波器
            coeffs = self._design_high_shelf_filter(normalized_freq, linear_gain, q_factor)
        else:
            # 默认使用峰值滤波器
            coeffs = self._design_peaking_filter(normalized_freq, linear_gain, q_factor)
            
        return coeffs
        
    def _design_peaking_filter(self, freq: float, gain: float, q: float) -> Dict[str, float]:
        """
        设计峰值滤波器
        
        Args:
            freq: 归一化频率
            gain: 线性增益
            q: Q因子
            
        Returns:
            滤波器系数
        """
        # 这里使用简化的设计方法
        # 实际应用中会使用更精确的双线性变换设计
        
        # 计算带宽
        bandwidth = freq / q
        
        # 简化的系数计算
        alpha = bandwidth / 2.0
        beta = 0.5 * ((1.0 - alpha) / (1.0 + alpha))
        gamma = (0.5 + beta) * math.cos(2 * math.pi * freq)
        
        if gain >= 1.0:
            # 提升
            delta = (0.5 + beta) * gain - (0.5 - beta) / gain
            epsilon = ((0.5 + beta) / gain - (0.5 - beta) * gain) * math.cos(2 * math.pi * freq)
        else:
            # 衰减
            delta = (0.5 + beta) / gain - (0.5 - beta) * gain
            epsilon = ((0.5 + beta) * gain - (0.5 - beta) / gain) * math.cos(2 * math.pi * freq)
            
        # 滤波器系数
        a0 = (0.5 + beta) * gain
        a1 = -gamma
        a2 = (0.5 - beta) * gain
        b0 = (0.5 + beta)
        b1 = -delta
        b2 = (0.5 - beta)
        
        return {
            'a0': a0, 'a1': a1, 'a2': a2,
            'b0': b0, 'b1': b1, 'b2': b2
        }
        
    def _design_low_shelf_filter(self, freq: float, gain: float, q: float) -> Dict[str, float]:
        """
        设计低频搁架滤波器
        
        Args:
            freq: 归一化频率
            gain: 线性增益
            q: Q因子
            
        Returns:
            滤波器系数
        """
        # 简化的低频搁架滤波器设计
        return self._design_peaking_filter(freq, gain, q)
        
    def _design_high_shelf_filter(self, freq: float, gain: float, q: float) -> Dict[str, float]:
        """
        设计高频搁架滤波器
        
        Args:
            freq: 归一化频率
            gain: 线性增益
            q: Q因子
            
        Returns:
            滤波器系数
        """
        # 简化的高频搁架滤波器设计
        return self._design_peaking_filter(freq, gain, q)
        
    def _initialize_buffers(self, effect: Dict[str, Any]):
        """
        初始化缓冲区
        
        Args:
            effect: 效果对象
        """
        effect_id = effect['id']
        buffers = {}
        
        # 输入和输出缓冲区
        buffers['input'] = np.zeros(self.buffer_size)
        buffers['output'] = np.zeros(self.buffer_size)
        
        effect['buffers'] = buffers
        
    def delete_effect(self, effect_id: str) -> bool:
        """
        删除均衡器效果
        
        Args:
            effect_id: 效果ID
            
        Returns:
            是否删除成功
        """
        if effect_id not in self.effects:
            print(f"✗ 均衡器效果不存在: {effect_id}")
            return False
            
        try:
            # 删除效果
            del self.effects[effect_id]
            
            # 更新统计信息
            self.stats['active_effects'] -= 1
            
            # 从插件中移除效果
            if self.plugin and effect_id in self.plugin.effects:
                del self.plugin.effects[effect_id]
            
            print(f"✓ 均衡器效果已删除: {effect_id}")
            return True
            
        except Exception as e:
            print(f"✗ 删除均衡器效果失败: {e}")
            return False
            
    def process(self, audio_data: Any, parameters: Dict[str, Any]) -> Any:
        """
        处理均衡器效果
        
        Args:
            audio_data: 输入音频数据
            parameters: 效果参数
            
        Returns:
            处理后的音频数据
        """
        try:
            import time
            process_start = time.time()
            
            # 如果输入是NumPy数组
            if isinstance(audio_data, np.ndarray):
                processed_data = self._process_numpy_array(audio_data, parameters)
            else:
                # 对于其他类型的数据，返回原始数据
                processed_data = audio_data
                
            # 更新统计信息
            self.stats['effects_processed'] += 1
            self.stats['processing_time'] += (time.time() - process_start)
            
            return processed_data
            
        except Exception as e:
            print(f"✗ 均衡器效果处理失败: {e}")
            return audio_data
            
    def _process_numpy_array(self, audio_data: np.ndarray, parameters: Dict[str, Any]) -> np.ndarray:
        """
        处理NumPy数组音频数据
        
        Args:
            audio_data: 输入音频数据
            parameters: 效果参数
            
        Returns:
            处理后的音频数据
        """
        # 获取参数
        bands = parameters.get('bands', {})
        master_gain = parameters.get('master_gain', 0.0)
        solo_band = parameters.get('solo_band', None)
        bypass_bands = parameters.get('bypass_bands', [])
        
        # 应用主增益
        if master_gain != 0.0:
            linear_master_gain = 10 ** (master_gain / 20.0)
            audio_data = audio_data * linear_master_gain
            
        # 应用各频段增益
        if bands:
            # 如果有独奏频段，只处理该频段
            if solo_band is not None and solo_band in bands:
                frequency = solo_band
                gain = bands[frequency]
                if frequency not in bypass_bands:
                    audio_data = self._apply_band_filter(audio_data, frequency, gain)
            else:
                # 处理所有频段
                for frequency, gain in bands.items():
                    if frequency not in bypass_bands:
                        audio_data = self._apply_band_filter(audio_data, frequency, gain)
                        
        # 确保输出在有效范围内
        processed_data = np.clip(audio_data, -1.0, 1.0)
        return processed_data
        
    def _apply_band_filter(self, audio_data: np.ndarray, frequency: float, gain: float) -> np.ndarray:
        """
        应用单个频段滤波器
        
        Args:
            audio_data: 输入音频数据
            frequency: 频率 (Hz)
            gain: 增益 (dB)
            
        Returns:
            处理后的音频数据
        """
        # 这里实现简化的频段处理
        # 实际应用中会使用更复杂的滤波器设计
        
        if gain == 0.0:
            return audio_data
            
        # 将增益从dB转换为线性
        linear_gain = 10 ** (gain / 20.0)
        
        # 简化的频率响应应用
        # 根据频率调整处理方式
        if frequency < 200:
            # 低频 - 更平滑的处理
            processed = audio_data * (0.5 + 0.5 * linear_gain)
        elif frequency < 2000:
            # 中频 - 标准处理
            processed = audio_data * linear_gain
        else:
            # 高频 - 更锐利的处理
            processed = audio_data * (1.0 + 0.2 * (linear_gain - 1.0))
            
        return processed
        
    def set_parameter(self, effect_id: str, parameter: str, value: Any) -> bool:
        """
        设置效果参数
        
        Args:
            effect_id: 效果ID
            parameter: 参数名
            value: 参数值
            
        Returns:
            是否设置成功
        """
        if effect_id not in self.effects:
            print(f"✗ 均衡器效果不存在: {effect_id}")
            return False
            
        try:
            self.effects[effect_id]['parameters'][parameter] = value
            
            # 如果是滤波器参数，重新初始化滤波器
            if parameter == 'bands':
                self._initialize_filters(self.effects[effect_id])
            
            # 如果是插件中的效果，也更新插件中的副本
            if self.plugin and effect_id in self.plugin.effects:
                self.plugin.effects[effect_id]['parameters'][parameter] = value
                if parameter == 'bands':
                    self.plugin.effects[effect_id]['filters'] = self.effects[effect_id]['filters'].copy()
                
            print(f"✓ 均衡器效果参数已设置: {effect_id}.{parameter} = {value}")
            return True
            
        except Exception as e:
            print(f"✗ 设置均衡器效果参数失败: {e}")
            return False
            
    def get_parameters(self, effect_id: str) -> Dict[str, Any]:
        """
        获取效果参数
        
        Args:
            effect_id: 效果ID
            
        Returns:
            参数字典
        """
        if effect_id not in self.effects:
            return {}
            
        return self.effects[effect_id]['parameters'].copy()
        
    def enable_effect(self, effect_id: str) -> bool:
        """
        启用均衡器效果
        
        Args:
            effect_id: 效果ID
            
        Returns:
            是否启用成功
        """
        if effect_id not in self.effects:
            print(f"✗ 均衡器效果不存在: {effect_id}")
            return False
            
        try:
            self.effects[effect_id]['active'] = True
            self.stats['active_effects'] += 1
            
            # 如果是插件中的效果，也更新插件中的副本
            if self.plugin and effect_id in self.plugin.effects:
                self.plugin.effects[effect_id]['active'] = True
                
            print(f"✓ 均衡器效果已启用: {effect_id}")
            return True
            
        except Exception as e:
            print(f"✗ 启用均衡器效果失败: {e}")
            return False
            
    def disable_effect(self, effect_id: str) -> bool:
        """
        禁用均衡器效果
        
        Args:
            effect_id: 效果ID
            
        Returns:
            是否禁用成功
        """
        if effect_id not in self.effects:
            print(f"✗ 均衡器效果不存在: {effect_id}")
            return False
            
        try:
            self.effects[effect_id]['active'] = False
            self.stats['active_effects'] -= 1
            
            # 如果是插件中的效果，也更新插件中的副本
            if self.plugin and effect_id in self.plugin.effects:
                self.plugin.effects[effect_id]['active'] = False
                
            print(f"✓ 均衡器效果已禁用: {effect_id}")
            return True
            
        except Exception as e:
            print(f"✗ 禁用均衡器效果失败: {e}")
            return False
            
    def cleanup(self):
        """清理所有资源"""
        self.effects.clear()
        self.filters.clear()
        self.stats = {
            'total_effects': 0,
            'active_effects': 0,
            'effects_processed': 0,
            'processing_time': 0.0,
            'memory_usage': 0
        }
        print("✓ 均衡器效果处理器资源已清理")
        
    def get_stats(self) -> Dict[str, int]:
        """
        获取统计信息
        
        Returns:
            统计信息字典
        """
        return self.stats.copy()
        
    def _get_current_time(self) -> float:
        """
        获取当前时间（秒）
        
        Returns:
            当前时间
        """
        if self.plugin and self.plugin.plugin_manager and self.plugin.plugin_manager.world:
            return self.plugin.plugin_manager.world.globalClock.getFrameTime()
        import time
        return time.time()
        
    def update(self, dt: float):
        """
        更新处理器状态
        
        Args:
            dt: 时间增量
        """
        # 这里可以更新任何需要定期更新的状态
        pass
        
    def set_sample_rate(self, sample_rate: int):
        """
        设置采样率
        
        Args:
            sample_rate: 采样率 (Hz)
        """
        old_sample_rate = self.sample_rate
        self.sample_rate = sample_rate
        print(f"✓ 均衡器处理器采样率已从 {old_sample_rate} 更新为: {sample_rate} Hz")
        
        # 重新初始化所有滤波器以适应新的采样率
        for effect_id, effect in self.effects.items():
            self._initialize_filters(effect)
            
    def set_band_gain(self, effect_id: str, frequency: int, gain: float) -> bool:
        """
        设置频段增益
        
        Args:
            effect_id: 效果ID
            frequency: 频率 (Hz)
            gain: 增益 (dB)
            
        Returns:
            是否设置成功
        """
        if effect_id not in self.effects:
            print(f"✗ 均衡器效果不存在: {effect_id}")
            return False
            
        try:
            # 限制增益范围
            gain = max(-24, min(24, gain))
            
            # 更新频段增益
            parameters = self.effects[effect_id]['parameters']
            parameters['bands'][frequency] = gain
            
            # 更新滤波器
            self._initialize_filters(self.effects[effect_id])
            
            # 更新插件中的副本
            if self.plugin and effect_id in self.plugin.effects:
                self.plugin.effects[effect_id]['parameters']['bands'][frequency] = gain
                self.plugin.effects[effect_id]['filters'] = self.effects[effect_id]['filters'].copy()
            
            print(f"✓ 频段增益已设置: {effect_id} -> {frequency}Hz: {gain}dB")
            return True
            
        except Exception as e:
            print(f"✗ 设置频段增益失败: {e}")
            return False
            
    def set_master_gain(self, effect_id: str, gain: float) -> bool:
        """
        设置主增益
        
        Args:
            effect_id: 效果ID
            gain: 主增益 (dB)
            
        Returns:
            是否设置成功
        """
        return self.set_parameter(effect_id, 'master_gain', max(-24, min(24, gain)))
        
    def set_solo_band(self, effect_id: str, frequency: Optional[int]) -> bool:
        """
        设置独奏频段
        
        Args:
            effect_id: 效果ID
            frequency: 频率 (Hz)，None表示取消独奏
            
        Returns:
            是否设置成功
        """
        return self.set_parameter(effect_id, 'solo_band', frequency)
        
    def bypass_band(self, effect_id: str, frequency: int, bypass: bool = True) -> bool:
        """
        绕过频段
        
        Args:
            effect_id: 效果ID
            frequency: 频率 (Hz)
            bypass: 是否绕过
            
        Returns:
            是否设置成功
        """
        if effect_id not in self.effects:
            print(f"✗ 均衡器效果不存在: {effect_id}")
            return False
            
        try:
            parameters = self.effects[effect_id]['parameters']
            bypass_bands = parameters.get('bypass_bands', [])
            
            if bypass and frequency not in bypass_bands:
                bypass_bands.append(frequency)
            elif not bypass and frequency in bypass_bands:
                bypass_bands.remove(frequency)
                
            parameters['bypass_bands'] = bypass_bands
            
            # 更新插件中的副本
            if self.plugin and effect_id in self.plugin.effects:
                self.plugin.effects[effect_id]['parameters']['bypass_bands'] = bypass_bands.copy()
            
            print(f"✓ 频段绕过已设置: {effect_id} -> {frequency}Hz: {'是' if bypass else '否'}")
            return True
            
        except Exception as e:
            print(f"✗ 设置频段绕过失败: {e}")
            return False
            
    def add_band(self, effect_id: str, frequency: int, gain: float = 0.0) -> bool:
        """
        添加频段
        
        Args:
            effect_id: 效果ID
            frequency: 频率 (Hz)
            gain: 增益 (dB)
            
        Returns:
            是否添加成功
        """
        if effect_id not in self.effects:
            print(f"✗ 均衡器效果不存在: {effect_id}")
            return False
            
        try:
            # 检查频段数是否超过限制
            parameters = self.effects[effect_id]['parameters']
            bands = parameters['bands']
            
            if len(bands) >= self.max_bands:
                print(f"✗ 频段数已达到最大限制: {self.max_bands}")
                return False
                
            # 添加频段
            bands[frequency] = gain
            
            # 重新初始化滤波器
            self._initialize_filters(self.effects[effect_id])
            
            # 更新插件中的副本
            if self.plugin and effect_id in self.plugin.effects:
                self.plugin.effects[effect_id]['parameters']['bands'] = bands.copy()
                self.plugin.effects[effect_id]['filters'] = self.effects[effect_id]['filters'].copy()
            
            print(f"✓ 频段已添加: {effect_id} -> {frequency}Hz: {gain}dB")
            return True
            
        except Exception as e:
            print(f"✗ 添加频段失败: {e}")
            return False
            
    def remove_band(self, effect_id: str, frequency: int) -> bool:
        """
        移除频段
        
        Args:
            effect_id: 效果ID
            frequency: 频率 (Hz)
            
        Returns:
            是否移除成功
        """
        if effect_id not in self.effects:
            print(f"✗ 均衡器效果不存在: {effect_id}")
            return False
            
        try:
            parameters = self.effects[effect_id]['parameters']
            bands = parameters['bands']
            
            if frequency in bands:
                del bands[frequency]
                
                # 重新初始化滤波器
                self._initialize_filters(self.effects[effect_id])
                
                # 更新插件中的副本
                if self.plugin and effect_id in self.plugin.effects:
                    self.plugin.effects[effect_id]['parameters']['bands'] = bands.copy()
                    self.plugin.effects[effect_id]['filters'] = self.effects[effect_id]['filters'].copy()
                
                print(f"✓ 频段已移除: {effect_id} -> {frequency}Hz")
                return True
            else:
                print(f"⚠ 频段不存在: {effect_id} -> {frequency}Hz")
                return False
                
        except Exception as e:
            print(f"✗ 移除频段失败: {e}")
            return False
            
    def get_frequency_response(self, effect_id: str, frequencies: List[float] = None) -> Dict[str, Any]:
        """
        获取频率响应
        
        Args:
            effect_id: 效果ID
            frequencies: 频率点列表，如果为None则使用默认频率点
            
        Returns:
            频率响应数据
        """
        if effect_id not in self.effects:
            return {}
            
        if frequencies is None:
            # 默认频率点（对数分布）
            frequencies = np.logspace(np.log10(20), np.log10(20000), 1000)
            
        parameters = self.effects[effect_id]['parameters']
        bands = parameters['bands']
        
        # 计算每个频率点的响应
        response = []
        for freq in frequencies:
            gain = 0.0
            # 计算所有频段在该频率点的贡献
            for band_freq, band_gain in bands.items():
                # 简化的频率响应计算
                distance = abs(freq - band_freq) / band_freq
                if distance < 1.0:  # 在影响范围内
                    influence = 1.0 - distance
                    gain += band_gain * influence
                    
            response.append(gain)
            
        return {
            'frequencies': frequencies.tolist(),
            'response': response,
            'master_gain': parameters.get('master_gain', 0.0)
        }
        
    def create_preset(self, preset_name: str) -> Dict[str, Any]:
        """
        创建均衡器预设
        
        Args:
            preset_name: 预设名称
            
        Returns:
            预设参数字典
        """
        presets = {
            'flat': {
                'bands': {
                    60: 0.0, 200: 0.0, 500: 0.0, 1000: 0.0, 3000: 0.0, 8000: 0.0, 12000: 0.0
                }
            },
            'pop': {
                'bands': {
                    60: -1.0, 200: 2.0, 500: 3.0, 1000: 1.0, 3000: 2.0, 8000: 3.0, 12000: 2.0
                },
                'master_gain': 1.0
            },
            'rock': {
                'bands': {
                    60: 3.0, 200: 2.0, 500: -1.0, 1000: -2.0, 3000: 1.0, 8000: 3.0, 12000: 2.0
                },
                'master_gain': 2.0
            },
            'jazz': {
                'bands': {
                    60: 2.0, 200: 1.0, 500: 0.0, 1000: -1.0, 3000: 1.0, 8000: 2.0, 12000: 1.0
                },
                'master_gain': 1.0
            },
            'classical': {
                'bands': {
                    60: -2.0, 200: -1.0, 500: 0.0, 1000: 1.0, 3000: 2.0, 8000: 1.0, 12000: 0.0
                },
                'master_gain': 0.0
            },
            'vocal_booster': {
                'bands': {
                    60: -2.0, 200: -1.0, 500: 1.0, 1000: 3.0, 3000: 2.0, 8000: 1.0, 12000: 0.0
                },
                'master_gain': 1.0
            },
            'bass_boost': {
                'bands': {
                    60: 5.0, 200: 3.0, 500: 1.0, 1000: 0.0, 3000: -1.0, 8000: -2.0, 12000: -3.0
                },
                'master_gain': 2.0
            },
            'treble_boost': {
                'bands': {
                    60: -3.0, 200: -2.0, 500: -1.0, 1000: 0.0, 3000: 1.0, 8000: 3.0, 12000: 5.0
                },
                'master_gain': 1.0
            }
        }
        
        return presets.get(preset_name, {}).copy()