
from __future__ import print_function

import _setup_test_env  # noqa

import logging
import sys
import os
import os.path
import tensorflow as tf
import numpy as np
import unittest
import contextlib

from returnn.tf.network import *
from returnn.tf.layers.signal_processing import *
from returnn.config import Config

from returnn.util import better_exchook
better_exchook.replace_traceback_format_tb()


@contextlib.contextmanager
def make_scope():
  """
  :rtype: tf.compat.v1.Session
  """
  with tf.Graph().as_default() as graph:
    with tf_compat.v1.Session(graph=graph) as session:
      yield session


def test_melFilterbankLayer():
  with make_scope() as session:
    n_in, n_out = 257, 3
    layer_name = "mel_filterbank_layer"
    config = Config()
    config.update({
      "num_outputs": n_out,
      "num_inputs": n_in,
      "network": {
        layer_name: {
          "class": "mel_filterbank", "fft_size": 512, "nr_of_filters": n_out, "n_out": n_out, "is_output_layer": True,
          "from": "data:data"}
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    layer = network.layers[layer_name]
    test_out = session.run(
      layer.output.placeholder,
      feed_dict={network.get_extern_data('data').placeholder: np.ones((1, 1, 257))})
    assert np.sum(test_out - np.asarray([28.27923584, 53.10634232, 99.71585846], dtype=np.float32)) < 1e-5


def test_complexLinearProjectionLayer():
  with make_scope() as session:
    n_in, n_out = 514, 128
    layer_name = "clp_layer"
    config = Config()
    config.update({
      "num_outputs": n_out,
      "num_inputs": n_in,
      "network": {
        layer_name: {
          "class": "complex_linear_projection", "nr_of_filters": n_out, "n_out": n_out, "is_output_layer": True,
          "from": "data:data"}
      }})
    network = TFNetwork(config=config, train_flag=True)
    network.construct_from_dict(config.typed_value("network"))
    layer = network.layers[layer_name]
    assert isinstance(layer, ComplexLinearProjectionLayer)
    i_r = np.ones((1, n_in // 2))
    i_i = np.ones((1, n_in // 2)) * 0.5
    test_input = np.expand_dims(np.reshape(np.transpose(
      np.reshape(np.concatenate([i_r, i_i], axis=1), (1, 2, 257)), [0, 2, 1]), (1, 514)), 0)
    test_clp_kernel = np.ones((2, n_in // 2, 128))
    network.initialize_params(session)
    test_clp_output = session.run(
      layer.output.placeholder,
      feed_dict={network.get_extern_data('data').placeholder: test_input, layer._clp_kernel: test_clp_kernel})
    assert test_clp_output[0, 0, 0] - 6.00722122 < 1e-5


def test_MultichannelStftLayer():
  def _get_ref_output(time_sig, fft_size, frame_size, frame_shift, window_name, frame_nr, channel_nr):
      import numpy as np
      frame_start = frame_nr * frame_shift
      frame_end = frame_start + frame_size
      frame = time_sig[0, frame_start:frame_end, channel_nr]
      if window_name == "hanning":
          window = np.hanning(frame_size)
      windowed_frame = window * frame
      out = np.fft.rfft(windowed_frame, fft_size)
      return out

  def test_rfftStftConfig_01():
    with make_scope() as session:
      layer_name = "stft_layer"
      fft_size = 400
      frame_size = 400
      frame_shift = 160
      window = "hanning"
      test_input = np.ones((1, 32000, 2), dtype=np.float32)
      config = Config()
      config.update({
        "num_outputs": int(fft_size / 2) + 1 * test_input.shape[2],
        "num_inputs": test_input.shape[2],
        "network": {
          layer_name: {
            "class": "multichannel_stft_layer",
            "frame_shift": frame_shift, "frame_size": frame_size, "window": window,
            "fft_size": fft_size, "use_rfft": True, "nr_of_channels": 2, "is_output_layer": True,
            "from": "data:data"}
        }})
      network = TFNetwork(config=config, train_flag=True)
      network.construct_from_dict(config.typed_value("network"))
      layer = network.layers[layer_name]
      test_output = session.run(layer.output.placeholder, {network.get_extern_data('data').placeholder: test_input})
      ref0 = _get_ref_output(test_input, fft_size, frame_size, frame_shift, window, 0, 0)
      # np.fft.rfft and tensorflow.python.ops.rfft differ a little bit in their
      # results, thus an error margin is allowed in the result
      resultDiff = np.abs(test_output[0, 0, 0:(int(fft_size / 2) + 1)] - ref0)
      assert np.mean(resultDiff) < 0.02
      assert np.max(resultDiff) < 1
      pass

  test_rfftStftConfig_01()


def test_MultichannelMultiResStftLayer():
  def _get_ref_output_single_res(time_sig, fft_size, frame_size, frame_shift, window_name, frame_nr, channel_nr):
      import numpy as np
      frame_start = frame_nr * frame_shift
      frame_end = frame_start + frame_size
      frame = time_sig[0, frame_start:frame_end, channel_nr]
      if window_name == "hanning":
          window = np.hanning(frame_size)
      windowed_frame = window * frame
      out = np.fft.rfft(windowed_frame, fft_size)
      return out

  def test_stftConfig_single_res_01():
    with make_scope() as session:
      layer_name = "stft_layer"
      fft_sizes = [400]
      frame_sizes = [400]
      frame_shift = 160
      window = "hanning"
      test_input = np.ones((1, 32000, 2), dtype=np.float32)
      num_outputs = (int(fft_sizes[0] / 2) + 1) * test_input.shape[2]
      config = Config()
      config.update({
        "num_outputs": num_outputs,
        "num_inputs": test_input.shape[2],
        "network": {
          layer_name: {
            "class": "multichannel_multiresolution_stft_layer",
            "frame_shift": frame_shift, "frame_sizes": frame_sizes, "window": window,
            "fft_sizes": fft_sizes, "use_rfft": True, "nr_of_channels": 2, "is_output_layer": True,
            "from": "data:data"}
        }})
      network = TFNetwork(config=config, train_flag=True)
      network.construct_from_dict(config.typed_value("network"))
      layer = network.layers[layer_name]
      test_output = session.run(layer.output.placeholder, {network.get_extern_data('data').placeholder: test_input})
      ref0 = _get_ref_output_single_res(test_input, fft_sizes[0], frame_sizes[0], frame_shift, window, 0, 0)
      resultDiff = np.abs(test_output[0, 0, 0:(int(fft_sizes[0] / 2) + 1)] - ref0)
      assert test_output.shape[2] == num_outputs
      assert np.mean(resultDiff) < 0.02
      assert np.max(resultDiff) < 1

  def test_stftConfig_multi_res_01():
    with make_scope() as session:
      layer_name = "stft_layer"
      fft_sizes = [400, 200]
      frame_sizes = [400, 200]
      frame_shift = 160
      window = "hanning"
      test_input = np.ones((1, 32000, 2), dtype=np.float32)
      test_input[0, 1000:1100, 1] = np.ones((100), dtype=np.float32) * 0.5
      num_outputs = int(np.sum([(int(fft_size / 2) + 1) * test_input.shape[2] for fft_size in fft_sizes]))
      config = Config()
      config.update({
        "num_outputs": num_outputs,
        "num_inputs": test_input.shape[2],
        "network": {
          layer_name: {
            "class": "multichannel_multiresolution_stft_layer",
            "frame_shift": frame_shift, "frame_sizes": frame_sizes, "window": window,
            "fft_sizes": fft_sizes, "use_rfft": True, "nr_of_channels": 2, "is_output_layer": True,
            "from": "data:data"}
        }})
      network = TFNetwork(config=config, train_flag=True)
      network.construct_from_dict(config.typed_value("network"))
      layer = network.layers[layer_name]
      test_output = session.run(layer.output.placeholder, {network.get_extern_data('data').placeholder: test_input})
      assert test_output.shape[2] == num_outputs
      comparison_frame = 6
      ref00 = _get_ref_output_single_res(test_input, fft_sizes[0], frame_sizes[0], frame_shift, window, comparison_frame, 0)
      ref01 = _get_ref_output_single_res(test_input, fft_sizes[0], frame_sizes[0], frame_shift, window, comparison_frame, 1)
      ref10 = _get_ref_output_single_res(test_input, fft_sizes[1], frame_sizes[1], frame_shift, window, comparison_frame, 0)
      ref11 = _get_ref_output_single_res(test_input, fft_sizes[1], frame_sizes[1], frame_shift, window, comparison_frame, 1)
      ref = np.concatenate([ref00, ref01, ref10, ref11], axis=0)
      resultDiff = np.abs(test_output[0, comparison_frame, :] - ref)
      assert np.mean(resultDiff) < 0.02
      assert np.max(resultDiff) < 1

  def test_stftConfig_multi_res_02():
    with make_scope() as session:
      layer_name = "stft_layer"
      fft_sizes = [400, 200, 800]
      frame_sizes = [400, 200, 800]
      frame_shift = 160
      window = "hanning"
      test_input = np.random.normal(0, 0.6, (1, 3200, 2))
      num_outputs = int(np.sum([(int(fft_size / 2) + 1) * test_input.shape[2] for fft_size in fft_sizes]))
      config = Config()
      config.update({
        "num_outputs": num_outputs,
        "num_inputs": test_input.shape[2],
        "network": {
          layer_name: {
            "class": "multichannel_multiresolution_stft_layer",
            "frame_shift": frame_shift, "frame_sizes": frame_sizes, "window": window,
            "fft_sizes": fft_sizes, "use_rfft": True, "nr_of_channels": 2, "is_output_layer": True,
            "from": "data:data"}
        }})
      network = TFNetwork(config=config, train_flag=True)
      network.construct_from_dict(config.typed_value("network"))
      layer = network.layers[layer_name]
      test_output = session.run(layer.output.placeholder, {network.get_extern_data('data').placeholder: test_input})
      assert test_output.shape[2] == num_outputs
      comparison_frame = 6
      ref00 = _get_ref_output_single_res(test_input, fft_sizes[0], frame_sizes[0], frame_shift, window, comparison_frame, 0)
      ref01 = _get_ref_output_single_res(test_input, fft_sizes[0], frame_sizes[0], frame_shift, window, comparison_frame, 1)
      ref10 = _get_ref_output_single_res(test_input, fft_sizes[1], frame_sizes[1], frame_shift, window, comparison_frame, 0)
      ref11 = _get_ref_output_single_res(test_input, fft_sizes[1], frame_sizes[1], frame_shift, window, comparison_frame, 1)
      ref20 = _get_ref_output_single_res(test_input, fft_sizes[2], frame_sizes[2], frame_shift, window, comparison_frame, 0)
      ref21 = _get_ref_output_single_res(test_input, fft_sizes[2], frame_sizes[2], frame_shift, window, comparison_frame, 1)
      ref = np.concatenate([ref00, ref01, ref10, ref11, ref20, ref21], axis=0)
      resultDiff = np.abs(test_output[0, comparison_frame, :] - ref)
      assert np.mean(resultDiff) < 0.06
      assert np.max(resultDiff) < 1

  test_stftConfig_single_res_01()
  test_stftConfig_multi_res_01()
  test_stftConfig_multi_res_02()


if __name__ == "__main__":
  better_exchook.install()
  if len(sys.argv) <= 1:
    for k, v in sorted(globals().items()):
      if k.startswith("test_"):
        print("-" * 40)
        print("Executing: %s" % k)
        try:
          v()
        except unittest.SkipTest as exc:
          print("SkipTest:", exc)
        print("-" * 40)
    print("Finished all tests.")
  else:
    assert len(sys.argv) >= 2
    for arg in sys.argv[1:]:
      print("Executing: %s" % arg)
      if arg in globals():
        globals()[arg]()  # assume function and execute
      else:
        eval(arg)  # assume Python code and execute
