import unittest
import numpy as np
from engine.processor import AudioEngine
from scipy import signal

class TestMultiband(unittest.TestCase):
    def setUp(self):
        self.engine = AudioEngine()
        self.engine.sample_rate = 44100

    def test_crossover_reconstruction(self):
        """Test that splitting and recombining without processing yields original signal (approximately)."""
        sr = 44100
        # White noise
        np.random.seed(42)
        audio = np.random.normal(0, 0.5, sr)  # 1 second
        
        # Define bands
        bands = {
            'low': {'freq': 200},
            'mid': {'freq': 5000},
            'high': {}
        }
        
        # We need to expose the crossover logic to test it, 
        # or just call apply_multiband_compression with 0 compression (threshold 0dB or ratio 1:1)
        
        # Using a pass-through setting
        band_settings = [
            {'threshold_db': 0.0, 'ratio': 1.0}, # Low
            {'threshold_db': 0.0, 'ratio': 1.0}, # Mid
            {'threshold_db': 0.0, 'ratio': 1.0}  # High
        ]
        
        processed = self.engine.apply_multiband_compression(
            audio,
            band_settings=band_settings
        )
        
        # Check reconstruction error
        # Note: IIR filters have phase shift, so time-domain waveform won't match perfectly.
        # But magnitude spectrum should be very close.
        # Ideally we'd use linear phase FIR for perfect reconstruction, but that's slow in Python.
        # For this test, we accept some difference but check RMS energy preservation.
        
        rms_orig = np.sqrt(np.mean(audio**2))
        rms_proc = np.sqrt(np.mean(processed**2))
        
        print(f"Original RMS: {rms_orig:.4f}")
        print(f"Processed RMS: {rms_proc:.4f}")
        
        # Should be within 1% (Linkwitz-Riley sums to unity gain)
        self.assertAlmostEqual(rms_proc, rms_orig, delta=0.01)

    def test_low_band_compression(self):
        """Test that compression on low band affects low frequencies but not high."""
        sr = 44100
        t = np.linspace(0, 1.0, sr, endpoint=False)
        
        # Mix 100Hz (Low) and 10kHz (High)
        low_tone = np.sin(2 * np.pi * 100 * t) * 0.5
        high_tone = np.sin(2 * np.pi * 10000 * t) * 0.5
        mix = low_tone + high_tone
        
        # Compress ONLY Low band
        band_settings = [
            {'threshold_db': -20.0, 'ratio': 8.0, 'makeup_db': 0.0}, # Low: Heavily compressed
            {'threshold_db': 0.0, 'ratio': 1.0, 'makeup_db': 0.0},   # Mid: Pass
            {'threshold_db': 0.0, 'ratio': 1.0, 'makeup_db': 0.0}    # High: Pass
        ]
        
        processed = self.engine.apply_multiband_compression(
            mix,
            band_settings=band_settings
        )
        
        # Analyze spectrum roughly using FFT
        def get_freq_mag(sig, freq):
            fft = np.fft.rfft(sig)
            freqs = np.fft.rfftfreq(len(sig), 1/sr)
            idx = np.argmin(np.abs(freqs - freq))
            return np.abs(fft[idx])
        
        orig_low_mag = get_freq_mag(mix, 100)
        orig_high_mag = get_freq_mag(mix, 10000)
        
        proc_low_mag = get_freq_mag(processed, 100)
        proc_high_mag = get_freq_mag(processed, 10000)
        
        print(f"Low Mag: {orig_low_mag:.2f} -> {proc_low_mag:.2f}")
        print(f"High Mag: {orig_high_mag:.2f} -> {proc_high_mag:.2f}")
        
        # Low should be reduced
        self.assertLess(proc_low_mag, orig_low_mag * 0.8)
        
        # High should be preserved (mostly - some crossover leakage is expected but small)
        self.assertAlmostEqual(proc_high_mag, orig_high_mag, delta=orig_high_mag * 0.1)

if __name__ == '__main__':
    unittest.main()