# 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.
"""Signal processing utility module.
"""

import array
import logging
import os
import sys
import enum

try:
    import numpy as np
except ImportError:
    logging.critical('Cannot import the third-party Python package numpy')
    sys.exit(1)

try:
    import pydub
    import pydub.generators
except ImportError:
    logging.critical('Cannot import the third-party Python package pydub')
    sys.exit(1)

try:
    import scipy.signal
    import scipy.fftpack
except ImportError:
    logging.critical('Cannot import the third-party Python package scipy')
    sys.exit(1)

from . import exceptions


class SignalProcessingUtils(object):
    """Collection of signal processing utilities.
  """

    @enum.unique
    class MixPadding(enum.Enum):
        NO_PADDING = 0
        ZERO_PADDING = 1
        LOOP = 2

    def __init__(self):
        pass

    @classmethod
    def LoadWav(cls, filepath, channels=1):
        """Loads wav file.

    Args:
      filepath: path to the wav audio track file to load.
      channels: number of channels (downmixing to mono by default).

    Returns:
      AudioSegment instance.
    """
        if not os.path.exists(filepath):
            logging.error('cannot find the <%s> audio track file', filepath)
            raise exceptions.FileNotFoundError()
        return pydub.AudioSegment.from_file(filepath,
                                            format='wav',
                                            channels=channels)

    @classmethod
    def SaveWav(cls, output_filepath, signal):
        """Saves wav file.

    Args:
      output_filepath: path to the wav audio track file to save.
      signal: AudioSegment instance.
    """
        return signal.export(output_filepath, format='wav')

    @classmethod
    def CountSamples(cls, signal):
        """Number of samples per channel.

    Args:
      signal: AudioSegment instance.

    Returns:
      An integer.
    """
        number_of_samples = len(signal.get_array_of_samples())
        assert signal.channels > 0
        assert number_of_samples % signal.channels == 0
        return number_of_samples / signal.channels

    @classmethod
    def GenerateSilence(cls, duration=1000, sample_rate=48000):
        """Generates silence.

    This method can also be used to create a template AudioSegment instance.
    A template can then be used with other Generate*() methods accepting an
    AudioSegment instance as argument.

    Args:
      duration: duration in ms.
      sample_rate: sample rate.

    Returns:
      AudioSegment instance.
    """
        return pydub.AudioSegment.silent(duration, sample_rate)

    @classmethod
    def GeneratePureTone(cls, template, frequency=440.0):
        """Generates a pure tone.

    The pure tone is generated with the same duration and in the same format of
    the given template signal.

    Args:
      template: AudioSegment instance.
      frequency: Frequency of the pure tone in Hz.

    Return:
      AudioSegment instance.
    """
        if frequency > template.frame_rate >> 1:
            raise exceptions.SignalProcessingException('Invalid frequency')

        generator = pydub.generators.Sine(sample_rate=template.frame_rate,
                                          bit_depth=template.sample_width * 8,
                                          freq=frequency)

        return generator.to_audio_segment(duration=len(template), volume=0.0)

    @classmethod
    def GenerateWhiteNoise(cls, template):
        """Generates white noise.

    The white noise is generated with the same duration and in the same format
    of the given template signal.

    Args:
      template: AudioSegment instance.

    Return:
      AudioSegment instance.
    """
        generator = pydub.generators.WhiteNoise(
            sample_rate=template.frame_rate,
            bit_depth=template.sample_width * 8)
        return generator.to_audio_segment(duration=len(template), volume=0.0)

    @classmethod
    def AudioSegmentToRawData(cls, signal):
        samples = signal.get_array_of_samples()
        if samples.typecode != 'h':
            raise exceptions.SignalProcessingException(
                'Unsupported samples type')
        return np.array(signal.get_array_of_samples(), np.int16)

    @classmethod
    def Fft(cls, signal, normalize=True):
        if signal.channels != 1:
            raise NotImplementedError('multiple-channel FFT not implemented')
        x = cls.AudioSegmentToRawData(signal).astype(np.float32)
        if normalize:
            x /= max(abs(np.max(x)), 1.0)
        y = scipy.fftpack.fft(x)
        return y[:len(y) / 2]

    @classmethod
    def DetectHardClipping(cls, signal, threshold=2):
        """Detects hard clipping.

    Hard clipping is simply detected by counting samples that touch either the
    lower or upper bound too many times in a row (according to `threshold`).
    The presence of a single sequence of samples meeting such property is enough
    to label the signal as hard clipped.

    Args:
      signal: AudioSegment instance.
      threshold: minimum number of samples at full-scale in a row.

    Returns:
      True if hard clipping is detect, False otherwise.
    """
        if signal.channels != 1:
            raise NotImplementedError(
                'multiple-channel clipping not implemented')
        if signal.sample_width != 2:  # Note that signal.sample_width is in bytes.
            raise exceptions.SignalProcessingException(
                'hard-clipping detection only supported for 16 bit samples')
        samples = cls.AudioSegmentToRawData(signal)

        # Detect adjacent clipped samples.
        samples_type_info = np.iinfo(samples.dtype)
        mask_min = samples == samples_type_info.min
        mask_max = samples == samples_type_info.max

        def HasLongSequence(vector, min_legth=threshold):
            """Returns True if there are one or more long sequences of True flags."""
            seq_length = 0
            for b in vector:
                seq_length = seq_length + 1 if b else 0
                if seq_length >= min_legth:
                    return True
            return False

        return HasLongSequence(mask_min) or HasLongSequence(mask_max)

    @classmethod
    def ApplyImpulseResponse(cls, signal, impulse_response):
        """Applies an impulse response to a signal.

    Args:
      signal: AudioSegment instance.
      impulse_response: list or numpy vector of float values.

    Returns:
      AudioSegment instance.
    """
        # Get samples.
        assert signal.channels == 1, (
            'multiple-channel recordings not supported')
        samples = signal.get_array_of_samples()

        # Convolve.
        logging.info(
            'applying %d order impulse response to a signal lasting %d ms',
            len(impulse_response), len(signal))
        convolved_samples = scipy.signal.fftconvolve(in1=samples,
                                                     in2=impulse_response,
                                                     mode='full').astype(
                                                         np.int16)
        logging.info('convolution computed')

        # Cast.
        convolved_samples = array.array(signal.array_type, convolved_samples)

        # Verify.
        logging.debug('signal length: %d samples', len(samples))
        logging.debug('convolved signal length: %d samples',
                      len(convolved_samples))
        assert len(convolved_samples) > len(samples)

        # Generate convolved signal AudioSegment instance.
        convolved_signal = pydub.AudioSegment(data=convolved_samples,
                                              metadata={
                                                  'sample_width':
                                                  signal.sample_width,
                                                  'frame_rate':
                                                  signal.frame_rate,
                                                  'frame_width':
                                                  signal.frame_width,
                                                  'channels': signal.channels,
                                              })
        assert len(convolved_signal) > len(signal)

        return convolved_signal

    @classmethod
    def Normalize(cls, signal):
        """Normalizes a signal.

    Args:
      signal: AudioSegment instance.

    Returns:
      An AudioSegment instance.
    """
        return signal.apply_gain(-signal.max_dBFS)

    @classmethod
    def Copy(cls, signal):
        """Makes a copy os a signal.

    Args:
      signal: AudioSegment instance.

    Returns:
      An AudioSegment instance.
    """
        return pydub.AudioSegment(data=signal.get_array_of_samples(),
                                  metadata={
                                      'sample_width': signal.sample_width,
                                      'frame_rate': signal.frame_rate,
                                      'frame_width': signal.frame_width,
                                      'channels': signal.channels,
                                  })

    @classmethod
    def MixSignals(cls,
                   signal,
                   noise,
                   target_snr=0.0,
                   pad_noise=MixPadding.NO_PADDING):
        """Mixes `signal` and `noise` with a target SNR.

    Mix `signal` and `noise` with a desired SNR by scaling `noise`.
    If the target SNR is +/- infinite, a copy of signal/noise is returned.
    If `signal` is shorter than `noise`, the length of the mix equals that of
    `signal`. Otherwise, the mix length depends on whether padding is applied.
    When padding is not applied, that is `pad_noise` is set to NO_PADDING
    (default), the mix length equals that of `noise` - i.e., `signal` is
    truncated. Otherwise, `noise` is extended and the resulting mix has the same
    length of `signal`.

    Args:
      signal: AudioSegment instance (signal).
      noise: AudioSegment instance (noise).
      target_snr: float, numpy.Inf or -numpy.Inf (dB).
      pad_noise: SignalProcessingUtils.MixPadding, default: NO_PADDING.

    Returns:
      An AudioSegment instance.
    """
        # Handle infinite target SNR.
        if target_snr == -np.Inf:
            # Return a copy of noise.
            logging.warning('SNR = -Inf, returning noise')
            return cls.Copy(noise)
        elif target_snr == np.Inf:
            # Return a copy of signal.
            logging.warning('SNR = +Inf, returning signal')
            return cls.Copy(signal)

        # Check signal and noise power.
        signal_power = float(signal.dBFS)
        noise_power = float(noise.dBFS)
        if signal_power == -np.Inf:
            logging.error('signal has -Inf power, cannot mix')
            raise exceptions.SignalProcessingException(
                'cannot mix a signal with -Inf power')
        if noise_power == -np.Inf:
            logging.error('noise has -Inf power, cannot mix')
            raise exceptions.SignalProcessingException(
                'cannot mix a signal with -Inf power')

        # Mix.
        gain_db = signal_power - noise_power - target_snr
        signal_duration = len(signal)
        noise_duration = len(noise)
        if signal_duration <= noise_duration:
            # Ignore `pad_noise`, `noise` is truncated if longer that `signal`, the
            # mix will have the same length of `signal`.
            return signal.overlay(noise.apply_gain(gain_db))
        elif pad_noise == cls.MixPadding.NO_PADDING:
            # `signal` is longer than `noise`, but no padding is applied to `noise`.
            # Truncate `signal`.
            return noise.overlay(signal, gain_during_overlay=gain_db)
        elif pad_noise == cls.MixPadding.ZERO_PADDING:
            # TODO(alessiob): Check that this works as expected.
            return signal.overlay(noise.apply_gain(gain_db))
        elif pad_noise == cls.MixPadding.LOOP:
            # `signal` is longer than `noise`, extend `noise` by looping.
            return signal.overlay(noise.apply_gain(gain_db), loop=True)
        else:
            raise exceptions.SignalProcessingException('invalid padding type')