# Copyright (c) 2017 The WebRTC project authors. All Rights Reserved.
#
# Use of this source code is governed by a BSD-style license
# that can be found in the LICENSE file in the root of the source
# tree. An additional intellectual property rights grant can be found
# in the file PATENTS.  All contributing project authors may
# be found in the AUTHORS file in the root of the source tree.
"""Unit tests for the signal_processing module.
"""

import unittest

import numpy as np
import pydub

from . import exceptions
from . import signal_processing


class TestSignalProcessing(unittest.TestCase):
    """Unit tests for the signal_processing module.
  """

    def testMixSignals(self):
        # Generate a template signal with which white noise can be generated.
        silence = pydub.AudioSegment.silent(duration=1000, frame_rate=48000)

        # Generate two distinct AudioSegment instances with 1 second of white noise.
        signal = signal_processing.SignalProcessingUtils.GenerateWhiteNoise(
            silence)
        noise = signal_processing.SignalProcessingUtils.GenerateWhiteNoise(
            silence)

        # Extract samples.
        signal_samples = signal.get_array_of_samples()
        noise_samples = noise.get_array_of_samples()

        # Test target SNR -Inf (noise expected).
        mix_neg_inf = signal_processing.SignalProcessingUtils.MixSignals(
            signal, noise, -np.Inf)
        self.assertTrue(len(noise), len(mix_neg_inf))  # Check duration.
        mix_neg_inf_samples = mix_neg_inf.get_array_of_samples()
        self.assertTrue(  # Check samples.
            all([x == y for x, y in zip(noise_samples, mix_neg_inf_samples)]))

        # Test target SNR 0.0 (different data expected).
        mix_0 = signal_processing.SignalProcessingUtils.MixSignals(
            signal, noise, 0.0)
        self.assertTrue(len(signal), len(mix_0))  # Check duration.
        self.assertTrue(len(noise), len(mix_0))
        mix_0_samples = mix_0.get_array_of_samples()
        self.assertTrue(
            any([x != y for x, y in zip(signal_samples, mix_0_samples)]))
        self.assertTrue(
            any([x != y for x, y in zip(noise_samples, mix_0_samples)]))

        # Test target SNR +Inf (signal expected).
        mix_pos_inf = signal_processing.SignalProcessingUtils.MixSignals(
            signal, noise, np.Inf)
        self.assertTrue(len(signal), len(mix_pos_inf))  # Check duration.
        mix_pos_inf_samples = mix_pos_inf.get_array_of_samples()
        self.assertTrue(  # Check samples.
            all([x == y for x, y in zip(signal_samples, mix_pos_inf_samples)]))

    def testMixSignalsMinInfPower(self):
        silence = pydub.AudioSegment.silent(duration=1000, frame_rate=48000)
        signal = signal_processing.SignalProcessingUtils.GenerateWhiteNoise(
            silence)

        with self.assertRaises(exceptions.SignalProcessingException):
            _ = signal_processing.SignalProcessingUtils.MixSignals(
                signal, silence, 0.0)

        with self.assertRaises(exceptions.SignalProcessingException):
            _ = signal_processing.SignalProcessingUtils.MixSignals(
                silence, signal, 0.0)

    def testMixSignalNoiseDifferentLengths(self):
        # Test signals.
        shorter = signal_processing.SignalProcessingUtils.GenerateWhiteNoise(
            pydub.AudioSegment.silent(duration=1000, frame_rate=8000))
        longer = signal_processing.SignalProcessingUtils.GenerateWhiteNoise(
            pydub.AudioSegment.silent(duration=2000, frame_rate=8000))

        # When the signal is shorter than the noise, the mix length always equals
        # that of the signal regardless of whether padding is applied.
        # No noise padding, length of signal less than that of noise.
        mix = signal_processing.SignalProcessingUtils.MixSignals(
            signal=shorter,
            noise=longer,
            pad_noise=signal_processing.SignalProcessingUtils.MixPadding.
            NO_PADDING)
        self.assertEqual(len(shorter), len(mix))
        # With noise padding, length of signal less than that of noise.
        mix = signal_processing.SignalProcessingUtils.MixSignals(
            signal=shorter,
            noise=longer,
            pad_noise=signal_processing.SignalProcessingUtils.MixPadding.
            ZERO_PADDING)
        self.assertEqual(len(shorter), len(mix))

        # When the signal is longer than the noise, the mix length depends on
        # whether padding is applied.
        # No noise padding, length of signal greater than that of noise.
        mix = signal_processing.SignalProcessingUtils.MixSignals(
            signal=longer,
            noise=shorter,
            pad_noise=signal_processing.SignalProcessingUtils.MixPadding.
            NO_PADDING)
        self.assertEqual(len(shorter), len(mix))
        # With noise padding, length of signal greater than that of noise.
        mix = signal_processing.SignalProcessingUtils.MixSignals(
            signal=longer,
            noise=shorter,
            pad_noise=signal_processing.SignalProcessingUtils.MixPadding.
            ZERO_PADDING)
        self.assertEqual(len(longer), len(mix))

    def testMixSignalNoisePaddingTypes(self):
        # Test signals.
        shorter = signal_processing.SignalProcessingUtils.GenerateWhiteNoise(
            pydub.AudioSegment.silent(duration=1000, frame_rate=8000))
        longer = signal_processing.SignalProcessingUtils.GeneratePureTone(
            pydub.AudioSegment.silent(duration=2000, frame_rate=8000), 440.0)

        # Zero padding: expect pure tone only in 1-2s.
        mix_zero_pad = signal_processing.SignalProcessingUtils.MixSignals(
            signal=longer,
            noise=shorter,
            target_snr=-6,
            pad_noise=signal_processing.SignalProcessingUtils.MixPadding.
            ZERO_PADDING)

        # Loop: expect pure tone plus noise in 1-2s.
        mix_loop = signal_processing.SignalProcessingUtils.MixSignals(
            signal=longer,
            noise=shorter,
            target_snr=-6,
            pad_noise=signal_processing.SignalProcessingUtils.MixPadding.LOOP)

        def Energy(signal):
            samples = signal_processing.SignalProcessingUtils.AudioSegmentToRawData(
                signal).astype(np.float32)
            return np.sum(samples * samples)

        e_mix_zero_pad = Energy(mix_zero_pad[-1000:])
        e_mix_loop = Energy(mix_loop[-1000:])
        self.assertLess(0, e_mix_zero_pad)
        self.assertLess(e_mix_zero_pad, e_mix_loop)

    def testMixSignalSnr(self):
        # Test signals.
        tone_low = signal_processing.SignalProcessingUtils.GeneratePureTone(
            pydub.AudioSegment.silent(duration=64, frame_rate=8000), 250.0)
        tone_high = signal_processing.SignalProcessingUtils.GeneratePureTone(
            pydub.AudioSegment.silent(duration=64, frame_rate=8000), 3000.0)

        def ToneAmplitudes(mix):
            """Returns the amplitude of the coefficients #16 and #192, which
         correspond to the tones at 250 and 3k Hz respectively."""
            mix_fft = np.absolute(
                signal_processing.SignalProcessingUtils.Fft(mix))
            return mix_fft[16], mix_fft[192]

        mix = signal_processing.SignalProcessingUtils.MixSignals(
            signal=tone_low, noise=tone_high, target_snr=-6)
        ampl_low, ampl_high = ToneAmplitudes(mix)
        self.assertLess(ampl_low, ampl_high)

        mix = signal_processing.SignalProcessingUtils.MixSignals(
            signal=tone_high, noise=tone_low, target_snr=-6)
        ampl_low, ampl_high = ToneAmplitudes(mix)
        self.assertLess(ampl_high, ampl_low)

        mix = signal_processing.SignalProcessingUtils.MixSignals(
            signal=tone_low, noise=tone_high, target_snr=6)
        ampl_low, ampl_high = ToneAmplitudes(mix)
        self.assertLess(ampl_high, ampl_low)

        mix = signal_processing.SignalProcessingUtils.MixSignals(
            signal=tone_high, noise=tone_low, target_snr=6)
        ampl_low, ampl_high = ToneAmplitudes(mix)
        self.assertLess(ampl_low, ampl_high)