import numpy as np
from scipy.signal import butter, sosfilt
from .base import BaseProcessorModule

try:
    from ..pipeline.saturation import oversampled
except ImportError:
    from automaster_app.pipeline.saturation import oversampled


class ExciterModule(BaseProcessorModule):
    """Pro Exciter Module: Harmonic Synthesis & Spectral Saturation."""
    
    def apply_dsp_exciter(self, audio: np.ndarray) -> np.ndarray:
        """
        Harmonic synthesis using high-band waveshaping.
        Generates natural 'air' without the artifacts of pitch shifting.
        """
        sr = self.sample_rate
        # Generate harmonics specifically above 7kHz
        sos_hp = butter(4, 7000, 'hp', fs=sr, output='sos')
        
        def process_channel(channel):
            highs = sosfilt(sos_hp, channel)
            
            # Non-linear waveshaper (Excitement)
            # We use a combination of squaring (even) and cubing (odd)
            # but we filter the result to keep only the high-frequency products
            excited = oversampled(
                lambda u: (np.sign(u) * (np.abs(u) ** 0.5)) * 0.2,
                highs,
            )
            excited_filt = sosfilt(sos_hp, excited)
            
            # Blend back very subtly
            return channel + (excited_filt * 0.15)

        if audio.ndim == 1: 
            return np.clip(process_channel(audio), -1.0, 1.0)
        else:
            left = process_channel(audio[:, 0])
            right = process_channel(audio[:, 1])
            return np.clip(np.column_stack((left, right)), -1.0, 1.0)

    def apply_spectral_exciter(self, audio: np.ndarray, highpass_freq=5000.0, blend_db=-15.0, drive=0.3) -> np.ndarray:
        """Saturation-based exciter refined for transparency."""
        sr = self.sample_rate
        sos_hp = butter(2, highpass_freq, 'hp', fs=sr, output='sos')
        
        def process_channel(channel):
            highs = sosfilt(sos_hp, channel)
            # Use tanh for smooth, analog-style saturation
            excited = oversampled(lambda u: np.tanh(u * (1 + drive * 2)), highs)
            blend_linear = 10 ** (blend_db / 20)
            return channel + (excited * blend_linear)
            
        if audio.ndim == 1: return process_channel(audio)
        else:
            left = process_channel(audio[:, 0])
            right = process_channel(audio[:, 1])
            return np.column_stack((left, right))

    def process(self, audio: np.ndarray, mode: str = "none", **kwargs) -> np.ndarray:
        if mode == "dsp_fast":
            return self.apply_dsp_exciter(audio)
        elif mode == "spectral":
            # Pass drive/blend from kwargs if they exist
            return self.apply_spectral_exciter(audio, **kwargs)
        return audio