audio_processing/effect_chain.py

"""
Effect Chain - Pedalboard 기반 오디오 이펙트 처리
=================================================
V2: 학습 데이터 파라미터 키와 호환성 개선
"""

import numpy as np
import soundfile as sf
from typing import Dict, List, Optional
from pedalboard import (
    Pedalboard, 
    Compressor, 
    Gain,
    HighShelfFilter,
    LowShelfFilter,
    PeakFilter,
    Delay,
    Reverb,
    Distortion,
    Limiter
)


class EffectChain:
    """Pedalboard 기반 이펙트 체인"""
    
    def __init__(self, sample_rate: int = 44100):
        self.sample_rate = sample_rate
        
        self.available_effects = [
            "eq_peak1", "eq_peak2", 
            "eq_lowshelf", "eq_highshelf",
            "distortion", "delay", "compressor",
            "reverb", "limiter"
        ]
    
    def get_available_effects(self) -> List[str]:
        """사용 가능한 이펙트 목록"""
        return self.available_effects
    
    def _get_param(self, params: Dict[str, float], *keys, default=0.0) -> float:
        """
        여러 키 중 하나라도 있으면 반환 (호환성 처리)
        예: 'eq_peak1.params.Q' 또는 'eq_peak1.params.q' 둘 다 지원
        """
        for key in keys:
            if key in params:
                return params[key]
        return default
    
    def _build_pedalboard(self, params: Dict[str, float]) -> Pedalboard:
        """파라미터로 Pedalboard 구성"""
        
        effects = []
        
        # Compressor (항상 적용)
        effects.append(Compressor(
            threshold_db=-18.0,
            ratio=2.0,
            attack_ms=10.0,
            release_ms=100.0
        ))
        
        # EQ Peak 1 (Q/q 둘 다 지원)
        freq1 = self._get_param(params, "eq_peak1.params.freq", default=1000.0)
        gain1 = self._get_param(params, "eq_peak1.params.gain", default=0.0)
        q1 = self._get_param(params, "eq_peak1.params.Q", "eq_peak1.params.q", default=1.0)
        if abs(gain1) > 0.1:
            effects.append(PeakFilter(
                cutoff_frequency_hz=max(20, min(20000, freq1)),
                gain_db=max(-12, min(12, gain1)),
                q=max(0.1, min(10, q1))
            ))
        
        # EQ Peak 2 (Q/q 둘 다 지원)
        freq2 = self._get_param(params, "eq_peak2.params.freq", default=4000.0)
        gain2 = self._get_param(params, "eq_peak2.params.gain", default=0.0)
        q2 = self._get_param(params, "eq_peak2.params.Q", "eq_peak2.params.q", default=1.0)
        if abs(gain2) > 0.1:
            effects.append(PeakFilter(
                cutoff_frequency_hz=max(20, min(20000, freq2)),
                gain_db=max(-12, min(12, gain2)),
                q=max(0.1, min(10, q2))
            ))
        
        # Low Shelf
        freq_low = self._get_param(params, "eq_lowshelf.params.freq", default=200.0)
        gain_low = self._get_param(params, "eq_lowshelf.params.gain", default=0.0)
        if abs(gain_low) > 0.1:
            effects.append(LowShelfFilter(
                cutoff_frequency_hz=max(20, min(2000, freq_low)),
                gain_db=max(-12, min(12, gain_low)),
                q=0.707
            ))
        
        # High Shelf
        freq_high = self._get_param(params, "eq_highshelf.params.freq", default=8000.0)
        gain_high = self._get_param(params, "eq_highshelf.params.gain", default=0.0)
        if abs(gain_high) > 0.1:
            effects.append(HighShelfFilter(
                cutoff_frequency_hz=max(1000, min(20000, freq_high)),
                gain_db=max(-12, min(12, gain_high)),
                q=0.707
            ))
        
        # Distortion
        dist_amount = self._get_param(params, "distortion_amount", default=0.0)
        if dist_amount > 0.01:
            # 학습 데이터는 sigmoid 후 0.1 곱하므로 최대 0.1
            # 여기서는 drive_db로 변환 (0~20dB 범위)
            effects.append(Distortion(
                drive_db=max(0, min(20, dist_amount * 100))
            ))
        
        # Delay
        delay_time = self._get_param(params, "delay.delay_time", default=0.02)
        delay_feedback = self._get_param(params, "delay.feedback", default=0.3)
        delay_mix = self._get_param(params, "delay.mix", default=0.2)
        if delay_mix > 0.01:
            effects.append(Delay(
                delay_seconds=max(0.01, min(1.0, delay_time)),
                feedback=max(0.0, min(0.9, delay_feedback)),
                mix=max(0.0, min(1.0, delay_mix))
            ))
        
        # Limiter (항상 마지막에)
        effects.append(Limiter(threshold_db=-1.0))
        
        return Pedalboard(effects)
    
    def process(
        self, 
        input_path: str, 
        output_path: str, 
        parameters: Dict[str, float]
    ) -> bool:
        """오디오 파일 처리"""
        try:
            # 오디오 로드
            audio, sr = sf.read(input_path)
            
            # 모노/스테레오 처리
            if len(audio.shape) == 1:
                audio = audio.reshape(-1, 1)
            
            # float32로 변환
            audio = audio.astype(np.float32)
            
            # Pedalboard 구성
            board = self._build_pedalboard(parameters)
            
            # 처리
            processed = board(audio, sr)
            
            # Wet/Dry 믹스 (final_wet_mix 파라미터 사용)
            wet_mix = self._get_param(parameters, "final_wet_mix", default=0.5)
            wet_mix = max(0.0, min(1.0, wet_mix))
            
            # 길이 맞추기
            min_len = min(len(audio), len(processed))
            output = audio[:min_len] * (1 - wet_mix) + processed[:min_len] * wet_mix
            
            # 클리핑 방지
            output = np.clip(output, -1.0, 1.0)
            
            # 저장
            sf.write(output_path, output, sr)
            
            print(f"[EffectChain] ✅ 처리 완료: {output_path}")
            print(f"[EffectChain] 적용된 파라미터:")
            print(f"  - EQ Peak1 Gain: {self._get_param(parameters, 'eq_peak1.params.gain'):.2f} dB")
            print(f"  - EQ Peak2 Gain: {self._get_param(parameters, 'eq_peak2.params.gain'):.2f} dB")
            print(f"  - Low Shelf Gain: {self._get_param(parameters, 'eq_lowshelf.params.gain'):.2f} dB")
            print(f"  - High Shelf Gain: {self._get_param(parameters, 'eq_highshelf.params.gain'):.2f} dB")
            print(f"  - Distortion: {self._get_param(parameters, 'distortion_amount'):.3f}")
            print(f"  - Delay Mix: {self._get_param(parameters, 'delay.mix'):.2f}")
            print(f"  - Wet Mix: {wet_mix:.2f}")
            
            return True
            
        except Exception as e:
            print(f"[EffectChain] ❌ 처리 실패: {e}")
            import traceback
            traceback.print_exc()
            raise e