"""
Tests for add_noise module.
"""

import os
import tempfile
import numpy as np
import soundfile as sf
import pytest

from amt_augmentor.add_noise import apply_noise


class TestAddNoise:
    """Test noise addition functionality."""

    def setup_method(self):
        """Set up test files."""
        self.test_dir = tempfile.mkdtemp()

        # Create test audio file
        self.audio_file = os.path.join(self.test_dir, "test.wav")
        self.sr = 22050
        duration = 1.0
        t = np.linspace(0, duration, int(self.sr * duration), endpoint=False)
        self.original_audio = 0.5 * np.sin(2 * np.pi * 440 * t)
        sf.write(self.audio_file, self.original_audio, self.sr)

        # Create test annotation file
        self.ann_file = os.path.join(self.test_dir, "test.ann")
        self.ann_content = "0.1\t0.5\t60\t100\n0.6\t1.0\t62\t100\n1.1\t1.5\t64\t100\n"
        with open(self.ann_file, 'w', encoding='utf-8') as f:
            f.write(self.ann_content)

    def teardown_method(self):
        """Clean up test files."""
        import shutil
        if os.path.exists(self.test_dir):
            shutil.rmtree(self.test_dir)

    def test_apply_noise_basic(self):
        """Test basic noise addition."""
        output_file = os.path.join(self.test_dir, "noisy.wav")
        intensity = 0.1

        # Apply noise
        output_ann = apply_noise(
            self.audio_file,
            self.ann_file,
            output_file,
            intensity
        )

        # Check output files exist
        assert os.path.exists(output_file)
        assert os.path.exists(output_ann)
        assert output_ann.endswith(".ann")

        # Check output audio was created and is different from original
        noisy_audio, sr = sf.read(output_file)
        assert sr == self.sr
        assert len(noisy_audio) == len(self.original_audio)
        # Audio should be different after adding noise
        assert not np.allclose(noisy_audio, self.original_audio)

        # Check annotation file is unchanged
        with open(output_ann, 'r', encoding='utf-8') as f:
            content = f.read()
        assert content == self.ann_content

    def test_apply_noise_intensity_levels(self):
        """Test different noise intensity levels."""
        intensities = [0.01, 0.1, 0.5, 1.0]
        outputs = []

        for i, intensity in enumerate(intensities):
            output_file = os.path.join(self.test_dir, f"noisy_{i}.wav")
            apply_noise(self.audio_file, self.ann_file, output_file, intensity)

            noisy_audio, _ = sf.read(output_file)
            outputs.append(noisy_audio)

            # Check that output is normalized (max absolute value <= 1)
            assert np.max(np.abs(noisy_audio)) <= 1.0

        # Higher intensity should generally produce more different results
        # Compare RMS difference from original
        original_audio, _ = sf.read(self.audio_file)
        diffs = [np.sqrt(np.mean((out - original_audio) ** 2)) for out in outputs]
        # Generally, higher intensity means more difference (with some tolerance for randomness)
        assert diffs[3] > diffs[0] * 0.5  # Rough check due to random nature

    def test_apply_noise_flac_format(self):
        """Test noise addition with FLAC output format."""
        # Create FLAC input file
        flac_file = os.path.join(self.test_dir, "test.flac")
        sf.write(flac_file, self.original_audio, self.sr, format='FLAC')

        output_file = os.path.join(self.test_dir, "noisy.flac")
        intensity = 0.2

        # Apply noise
        output_ann = apply_noise(
            flac_file,
            self.ann_file,
            output_file,
            intensity
        )

        # Check output files exist
        assert os.path.exists(output_file)
        assert os.path.exists(output_ann)

        # Verify it's a FLAC file
        info = sf.info(output_file)
        assert info.format == 'FLAC'

        # Check audio was modified
        noisy_audio, sr = sf.read(output_file)
        assert sr == self.sr
        assert not np.allclose(noisy_audio, self.original_audio)

    def test_apply_noise_zero_intensity(self):
        """Test noise addition with zero intensity (should be minimal change)."""
        output_file = os.path.join(self.test_dir, "zero_noise.wav")
        intensity = 0.0

        # Apply noise with zero intensity
        apply_noise(self.audio_file, self.ann_file, output_file, intensity)

        # Audio should be very similar to original (just normalized)
        noisy_audio, _ = sf.read(output_file)
        original_normalized = self.original_audio / np.max(np.abs(self.original_audio))
        # With intensity=0, should be very close to normalized original
        # Use atol instead of rtol for small values
        assert np.allclose(noisy_audio, original_normalized, atol=1e-4)

    def test_apply_noise_high_intensity(self):
        """Test noise addition with very high intensity."""
        output_file = os.path.join(self.test_dir, "high_noise.wav")
        intensity = 10.0

        # Apply high intensity noise
        output_ann = apply_noise(
            self.audio_file,
            self.ann_file,
            output_file,
            intensity
        )

        # Check output files exist
        assert os.path.exists(output_file)
        assert os.path.exists(output_ann)

        # Output should still be normalized
        noisy_audio, _ = sf.read(output_file)
        assert np.max(np.abs(noisy_audio)) <= 1.0

        # Should be very different from original due to high noise
        correlation = np.corrcoef(noisy_audio, self.original_audio)[0, 1]
        assert abs(correlation) < 0.5  # Low correlation due to high noise

    def test_mono_audio_handling(self):
        """Test noise addition with mono audio."""
        output_file = os.path.join(self.test_dir, "mono_noisy.wav")
        result = apply_noise(self.audio_file, self.ann_file, output_file, 0.15)

        assert os.path.exists(output_file)
        assert os.path.exists(result)

        # Verify mono is preserved
        output_audio, _ = sf.read(output_file)
        assert output_audio.ndim == 1  # Should be mono

    def test_stereo_audio_handling(self):
        """Test noise addition with stereo audio."""
        # Create stereo audio file
        stereo_file = os.path.join(self.test_dir, "stereo.wav")
        t = np.linspace(0, 1, self.sr)
        left = 0.5 * np.sin(2 * np.pi * 440 * t)
        right = 0.5 * np.sin(2 * np.pi * 550 * t)
        stereo = np.column_stack((left, right))
        sf.write(stereo_file, stereo, self.sr)

        output_file = os.path.join(self.test_dir, "stereo_noisy.wav")
        result = apply_noise(stereo_file, self.ann_file, output_file, 0.2)

        assert os.path.exists(output_file)
        assert os.path.exists(result)

        # Note: The current implementation may flatten stereo to mono
        # due to the way noise is generated (samples.size instead of samples.shape)

    def test_annotation_preservation(self):
        """Test that annotation files are preserved exactly."""
        # Create annotation with various formats
        complex_ann = os.path.join(self.test_dir, "complex.ann")
        complex_content = "0.1\t0.5\t60\t100\n# Comment line\n0.6\t1.0\t62\t100\n\n1.1\t1.5\t64\t100\n"
        with open(complex_ann, 'w', encoding='utf-8') as f:
            f.write(complex_content)

        output_file = os.path.join(self.test_dir, "output.wav")
        output_ann = apply_noise(self.audio_file, complex_ann, output_file, 0.1)

        # Check annotation is preserved exactly
        with open(output_ann, 'r', encoding='utf-8') as f:
            output_content = f.read()
        assert output_content == complex_content

    def test_reproducibility_with_seed(self):
        """Test that noise is random (different each time)."""
        output1 = os.path.join(self.test_dir, "noisy1.wav")
        output2 = os.path.join(self.test_dir, "noisy2.wav")

        intensity = 0.3
        apply_noise(self.audio_file, self.ann_file, output1, intensity)
        apply_noise(self.audio_file, self.ann_file, output2, intensity)

        audio1, _ = sf.read(output1)
        audio2, _ = sf.read(output2)

        # The two outputs should be different due to random noise
        assert not np.allclose(audio1, audio2)

    def test_empty_annotation_file(self):
        """Test handling of empty annotation file."""
        empty_ann = os.path.join(self.test_dir, "empty.ann")
        with open(empty_ann, 'w', encoding='utf-8') as f:
            f.write("")

        output_file = os.path.join(self.test_dir, "output.wav")
        output_ann = apply_noise(self.audio_file, empty_ann, output_file, 0.1)

        assert os.path.exists(output_file)
        assert os.path.exists(output_ann)

        # Empty annotation should remain empty
        with open(output_ann, 'r', encoding='utf-8') as f:
            content = f.read()
        assert content == ""