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

class ImagingModule(BaseProcessorModule):
    """Stereo Imaging Module."""
    
    def process(self, audio: np.ndarray, **kwargs) -> np.ndarray:
        return audio

    def apply_balancing(self, audio: np.ndarray, balance_db: float) -> np.ndarray:
        """Surgically correct stereo leaning."""
        if audio.ndim == 1: return audio
        if abs(balance_db) < 0.1: return audio
        
        # balance_db > 0 means left is louder. We need to boost right or cut left.
        # We apply half the correction to each side for transparency.
        # adjustment_lin = 10 ** (abs(balance_db)/2 / 20)
        
        l_mult = 10 ** (-balance_db / 2 / 20)
        r_mult = 10 ** (balance_db / 2 / 20)
        
        out = audio.copy()
        out[:, 0] *= l_mult
        out[:, 1] *= r_mult
        
        return out

    def apply_haas_widening(self, audio: np.ndarray, amount: float = 0.0) -> np.ndarray:
        """Synthesize stereo width using the Haas effect (Micro-delay)."""
        if audio.ndim == 1 or amount <= 0: return audio
        sr = self.sample_rate
        
        # 1. High-Pass the processing (Don't touch bass width)
        sos_hp = butter(4, 500, 'hp', fs=sr, output='sos')
        hp_audio = sosfilt(sos_hp, audio, axis=0)
        
        # 2. Apply 15ms delay to the Right channel of the HP signal
        delay_samples = int(sr * 0.015) 
        delayed_r = np.roll(hp_audio[:, 1], delay_samples)
        
        # 3. Create 'Side' energy from the delayed signal
        # We mix the delayed R back into the original
        out = audio.copy()
        # amount scales the injected 'stereo' information
        out[:, 1] = (audio[:, 1] * (1.0 - amount)) + (delayed_r * amount)
        
        return out

    def make_mono_bass(self, audio: np.ndarray, freq: float = 120.0) -> np.ndarray:
        if audio.ndim == 1: return audio
        sr = self.sample_rate
        
        mid = (audio[:, 0] + audio[:, 1]) / 2
        side = (audio[:, 0] - audio[:, 1]) / 2
        
        sos = butter(4, freq, 'hp', fs=sr, output='sos')
        side_high = sosfilt(sos, side)
        
        left = mid + side_high
        right = mid - side_high
        return np.column_stack((left, right))

    def apply_width(self, audio: np.ndarray, low_w=1.0, mid_w=1.0, high_w=1.0) -> np.ndarray:
        if audio.ndim == 1: return audio
        sr = self.sample_rate
        
        # 3-Band Split
        sos_low = butter(4, 200, 'lp', fs=sr, output='sos')
        sos_mid = butter(4, [200, 5000], 'bp', fs=sr, output='sos')
        sos_high = butter(4, 5000, 'hp', fs=sr, output='sos')
        
        def process_band(sos, width):
            l = sosfilt(sos, audio[:, 0])
            r = sosfilt(sos, audio[:, 1])
            m = (l + r) / 2
            s = (l - r) / 2
            s *= width
            return np.column_stack((m + s, m - s))
            
        out = process_band(sos_low, low_w)
        out += process_band(sos_mid, mid_w)
        out += process_band(sos_high, high_w)
        
        return out