# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for spectral_ops."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl.testing import parameterized
import numpy as np

from tensorflow.python.compat import compat
from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradient_checker_v2
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.signal import spectral_ops
from tensorflow.python.ops.signal import window_ops
from tensorflow.python.platform import test


@test_util.run_all_in_graph_and_eager_modes
class SpectralOpsTest(test.TestCase, parameterized.TestCase):

  @staticmethod
  def _np_hann_periodic_window(length):
    if length == 1:
      return np.ones(1)
    odd = length % 2
    if not odd:
      length += 1
    window = 0.5 - 0.5 * np.cos(2.0 * np.pi * np.arange(length) / (length - 1))
    if not odd:
      window = window[:-1]
    return window

  @staticmethod
  def _np_frame(data, window_length, hop_length):
    num_frames = 1 + int(np.floor((len(data) - window_length) // hop_length))
    shape = (num_frames, window_length)
    strides = (data.strides[0] * hop_length, data.strides[0])
    return np.lib.stride_tricks.as_strided(data, shape=shape, strides=strides)

  @staticmethod
  def _np_stft(data, fft_length, hop_length, window_length):
    frames = SpectralOpsTest._np_frame(data, window_length, hop_length)
    window = SpectralOpsTest._np_hann_periodic_window(window_length)
    return np.fft.rfft(frames * window, fft_length)

  @staticmethod
  def _np_inverse_stft(stft, fft_length, hop_length, window_length):
    frames = np.fft.irfft(stft, fft_length)
    # Pad or truncate frames's inner dimension to window_length.
    frames = frames[..., :window_length]
    frames = np.pad(frames, [[0, 0]] * (frames.ndim - 1) +
                    [[0, max(0, window_length - frames.shape[-1])]], "constant")
    window = SpectralOpsTest._np_hann_periodic_window(window_length)
    return SpectralOpsTest._np_overlap_add(frames * window, hop_length)

  @staticmethod
  def _np_overlap_add(stft, hop_length):
    num_frames, window_length = np.shape(stft)
    # Output length will be one complete window, plus another hop_length's
    # worth of points for each additional window.
    output_length = window_length + (num_frames - 1) * hop_length
    output = np.zeros(output_length)
    for i in range(num_frames):
      output[i * hop_length:i * hop_length + window_length] += stft[i,]
    return output

  def _compare(self, signal, frame_length, frame_step, fft_length, tol):
    actual_stft = spectral_ops.stft(
        signal, frame_length, frame_step, fft_length, pad_end=False)
    signal_ph = array_ops.placeholder_with_default(signal, shape=signal.shape)
    actual_stft_from_ph = spectral_ops.stft(
        signal_ph, frame_length, frame_step, fft_length, pad_end=False)

    actual_inverse_stft = spectral_ops.inverse_stft(
        actual_stft, frame_length, frame_step, fft_length)

    actual_stft, actual_stft_from_ph, actual_inverse_stft = self.evaluate(
        [actual_stft, actual_stft_from_ph, actual_inverse_stft])

    actual_stft_ph = array_ops.placeholder_with_default(
        actual_stft, shape=actual_stft.shape)
    actual_inverse_stft_from_ph = self.evaluate(
        spectral_ops.inverse_stft(
            actual_stft_ph, frame_length, frame_step, fft_length))

    # Confirm that there is no difference in output when shape/rank is fully
    # unknown or known.
    self.assertAllClose(actual_stft, actual_stft_from_ph)
    self.assertAllClose(actual_inverse_stft, actual_inverse_stft_from_ph)

    expected_stft = SpectralOpsTest._np_stft(
        signal, fft_length, frame_step, frame_length)
    self.assertAllClose(expected_stft, actual_stft, rtol=tol, atol=tol)

    expected_inverse_stft = SpectralOpsTest._np_inverse_stft(
        expected_stft, fft_length, frame_step, frame_length)
    self.assertAllClose(
        expected_inverse_stft, actual_inverse_stft, rtol=tol, atol=tol)

  def test_shapes(self):
    signal = np.zeros((512,)).astype(np.float32)

    # If fft_length is not provided, the smallest enclosing power of 2 of
    # frame_length (8) is used.
    stft = spectral_ops.stft(signal, frame_length=7, frame_step=8,
                             pad_end=True)
    self.assertAllEqual([64, 5], stft.shape.as_list())
    self.assertAllEqual([64, 5], self.evaluate(stft).shape)

    stft = spectral_ops.stft(signal, frame_length=8, frame_step=8,
                             pad_end=True)
    self.assertAllEqual([64, 5], stft.shape.as_list())
    self.assertAllEqual([64, 5], self.evaluate(stft).shape)

    stft = spectral_ops.stft(signal, frame_length=8, frame_step=8,
                             fft_length=16, pad_end=True)
    self.assertAllEqual([64, 9], stft.shape.as_list())
    self.assertAllEqual([64, 9], self.evaluate(stft).shape)

    stft = spectral_ops.stft(signal, frame_length=16, frame_step=8,
                             fft_length=8, pad_end=True)
    self.assertAllEqual([64, 5], stft.shape.as_list())
    self.assertAllEqual([64, 5], self.evaluate(stft).shape)

    stft = np.zeros((32, 9)).astype(np.complex64)

    inverse_stft = spectral_ops.inverse_stft(stft, frame_length=8,
                                             fft_length=16, frame_step=8)
    expected_length = (stft.shape[0] - 1) * 8 + 8
    self.assertAllEqual([256], inverse_stft.shape.as_list())
    self.assertAllEqual([expected_length], self.evaluate(inverse_stft).shape)

  @parameterized.parameters(
      (512, 64, 32, 64, np.float32, 1e-4),
      (512, 64, 32, 64, np.float64, 1e-8),
      (512, 64, 64, 64, np.float32, 1e-4),
      (512, 64, 64, 64, np.float64, 1e-8),
      (512, 72, 64, 64, np.float32, 1e-4),
      (512, 72, 64, 64, np.float64, 1e-8),
      (512, 64, 25, 64, np.float32, 1e-4),
      (512, 64, 25, 64, np.float64, 1e-8),
      (512, 25, 15, 36, np.float32, 1e-4),
      (512, 25, 15, 36, np.float64, 1e-8),
      (123, 23, 5, 42, np.float32, 1e-4),
      (123, 23, 5, 42, np.float64, 1e-8))
  def test_stft_and_inverse_stft(self, signal_length, frame_length,
                                 frame_step, fft_length, np_rtype, tol):
    """Test that spectral_ops.stft/inverse_stft match a NumPy implementation."""
    # Enable float64 support for RFFTs.
    with compat.forward_compatibility_horizon(2019, 10, 13):
      signal = np.random.random(signal_length).astype(np_rtype)
      self._compare(signal, frame_length, frame_step, fft_length, tol)

  @parameterized.parameters(
      # 87.5% overlap.
      (4096, 256, 32, 256, np.float32, 1e-5, 1e-6),
      (4096, 256, 32, 256, np.float64, 1e-8, 1e-8),
      # 75% overlap.
      (4096, 256, 64, 256, np.float32, 1e-5, 1e-6),
      (4096, 256, 64, 256, np.float64, 1e-8, 1e-8),
      # Odd frame hop.
      (4096, 128, 25, 128, np.float32, 1e-3, 1e-6),
      (4096, 128, 25, 128, np.float64, 5e-4, 1e-8),
      # Odd frame length.
      (4096, 127, 32, 128, np.float32, 1e-3, 1e-6),
      (4096, 127, 32, 128, np.float64, 1e-3, 1e-8),
      # 50% overlap.
      (4096, 128, 64, 128, np.float32, 0.4, 1e-6),
      (4096, 128, 64, 128, np.float64, 0.4, 1e-8))
  def test_stft_round_trip(self, signal_length, frame_length, frame_step,
                           fft_length, np_rtype, threshold,
                           corrected_threshold):
    # Enable float64 support for RFFTs.
    with compat.forward_compatibility_horizon(2019, 10, 13):
      # Generate a random white Gaussian signal.
      signal = np.random.normal(size=signal_length).astype(np_rtype)

      stft = spectral_ops.stft(signal, frame_length, frame_step, fft_length,
                               pad_end=False)
      inverse_stft = spectral_ops.inverse_stft(stft, frame_length, frame_step,
                                               fft_length)
      inverse_stft_corrected = spectral_ops.inverse_stft(
          stft, frame_length, frame_step, fft_length,
          window_fn=spectral_ops.inverse_stft_window_fn(frame_step))
      inverse_stft, inverse_stft_corrected = self.evaluate(
          [inverse_stft, inverse_stft_corrected])

      # Truncate signal to the size of inverse stft.
      signal = signal[:inverse_stft.shape[0]]

      # Ignore the frame_length samples at either edge.
      signal = signal[frame_length:-frame_length]
      inverse_stft = inverse_stft[frame_length:-frame_length]
      inverse_stft_corrected = inverse_stft_corrected[
          frame_length:-frame_length]

      # Check that the inverse and original signal are close up to a scale
      # factor.
      inverse_stft_scaled = inverse_stft / np.mean(np.abs(inverse_stft))
      signal_scaled = signal / np.mean(np.abs(signal))
      self.assertLess(np.std(inverse_stft_scaled - signal_scaled), threshold)

      # Check that the inverse with correction and original signal are close.
      self.assertLess(np.std(inverse_stft_corrected - signal),
                      corrected_threshold)

  @parameterized.parameters(
      (256, 32),
      (256, 64),
      (128, 25),
      (127, 32),
      (128, 64))
  def test_inverse_stft_window_fn(self, frame_length, frame_step):
    """Test that inverse_stft_window_fn has unit gain at each window phase."""
    hann_window = window_ops.hann_window(frame_length, dtype=dtypes.float32)
    inverse_window_fn = spectral_ops.inverse_stft_window_fn(frame_step)
    inverse_window = inverse_window_fn(frame_length, dtype=dtypes.float32)
    hann_window, inverse_window = self.evaluate([hann_window, inverse_window])

    # Expect unit gain at each phase of the window.
    product_window = hann_window * inverse_window
    for i in range(frame_step):
      self.assertAllClose(1.0, np.sum(product_window[i::frame_step]))

  @parameterized.parameters((256, 64), (128, 32))
  def test_inverse_stft_window_fn_special_case(self, frame_length, frame_step):
    """Test inverse_stft_window_fn in special overlap = 3/4 case."""
    # Cases in which frame_length is an integer multiple of 4 * frame_step are
    # special because they allow exact reproduction of the waveform with a
    # squared Hann window (Hann window in both forward and reverse transforms).
    # In the case where frame_length = 4 * frame_step, that combination
    # produces a constant gain of 1.5, and so the corrected window will be the
    # Hann window / 1.5.
    hann_window = window_ops.hann_window(frame_length, dtype=dtypes.float32)
    inverse_window_fn = spectral_ops.inverse_stft_window_fn(frame_step)
    inverse_window = inverse_window_fn(frame_length, dtype=dtypes.float32)
    self.assertAllClose(hann_window, inverse_window * 1.5)

  @staticmethod
  def _compute_stft_gradient(signal, frame_length=32, frame_step=16,
                             fft_length=32):
    """Computes the gradient of the STFT with respect to `signal`."""
    # Enable float64 support for RFFTs.
    with compat.forward_compatibility_horizon(2019, 10, 13):
      stft = spectral_ops.stft(signal, frame_length, frame_step, fft_length)
      magnitude_stft = math_ops.abs(stft)
      loss = math_ops.reduce_sum(magnitude_stft)
      return gradients_impl.gradients([loss], [signal])[0]

  def test_gradients(self):
    """Test that spectral_ops.stft has a working gradient."""
    # TODO(rjryan): Update gradient tests for Eager.
    if context.executing_eagerly():
      return
    with self.session(use_gpu=True) as sess:
      signal_length = 512

      # An all-zero signal has all zero gradients with respect to the sum of the
      # magnitude STFT.
      empty_signal = array_ops.zeros([signal_length], dtype=dtypes.float32)
      empty_signal_gradient = sess.run(
          self._compute_stft_gradient(empty_signal))
      self.assertTrue((empty_signal_gradient == 0.0).all())

      # A sinusoid will have non-zero components of its gradient with respect to
      # the sum of the magnitude STFT.
      sinusoid = math_ops.sin(
          2 * np.pi * math_ops.linspace(0.0, 1.0, signal_length))
      sinusoid_gradient = self.evaluate(self._compute_stft_gradient(sinusoid))
      self.assertFalse((sinusoid_gradient == 0.0).all())

  @parameterized.parameters(
      (64, 16, 8, 16, np.float32, 2e-3, 5e-4),
      (64, 16, 8, 16, np.float64, 1e-8, 1e-8),
      (64, 16, 16, 16, np.float32, 2e-3, 5e-4),
      (64, 16, 16, 16, np.float64, 1e-8, 1e-8),
      (64, 16, 7, 16, np.float32, 2e-3, 5e-4),
      (64, 16, 7, 16, np.float64, 1e-8, 1e-8),
      (64, 7, 4, 9, np.float32, 2e-3, 5e-4),
      (64, 7, 4, 9, np.float64, 1e-8, 1e-8),
      (29, 5, 1, 10, np.float32, 2e-3, 5e-4),
      (29, 5, 1, 10, np.float64, 1e-8, 1e-8))
  def test_gradients_numerical(self, signal_length, frame_length, frame_step,
                               fft_length, np_rtype, forward_tol, backward_tol):
    # Enable float64 support for RFFTs.
    with compat.forward_compatibility_horizon(2019, 10, 13):
      # TODO(rjryan): Investigate why STFT gradient error is so high.
      signal = np.random.rand(signal_length).astype(np_rtype) * 2 - 1

      def forward(signal):
        return spectral_ops.stft(
            signal, frame_length, frame_step, fft_length, pad_end=False)
      ((f_jacob_t,), (f_jacob_n,)) = gradient_checker_v2.compute_gradient(
          forward, [signal])
      self.assertAllClose(f_jacob_t, f_jacob_n,
                          rtol=forward_tol, atol=forward_tol)

      def backward(stft):
        return spectral_ops.inverse_stft(
            stft, frame_length, frame_step, fft_length)

      stft = forward(signal)
      ((b_jacob_t,), (b_jacob_n,)) = gradient_checker_v2.compute_gradient(
          backward, [stft])
      self.assertAllClose(b_jacob_t, b_jacob_n,
                          rtol=backward_tol, atol=backward_tol)


if __name__ == "__main__":
  test.main()
