"""
test_new_dsp.py - Unit tests for new DSP methods

Tests for:
1. apply_stereo_imaging()
2. calculate_phase_correlation()
3. apply_harmonic_enhancement()
4. apply_dereverb()

Verifies correct behavior, output shapes, and audio processing integrity.
"""

import sys
import os

# Add src to path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))

import numpy as np
import pytest

# Import after path setup
from engine.processor import AudioEngine


class TestStereoImaging:
    """Test stereo imaging DSP method."""

    def setup_method(self):
        """Set up test fixtures."""
        self.engine = AudioEngine(temp_dir="temp")
        self.engine.sample_rate = 48000

        # Create 1 second stereo test signal
        duration = 1.0
        sr = self.engine.sample_rate
        t = np.linspace(0, duration, int(sr * duration))

        # Wide stereo content (different L/R)
        left = np.sin(2 * np.pi * 1000 * t).astype(np.float32)
        right = np.sin(2 * np.pi * 1000 * t + np.pi/4).astype(np.float32)  # 45° phase diff

        self.stereo_audio = np.column_stack((left, right))

    def test_ms_conversion_reversible(self):
        """M/S conversion with width=1.0 should return nearly identical audio."""
        # Width 1.0 for all bands should be transparent
        processed = self.engine.apply_stereo_imaging(
            self.stereo_audio,
            low_width=1.0,
            mid_width=1.0,
            high_width=1.0
        )

        # Should be very close to original
        diff_rms = np.sqrt(np.mean((processed - self.stereo_audio) ** 2))
        assert diff_rms < 0.01, f"Width=1.0 changed audio too much (RMS diff: {diff_rms})"

    def test_mono_bass_creation(self):
        """Low_width=0.0 should create mono below crossover."""
        # Create low frequency stereo content
        sr = self.engine.sample_rate
        duration = 1.0
        t = np.linspace(0, duration, int(sr * duration))

        # 100Hz with stereo difference
        left_low = np.sin(2 * np.pi * 100 * t).astype(np.float32) * 0.5
        right_low = np.sin(2 * np.pi * 100 * t + np.pi/2).astype(np.float32) * 0.5  # 90° phase diff

        low_freq_audio = np.column_stack((left_low, right_low))

        # Apply with mono bass (crossover at 200Hz)
        processed = self.engine.apply_stereo_imaging(
            low_freq_audio,
            low_width=0.0,
            mid_width=1.0,
            high_width=1.0,
            crossover_low=200.0,
            crossover_high=5000.0
        )

        # Low frequencies should be more mono (L/R more similar)
        original_diff = np.abs(low_freq_audio[:, 0] - low_freq_audio[:, 1]).mean()
        processed_diff = np.abs(processed[:, 0] - processed[:, 1]).mean()

        assert processed_diff < original_diff, \
            "Low_width=0.0 should make low frequencies more mono"

    def test_width_expansion(self):
        """High_width=1.5 should increase stereo width."""
        # Create high frequency content
        sr = self.engine.sample_rate
        duration = 1.0
        t = np.linspace(0, duration, int(sr * duration))

        # 10kHz with moderate stereo
        left_high = np.sin(2 * np.pi * 10000 * t).astype(np.float32) * 0.4
        right_high = np.sin(2 * np.pi * 10000 * t + np.pi/6).astype(np.float32) * 0.4

        high_freq_audio = np.column_stack((left_high, right_high))

        # Apply with expanded high width
        processed = self.engine.apply_stereo_imaging(
            high_freq_audio,
            low_width=1.0,
            mid_width=1.0,
            high_width=1.5,
            crossover_low=200.0,
            crossover_high=5000.0
        )

        # High frequencies should have greater stereo width (larger L/R difference)
        original_diff = np.abs(high_freq_audio[:, 0] - high_freq_audio[:, 1]).mean()
        processed_diff = np.abs(processed[:, 0] - processed[:, 1]).mean()

        assert processed_diff > original_diff, \
            "High_width=1.5 should increase stereo width"

    def test_output_shape_matches_input(self):
        """Output shape should match input shape."""
        processed = self.engine.apply_stereo_imaging(self.stereo_audio)

        assert processed.shape == self.stereo_audio.shape, \
            f"Output shape {processed.shape} doesn't match input {self.stereo_audio.shape}"

    def test_no_clipping(self):
        """Output values should stay within reasonable bounds (allow some headroom)."""
        # Use moderate input level to test clipping prevention
        moderate_audio = self.stereo_audio * 0.5

        processed = self.engine.apply_stereo_imaging(
            moderate_audio,
            low_width=0.0,
            mid_width=1.0,
            high_width=1.5
        )

        # Allow some overshoot but should be mostly contained
        # Width expansion can cause some peaks but shouldn't be excessive
        max_abs = np.max(np.abs(processed))
        assert max_abs < 2.0, \
            f"Audio clipped excessively: max={max_abs}"


class TestPhaseCorrelation:
    """Test phase correlation calculation."""

    def setup_method(self):
        """Set up test fixtures."""
        self.engine = AudioEngine(temp_dir="temp")
        self.engine.sample_rate = 48000

    def test_returns_expected_keys(self):
        """Should return dict with overall and band-specific keys."""
        sr = self.engine.sample_rate
        duration = 1.0
        samples = int(sr * duration)

        # Create stereo audio
        audio = np.random.randn(samples, 2).astype(np.float32) * 0.3

        result = self.engine.calculate_phase_correlation(audio)

        # Check for overall key
        assert "overall" in result, "Missing 'overall' key"

        # Check for band keys (default bands)
        # Band format: "band_{i+1}_{low}-{high}Hz"
        assert any("band_1" in key for key in result.keys()), "Missing low band key"
        assert any("band_2" in key for key in result.keys()), "Missing mid band key"
        assert any("band_3" in key for key in result.keys()), "Missing high band key"

    def test_values_in_valid_range(self):
        """All correlation values should be in [-1, 1] range."""
        sr = self.engine.sample_rate
        duration = 1.0
        samples = int(sr * duration)

        audio = np.random.randn(samples, 2).astype(np.float32) * 0.3

        result = self.engine.calculate_phase_correlation(audio)

        for key, value in result.items():
            assert -1.0 <= value <= 1.0, \
                f"Correlation value {key}={value} is out of range [-1, 1]"

    def test_mono_signal_returns_high_correlation(self):
        """Mono signal (identical L/R) should return correlation ~1.0."""
        sr = self.engine.sample_rate
        duration = 1.0
        t = np.linspace(0, duration, int(sr * duration))

        # Perfect mono - identical channels
        mono_signal = np.sin(2 * np.pi * 440 * t).astype(np.float32)
        stereo_mono = np.column_stack((mono_signal, mono_signal))

        result = self.engine.calculate_phase_correlation(stereo_mono)

        # Overall correlation should be very close to 1.0
        assert result["overall"] > 0.99, \
            f"Mono signal should have correlation ~1.0, got {result['overall']}"

    def test_out_of_phase_signal_returns_negative_correlation(self):
        """Out-of-phase signal should return negative correlation."""
        sr = self.engine.sample_rate
        duration = 1.0
        t = np.linspace(0, duration, int(sr * duration))

        # 180° out of phase
        left = np.sin(2 * np.pi * 440 * t).astype(np.float32)
        right = -left  # Inverted

        out_of_phase = np.column_stack((left, right))

        result = self.engine.calculate_phase_correlation(out_of_phase)

        # Overall correlation should be very close to -1.0
        assert result["overall"] < -0.99, \
            f"Out-of-phase signal should have correlation ~-1.0, got {result['overall']}"


class TestHarmonicEnhancement:
    """Test harmonic enhancement DSP method."""

    def setup_method(self):
        """Set up test fixtures."""
        self.engine = AudioEngine(temp_dir="temp")
        self.engine.sample_rate = 48000

    def test_generates_2nd_harmonics(self):
        """Warmth parameter should generate 2nd harmonics (2x frequency)."""
        sr = self.engine.sample_rate
        duration = 0.5
        t = np.linspace(0, duration, int(sr * duration))

        # Pure 1kHz tone
        fundamental_freq = 1000.0
        left = np.sin(2 * np.pi * fundamental_freq * t).astype(np.float32) * 0.3
        right = left.copy()
        audio = np.column_stack((left, right))

        # Apply warmth enhancement
        processed = self.engine.apply_harmonic_enhancement(
            audio,
            warmth_amount=0.5,
            edge_amount=0.0,
            frequency_range=(200.0, 5000.0)
        )

        # FFT to check for 2nd harmonic (2kHz)
        fft_original = np.fft.rfft(audio[:, 0])
        fft_processed = np.fft.rfft(processed[:, 0])
        freqs = np.fft.rfftfreq(len(audio), 1/sr)

        # Find magnitude at 2nd harmonic (2kHz)
        harmonic_2_idx = np.argmin(np.abs(freqs - fundamental_freq * 2))

        original_2nd = np.abs(fft_original[harmonic_2_idx])
        processed_2nd = np.abs(fft_processed[harmonic_2_idx])

        # Processed should have stronger 2nd harmonic
        assert processed_2nd > original_2nd, \
            "Warmth should increase 2nd harmonic content"

    def test_generates_3rd_harmonics(self):
        """Edge parameter should generate 3rd harmonics (3x frequency)."""
        sr = self.engine.sample_rate
        duration = 0.5
        t = np.linspace(0, duration, int(sr * duration))

        # Pure 1kHz tone
        fundamental_freq = 1000.0
        left = np.sin(2 * np.pi * fundamental_freq * t).astype(np.float32) * 0.3
        right = left.copy()
        audio = np.column_stack((left, right))

        # Apply edge enhancement
        processed = self.engine.apply_harmonic_enhancement(
            audio,
            warmth_amount=0.0,
            edge_amount=0.5,
            frequency_range=(200.0, 5000.0)
        )

        # FFT to check for 3rd harmonic (3kHz)
        fft_original = np.fft.rfft(audio[:, 0])
        fft_processed = np.fft.rfft(processed[:, 0])
        freqs = np.fft.rfftfreq(len(audio), 1/sr)

        # Find magnitude at 3rd harmonic (3kHz)
        harmonic_3_idx = np.argmin(np.abs(freqs - fundamental_freq * 3))

        original_3rd = np.abs(fft_original[harmonic_3_idx])
        processed_3rd = np.abs(fft_processed[harmonic_3_idx])

        # Processed should have stronger 3rd harmonic
        assert processed_3rd > original_3rd, \
            "Edge should increase 3rd harmonic content"

    def test_warmth_parameter_increases_2nd_harmonics(self):
        """Higher warmth amount should increase 2nd harmonic content more."""
        sr = self.engine.sample_rate
        duration = 0.5
        t = np.linspace(0, duration, int(sr * duration))

        fundamental_freq = 1000.0
        audio = np.column_stack([
            np.sin(2 * np.pi * fundamental_freq * t).astype(np.float32) * 0.3,
            np.sin(2 * np.pi * fundamental_freq * t).astype(np.float32) * 0.3
        ])

        # Low warmth
        processed_low = self.engine.apply_harmonic_enhancement(
            audio.copy(), warmth_amount=0.1, edge_amount=0.0
        )

        # High warmth
        processed_high = self.engine.apply_harmonic_enhancement(
            audio.copy(), warmth_amount=0.8, edge_amount=0.0
        )

        # FFT analysis
        fft_low = np.fft.rfft(processed_low[:, 0])
        fft_high = np.fft.rfft(processed_high[:, 0])
        freqs = np.fft.rfftfreq(len(audio), 1/sr)

        harmonic_2_idx = np.argmin(np.abs(freqs - fundamental_freq * 2))

        magnitude_low = np.abs(fft_low[harmonic_2_idx])
        magnitude_high = np.abs(fft_high[harmonic_2_idx])

        assert magnitude_high > magnitude_low, \
            "Higher warmth should produce more 2nd harmonic content"

    def test_edge_parameter_increases_3rd_harmonics(self):
        """Higher edge amount should increase 3rd harmonic content more."""
        sr = self.engine.sample_rate
        duration = 0.5
        t = np.linspace(0, duration, int(sr * duration))

        fundamental_freq = 1000.0
        audio = np.column_stack([
            np.sin(2 * np.pi * fundamental_freq * t).astype(np.float32) * 0.3,
            np.sin(2 * np.pi * fundamental_freq * t).astype(np.float32) * 0.3
        ])

        # Low edge
        processed_low = self.engine.apply_harmonic_enhancement(
            audio.copy(), warmth_amount=0.0, edge_amount=0.1
        )

        # High edge
        processed_high = self.engine.apply_harmonic_enhancement(
            audio.copy(), warmth_amount=0.0, edge_amount=0.8
        )

        # FFT analysis
        fft_low = np.fft.rfft(processed_low[:, 0])
        fft_high = np.fft.rfft(processed_high[:, 0])
        freqs = np.fft.rfftfreq(len(audio), 1/sr)

        harmonic_3_idx = np.argmin(np.abs(freqs - fundamental_freq * 3))

        magnitude_low = np.abs(fft_low[harmonic_3_idx])
        magnitude_high = np.abs(fft_high[harmonic_3_idx])

        assert magnitude_high > magnitude_low, \
            "Higher edge should produce more 3rd harmonic content"

    def test_output_doesnt_clip(self):
        """Enhanced audio should apply soft clipping to stay within bounds."""
        sr = self.engine.sample_rate
        duration = 1.0
        samples = int(sr * duration)

        # Moderate-level input (harmonic enhancement adds energy)
        audio = np.random.randn(samples, 2).astype(np.float32) * 0.5

        processed = self.engine.apply_harmonic_enhancement(
            audio,
            warmth_amount=0.8,
            edge_amount=0.8
        )

        # The tanh soft clipper should keep output bounded
        # tanh(x * 0.9) / 0.9 constrains output to approximately [-1.11, 1.11]
        assert np.all(np.abs(processed) <= 1.2), \
            f"Audio clipped excessively: max={np.abs(processed).max()}"


class TestDereverb:
    """Test de-reverb DSP method."""

    def setup_method(self):
        """Set up test fixtures."""
        self.engine = AudioEngine(temp_dir="temp")
        self.engine.sample_rate = 48000

    def test_preserves_transients(self):
        """Dereverb should preserve loud regions better than quiet regions."""
        sr = self.engine.sample_rate
        duration = 1.0
        samples = int(sr * duration)

        # Create audio with loud and quiet sections
        audio = np.zeros((samples, 2), dtype=np.float32)

        # Loud section (above threshold)
        loud_start = 0
        loud_end = int(sr * 0.3)
        audio[loud_start:loud_end, :] = 0.5  # Loud, above threshold

        # Quiet section (below threshold) - this is what should be reduced
        quiet_start = int(sr * 0.5)
        quiet_end = int(sr * 0.8)
        audio[quiet_start:quiet_end, :] = 0.01  # Very quiet, below threshold

        # Apply dereverb
        processed = self.engine.apply_dereverb(
            audio,
            reduction_amount=0.8,
            gate_threshold_db=-30.0  # ~0.032 linear
        )

        # Calculate RMS for each region
        loud_region = slice(loud_start + 1000, loud_end - 1000)  # Avoid edges
        quiet_region = slice(quiet_start + 1000, quiet_end - 1000)

        original_loud_rms = np.sqrt(np.mean(audio[loud_region] ** 2))
        processed_loud_rms = np.sqrt(np.mean(processed[loud_region] ** 2))

        original_quiet_rms = np.sqrt(np.mean(audio[quiet_region] ** 2))
        processed_quiet_rms = np.sqrt(np.mean(processed[quiet_region] ** 2))

        # Calculate retention ratios
        loud_retention = processed_loud_rms / (original_loud_rms + 1e-9)
        quiet_retention = processed_quiet_rms / (original_quiet_rms + 1e-9)

        # Loud regions should be retained better than quiet regions
        assert loud_retention > quiet_retention, \
            f"Loud regions not preserved better: loud={loud_retention:.3f}, quiet={quiet_retention:.3f}"

    def test_reduces_reverb_tails(self):
        """RMS should decrease in reverb tail regions."""
        sr = self.engine.sample_rate
        duration = 1.0
        samples = int(sr * duration)

        # Create signal with reverb tail
        audio = np.zeros((samples, 2), dtype=np.float32)

        # Initial transient
        audio[0:100, :] = 0.8

        # Long reverb tail (quiet sustained signal)
        tail_start = 1000
        tail_length = int(sr * 0.5)
        audio[tail_start:tail_start+tail_length, :] = 0.1

        # Apply dereverb
        processed = self.engine.apply_dereverb(
            audio,
            reduction_amount=0.8,
            gate_threshold_db=-30.0
        )

        # Calculate RMS in tail region
        tail_region = slice(tail_start, tail_start + tail_length)
        original_rms = np.sqrt(np.mean(audio[tail_region] ** 2))
        processed_rms = np.sqrt(np.mean(processed[tail_region] ** 2))

        assert processed_rms < original_rms, \
            f"Reverb tail not reduced: original RMS={original_rms}, processed={processed_rms}"

    def test_gate_threshold_works(self):
        """Gate threshold should control what gets reduced."""
        sr = self.engine.sample_rate
        duration = 1.0
        samples = int(sr * duration)

        # Create quiet sustained signal
        audio = np.ones((samples, 2), dtype=np.float32) * 0.05  # -26 dB

        # Apply with high threshold (should reduce)
        processed_high = self.engine.apply_dereverb(
            audio.copy(),
            reduction_amount=0.8,
            gate_threshold_db=-20.0  # Above signal level
        )

        # Apply with low threshold (should not reduce much)
        processed_low = self.engine.apply_dereverb(
            audio.copy(),
            reduction_amount=0.8,
            gate_threshold_db=-60.0  # Below signal level
        )

        rms_high = np.sqrt(np.mean(processed_high ** 2))
        rms_low = np.sqrt(np.mean(processed_low ** 2))

        # High threshold should reduce more
        assert rms_high < rms_low, \
            "Gate threshold not working correctly"

    def test_reduction_amount_scales_properly(self):
        """Higher reduction amount should reduce more."""
        sr = self.engine.sample_rate
        duration = 1.0
        samples = int(sr * duration)

        # Create quiet sustained signal (simulating reverb tail)
        # Must be below threshold to trigger gating
        audio = np.ones((samples, 2), dtype=np.float32) * 0.02  # ~-34 dB

        # Use a threshold that will catch this signal
        gate_threshold = -25.0  # dB

        # Low reduction
        processed_low = self.engine.apply_dereverb(
            audio.copy(),
            reduction_amount=0.3,
            gate_threshold_db=gate_threshold
        )

        # High reduction
        processed_high = self.engine.apply_dereverb(
            audio.copy(),
            reduction_amount=0.9,
            gate_threshold_db=gate_threshold
        )

        rms_original = np.sqrt(np.mean(audio ** 2))
        rms_low = np.sqrt(np.mean(processed_low ** 2))
        rms_high = np.sqrt(np.mean(processed_high ** 2))

        # Calculate reduction percentages
        reduction_low = (rms_original - rms_low) / rms_original
        reduction_high = (rms_original - rms_high) / rms_original

        # Both should reduce, and high reduction should reduce more
        assert reduction_low > 0.1, f"Low reduction didn't reduce much: {reduction_low:.3f}"
        assert reduction_high > reduction_low, \
            f"High reduction ({reduction_high:.3f}) didn't reduce more than low ({reduction_low:.3f})"


def run_tests():
    """Run all tests and print results."""
    print("=" * 70)
    print("AutoRemaster New DSP Unit Tests")
    print("=" * 70)

    # Collect test classes
    test_classes = [
        TestStereoImaging,
        TestPhaseCorrelation,
        TestHarmonicEnhancement,
        TestDereverb,
    ]

    passed = 0
    failed = 0

    for test_class in test_classes:
        print(f"\n{test_class.__name__}:")

        instance = test_class()
        for method_name in dir(instance):
            if method_name.startswith('test_'):
                try:
                    if hasattr(instance, 'setup_method'):
                        instance.setup_method()

                    getattr(instance, method_name)()
                    print(f"  ✓ {method_name}")
                    passed += 1
                except AssertionError as e:
                    print(f"  ✗ {method_name}: {e}")
                    failed += 1
                except Exception as e:
                    print(f"  ✗ {method_name}: Exception - {type(e).__name__}: {e}")
                    failed += 1

    print("\n" + "=" * 70)
    print(f"Results: {passed} passed, {failed} failed")
    print("=" * 70)

    return failed == 0


if __name__ == "__main__":
    success = run_tests()
    sys.exit(0 if success else 1)