Audio_Effects_SDK/samples/effects_demo/effects_demo.cpp

/*
Copyright (c) 2021-2022, NVIDIA CORPORATION. All rights reserved.

NVIDIA CORPORATION and its licensors retain all intellectual property
and proprietary rights in and to this software, related documentation
and any modifications thereto. Any use, reproduction, disclosure or
distribution of this software and related documentation without an express
license agreement from NVIDIA CORPORATION is strictly prohibited.
*/
#include <math.h>

#include <algorithm>
#include <chrono>
#include <cstring>
#include <iostream>
#include <iomanip>
#include <memory>
#include <set>
#include <sstream>
#include <thread>
#include <tuple>

#include "utils/wave_reader/waveReadWrite.hpp"
#include "utils/ConfigReader.hpp"

#include <nvAudioEffects.h>
#include <map>

namespace {

const char kConfigEffectVariable[] = "effect";
const char kConfigSampleRateVariable[] = "sample_rate";
const char kConfigFileInputVariable[] = "input_wav_list";
const char kConfigFileInputFarendVariable[] = "input_farend_wav_list";
const char kConfigFileOutputVariable[] = "output_wav_list";
const char kConfigFileRTVariable[] = "real_time";
const char kConfigResetVariable[] = "reset";
const char kConfigFrameSize[] = "frame_size";
const char kConfigUseDefaultGpu[] = "use_default_gpu";
const char kConfigLogTarget[] = "log_target_list";
const char kConfigLogTargetFile[] = "log_target_file";
const char kConfigLogLevel[] = "log_level";
const char kConfigLogTargetFileDefault[] = "/tmp/nvAudioEffects_log.txt";
const char kConfigFileModelVariable[] = "model";
const char kConfigIntensityRatioVariable[] = "intensity_ratio";
const char kConfigVadEnable[] = "enable_vad";
const char kConfigChainedEffectGpuList[] = "chained_effect_gpu_list";

/* allowed sample rates */
const std::vector<uint32_t> kAllowedSampleRates = { 8000, 16000, 48000 };

}  // namespace

struct StreamData {
  // Wav file name(s) separated by ';'
  std::string input_wav_file;
  // Total number of samples in all input wav files together
  unsigned num_samples = 0;
  // Number of audio samples written to output file
  unsigned already_written = 0;
  // 32-bit float audio samples read from all input wav files
  std::vector<float>* input_wav_samples;
  // (AEC only) 32-bit float audio samples for farend input wav files
  std::vector<float>* input_farend_wav_samples;
  // Wav file name(s) to write the output data
  std::string output_wav_file;
  // Object for wav writing library
  std::vector<std::unique_ptr<CWaveFileWrite>> wav_write;
  // Index of each input file end in the audio data vector
  std::vector<int> file_end_offsets;
  // Index of current output file used as suffix to filename
  int output_file_num;
};

class EffectsDemoApp {
 public:
  bool run(const ConfigReader& config_reader);

 private:
  // Validate configuration data.
  bool validateConfig(const ConfigReader& config_reader);
  // write to output wav
  bool writeOutputWav(const std::vector<float*>& output_samples, size_t num_samples_per_channel,
                      int stream_index);
  // sample rate config
  uint32_t input_sample_rate_ = 0;
  // output sample rate config
  uint32_t output_sample_rate_ = 0;
  // stream data
  std::vector<StreamData> stream_data_;
  // GPU selection flag - by default set to false
  uint32_t use_default_gpu_ = 0;
  // inited from configuration
  bool real_time_ = false;
  // intensity_ratio_ config, lies between 0 (no effect applied) to 1 (fully turned on)
  float intensity_ratio_ = 1.0f;
  // VAD support (denoiser/dereverb+denoiser only)
  bool vad_supported_ = false;

  unsigned int output_channels_ = 1;
};

bool EffectsDemoApp::validateConfig(const ConfigReader& config_reader) {
  if (config_reader.IsConfigValueAvailable(kConfigEffectVariable) == false) {
    std::cerr << "No " << kConfigEffectVariable << " variable found" << std::endl;
    return false;
  }

  if (config_reader.IsConfigValueAvailable(kConfigFileModelVariable) == false) {
    std::cerr << "No " << kConfigFileModelVariable << " variable found" << std::endl;
    return false;
  }

  if (config_reader.IsConfigValueAvailable(kConfigFileInputVariable) == false) {
    std::cerr << "No " << kConfigFileInputVariable << " variable found" << std::endl;
    return false;
  }

  if (config_reader.IsConfigValueAvailable(kConfigFileOutputVariable) == false) {
    std::cerr << "No " << kConfigFileOutputVariable << " variable found" << std::endl;
    return false;
  }

  std::string real_time;
  if (config_reader.GetConfigValue(kConfigFileRTVariable, &real_time) == false) {
    std::cerr << "No " << kConfigFileRTVariable << " variable found" << std::endl;
    return false;
  }

  if (real_time[0] != '0') {
    real_time_ = true;
  }

  if (config_reader.IsConfigValueAvailable(kConfigResetVariable)) {
    std::cout << "Reset available" << std::endl;
  }

  std::string use_default_gpu;
  if (config_reader.IsConfigValueAvailable(kConfigUseDefaultGpu) &&
      config_reader.GetConfigValue(kConfigUseDefaultGpu, &use_default_gpu)) {
    use_default_gpu_ = std::strtoul(use_default_gpu.c_str(), nullptr, 0);
    std::cout << "Use default GPU: " << use_default_gpu_ << std::endl;
  }

  std::string intensity_ratio;
  float intensity_ratio_local;
  if (config_reader.GetConfigValue(kConfigIntensityRatioVariable, &intensity_ratio)) {
    intensity_ratio_local = std::strtof(intensity_ratio.c_str(), nullptr);
    if (intensity_ratio_local < 0.0f || intensity_ratio_local > 1.0f) {
      std::cerr << kConfigIntensityRatioVariable << " not supported" << std::endl;
      return false;
    }
  } else {
    intensity_ratio_local = 1.0f;
  }

  intensity_ratio_ = intensity_ratio_local;
  std::cout << "intensity ratio is ::" << intensity_ratio_ << "!!" << std::endl;
  return true;
}

bool EffectsDemoApp::writeOutputWav(const std::vector<float*>& output_samples,
                                    size_t num_samples_per_channel, int stream_index) {
  size_t to_write = num_samples_per_channel;
  size_t max_allowed = stream_data_[stream_index].num_samples *
      (output_sample_rate_/input_sample_rate_);
  size_t already_written = stream_data_[stream_index].already_written;
  size_t can_write = static_cast<size_t>(max_allowed - already_written);
  if (can_write == 0) { return true; }
  size_t will_write = std::min(to_write, can_write);

  for (int i = 0; i < output_channels_; i++) {
    if (!stream_data_[stream_index].wav_write[i]->writeChunk(output_samples[i],
                                                             will_write * sizeof(float))) {
      std::cerr << "Could not write output to file "
                << stream_data_[stream_index].output_wav_file << std::endl;
      return false;
    }
  }
  stream_data_[stream_index].already_written += will_write;
  return true;
}

void logger_cb(LoggingSeverity level, const char* log, void* userdata) {
  std::cout << "LOG" << '(' << LogSeverityToString(level) << ") " << log << '\n';
}

inline LoggingSeverity StringToLogSeverity(std::string &severity) {
  if ("ERROR" == severity) return LOG_LEVEL_ERROR;
  if ("WARNING" == severity) return LOG_LEVEL_WARNING;
  if ("INFO" == severity) return LOG_LEVEL_INFO;
  return LOG_LEVEL_ERROR;
}

inline LoggingTarget StringToLogTarget(std::string &severity) {
  if ("NONE" == severity) return LOG_TARGET_NONE;
  if ("STDERR" == severity) return LOG_TARGET_STDERR;
  if ("FILE" == severity) return LOG_TARGET_FILE;
  if ("CALLBACK" == severity) return LOG_TARGET_CALLBACK;
  return LOG_TARGET_NONE;
}

bool EffectsDemoApp::run(const ConfigReader& config_reader) {
  if (validateConfig(config_reader) == false) {
    return false;
  }

  auto input_wav_list = config_reader.GetConfigValueList(kConfigFileInputVariable);
  auto output_wav_list = config_reader.GetConfigValueList(kConfigFileOutputVariable);
  // Support 1024 max inputs to simplify
  const int kMaxSupportedInputs = 1024;
  if (input_wav_list.size() > kMaxSupportedInputs) {
    std::cerr << "This sample application supports a maximum of " << kMaxSupportedInputs <<
                 " input files (although more may be supported by the SDK depending on GPU,"
                 " please refer to the programming guide for limits and procedure for"
                 " setting batch size in SDK)." << std::endl;
    return false;
  }
  std::vector<std::string> log_targets;
  if (config_reader.IsConfigValueAvailable(kConfigLogTarget)) {
    log_targets = config_reader.GetConfigValueList(kConfigLogTarget);
  }
  LoggingSeverity severity = LOG_LEVEL_INFO;
  int target = LOG_TARGET_STDERR;
  std::string log_file(kConfigLogTargetFileDefault);

  if (!log_targets.empty()) {
    target = LOG_TARGET_NONE;
    for (auto& item : log_targets) {
      target |= StringToLogTarget(item);
    }
  }

  if (target & LOG_TARGET_FILE) {
    if (config_reader.IsConfigValueAvailable(kConfigLogTargetFile)) {
      log_file = config_reader.GetConfigValue(kConfigLogTargetFile);
    }
  }

  if (config_reader.IsConfigValueAvailable(kConfigLogLevel)) {
    auto log_level = config_reader.GetConfigValue(kConfigLogLevel);
    if (!log_level.empty()) {
      severity = StringToLogSeverity(log_level);
    }
  }

  if (input_wav_list.size() > output_wav_list.size()) {
    std::cout << "Error: Input and output wav files list size mismatch found" << std::endl;
    return false;
  }

  std::vector<bool> reset_list(input_wav_list.size(), false);
  if (config_reader.IsConfigValueAvailable(kConfigResetVariable)) {
    auto config_list = config_reader.GetConfigValueList(kConfigResetVariable);
    for (auto& item : config_list) {
      char* p;
      unsigned int num = strtol(item.c_str(), &p, 10);
      if (*p != 0) {
        // Invalid?
        continue;
      }

      if (num < 1 || num > input_wav_list.size()) {
        std::cerr << "Error: Invalid stream specified for reset" <<std::endl;
        return false;
      }

      reset_list[num-1] = true;
    }
  }

  const unsigned num_streams = input_wav_list.size();

  if (real_time_) {
    std::cout << "App will run in real time mode ..." << std::endl;
  }

  NvAFX_Status log_status;
  log_status = NvAFX_InitializeLogger(severity, target, log_file.c_str(),
                                      target & LOG_TARGET_CALLBACK ? logger_cb : nullptr,
                                      nullptr);

  if (log_status != NVAFX_STATUS_SUCCESS) {
    std::cerr << "NvAFX_InitializeLogger() failed" << std::endl;
    return false;
  }

  int num_effects;
  NvAFX_EffectSelector* supported_effects;
  if (NvAFX_GetEffectList(&num_effects, &supported_effects) != NVAFX_STATUS_SUCCESS) {
    std::cerr << "NvAFX_GetEffectList() failed" << std::endl;
    return false;
  }

  std::cout << "Total Effects supported: " << num_effects << std::endl;
  for (int i = 0; i < num_effects; ++i) {
    std::cout << "(" << i + 1 << ") " << supported_effects[i] << std::endl;
  }

  bool is_aec = false;

  NvAFX_Handle handle;
  std::vector<std::string> effects = config_reader.GetConfigValueList(kConfigEffectVariable);
  std::vector<std::string> sample_rate_str = config_reader.GetConfigValueList(kConfigSampleRateVariable);
  std::vector<uint32_t> sample_rates;
  for (auto& s: sample_rate_str) {
    sample_rates.push_back(std::strtoul(s.c_str(), nullptr, 0));
  }

  std::string effect_name;
  if (effects.size() == 1) {
    // Single effect
    if (sample_rates.size() != effects.size()) {
      std::cerr << "Expected single sample rate for single effect" << std::endl;
      return false;
    }
    input_sample_rate_ = sample_rates[0];
    if (std::find(kAllowedSampleRates.begin(), kAllowedSampleRates.end(), input_sample_rate_) ==
        kAllowedSampleRates.end()) {
      std::cerr << "Sample rate " << input_sample_rate_ << " not supported" << std::endl;
      return false;
    }
    NvAFX_Status status = NVAFX_STATUS_FAILED;
    if (effects[0] == "denoiser") {
      status = NvAFX_CreateEffect(NVAFX_EFFECT_DENOISER, &handle);
    } else if (effects[0] == "dereverb") {
      status = NvAFX_CreateEffect(NVAFX_EFFECT_DEREVERB, &handle);
    } else if (effects[0] == "dereverb_denoiser") {
      status = NvAFX_CreateEffect(NVAFX_EFFECT_DEREVERB_DENOISER, &handle);
    } else if (effects[0] == "aec") {
      status = NvAFX_CreateEffect(NVAFX_EFFECT_AEC, &handle);
      is_aec = true;
    } else if (effects[0] == "superres") {
      status = NvAFX_CreateEffect(NVAFX_EFFECT_SUPERRES, &handle);
    } else {
      std::cerr << "NvAFX_CreateEffect() failed. Invalid Effect Value : " << effects[0] << std::endl;
      return false;
    }
    if (status == NVAFX_UNSUPPORTED_RUNTIME) {
      float version = (CUDA_SUPPORTED_RUNTIME / 1000) + (CUDA_SUPPORTED_RUNTIME%100)/100.f;
      std::cerr << "Unsupported CUDA runtime (requires >= " <<  version << "). "
                   "Please ensure that a driver supporting the required CUDA version is installed (or if using FCU, library path "
                   "contains the correct CUDA compat libraries). For more details, please refer to the " 
                   "programming guide." << std::endl;
      return false;
    } else if (status != NVAFX_STATUS_SUCCESS) {
      std::cerr << "NvAFX_CreateEffect() failed" << std::endl;
      return false;
    }

    effect_name = effects[0];
    const std::set<std::string> kVadSupportedEffects = {"denoiser", "dereverb_denoiser"};
    vad_supported_ = (kVadSupportedEffects.find(effects[0]) != kVadSupportedEffects.end()) &&
        config_reader.IsConfigValueAvailable(kConfigVadEnable) &&
        std::strtoul(config_reader.GetConfigValue(kConfigVadEnable).c_str(), nullptr, 0);


    if (vad_supported_) {
      std::cout << "Enabling VAD" << std::endl;
      if (NvAFX_SetU32(handle, NVAFX_PARAM_ENABLE_VAD, 1) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "Could not enable VAD" << std::endl;
        return false;
      }
    }

    // If the system has multiple supported GPUs, then the application can either
    // use CUDA driver APIs or CUDA runtime APIs to enumerate the GPUs and select one based on the
    // application's requirements or offload the responsibility to SDK to select the GPU by setting
    // NVAFX_PARAM_USE_DEFAULT_GPU as 1
    if (NvAFX_SetU32(handle, NVAFX_PARAM_USE_DEFAULT_GPU, use_default_gpu_) != NVAFX_STATUS_SUCCESS) {
      std::cerr << "NvAFX_SetBool(NVAFX_PARAM_USE_DEFAULT_GPU " << ") failed" << std::endl;
    }

    status = NvAFX_SetU32(handle, NVAFX_PARAM_INPUT_SAMPLE_RATE, input_sample_rate_);
    if (status == NVAFX_STATUS_INVALID_PARAM) {
      // Try depreciated param
      status = NvAFX_SetU32(handle, NVAFX_PARAM_SAMPLE_RATE, input_sample_rate_);
    }
    if (status!= NVAFX_STATUS_SUCCESS) {
      std::cerr << "NvAFX_SetU32(Sample Rate: " << input_sample_rate_ << ") failed" << std::endl;
      return false;
    }

  } else {
    const std::map<std::tuple<std::string, std::string, uint32_t, uint32_t>, NvAFX_EffectSelector> supported_configs =
    {
        // 16k Effect + Superres 16k-48k
        {std::make_tuple(std::string("denoiser"), std::string("superres"), 16000, 16000), NVAFX_CHAINED_EFFECT_DENOISER_16k_SUPERRES_16k_TO_48k},
        {std::make_tuple(std::string("dereverb"), std::string("superres"), 16000, 16000), NVAFX_CHAINED_EFFECT_DEREVERB_16k_SUPERRES_16k_TO_48k},
        {std::make_tuple(std::string("dereverb_denoiser"), std::string("superres"), 16000, 16000),
         NVAFX_CHAINED_EFFECT_DEREVERB_DENOISER_16k_SUPERRES_16k_TO_48k},

         // Superres 8k-16k + 16k Effect
        {std::make_tuple("superres", "denoiser", 8000, 16000), NVAFX_CHAINED_EFFECT_SUPERRES_8k_TO_16k_DENOISER_16k },
        {std::make_tuple("superres", "dereverb", 8000, 16000), NVAFX_CHAINED_EFFECT_SUPERRES_8k_TO_16k_DEREVERB_16k },
        {std::make_tuple("superres", "dereverb_denoiser", 8000, 16000),
         NVAFX_CHAINED_EFFECT_SUPERRES_8k_TO_16k_DEREVERB_DENOISER_16k },
        };
    auto effect_config = supported_configs.find(std::make_tuple(effects[0], effects[1], sample_rates[0], sample_rates[1]));
    if (effect_config == supported_configs.end()) {
      std::cerr << "Unsupported effect chain" <<std::endl;
      return false;
    }
    auto status = NvAFX_CreateChainedEffect(effect_config->second,&handle);
    if (status == NVAFX_UNSUPPORTED_RUNTIME) {
      float version = (CUDA_SUPPORTED_RUNTIME / 1000) + (CUDA_SUPPORTED_RUNTIME%100)/100.f;
      std::cerr << "Unsupported CUDA runtime (requires >= " <<  version << "). "
                   "Please ensure that a driver supporting the required CUDA version is installed (or if using FCU, library path "
                   "contains the correct CUDA compat libraries). For more details, please refer to the " 
                   "programming guide." << std::endl;
      return false;
    } else if (status != NVAFX_STATUS_SUCCESS) {
      std::cerr << "Could not create effect" << std::endl;
      return false;
    }
    input_sample_rate_ = sample_rates[0];

    effect_name = "Chained Effect (" + effects[0] + " + " + effects[1] + ")";

    if (config_reader.IsConfigValueAvailable(kConfigChainedEffectGpuList)) {
      std::vector<std::string> gpu_list = config_reader.GetConfigValueList(kConfigChainedEffectGpuList);
      std::vector<uint32_t> gpus;
      for (auto& s: gpu_list) {
        gpus.push_back(std::strtoul(s.c_str(), nullptr, 0));
      }

      NvAFX_SetU32List(handle, NVAFX_PARAM_CHAINED_EFFECT_GPU_LIST, gpus.data(), gpus.size());
    }

  }

  unsigned num_input_samples_per_frame = 0;
  //Superres has more number of output samples
  // Note: Number of samples are per channel
  unsigned num_output_samples_per_frame = 0;
  if (config_reader.IsConfigValueAvailable(kConfigFrameSize)) {
    auto config_value = config_reader.GetConfigValue(kConfigFrameSize);
    unsigned frame_size = std::strtoul(config_value.c_str(), nullptr, 0);
    // frame_size is in milliseconds.
    num_input_samples_per_frame = (input_sample_rate_ * frame_size) / 1000;
  }

    // Obtain farend audio if effect is AEC
  std::vector<std::string> input_farend_wav_list;
  if (is_aec) {
    input_farend_wav_list = config_reader.GetConfigValueList(kConfigFileInputFarendVariable);
    if (input_farend_wav_list.size() != input_wav_list.size()) {
      std::cerr << "AEC effect requires farend audio as input in addition to nearend audio."
                   "Please ensure that the config value \"" << kConfigFileInputFarendVariable <<
                   "\" is present and has correct number of files" << std::endl;
      return false;
    }
  }

  NvAFX_Status status;

  std::vector<std::string> model_files = config_reader.GetConfigValueList(kConfigFileModelVariable);
  // Note: For single effect, model can be passed in as a single string or as a list of strings with size 1
  std::unique_ptr<char*[]> model_files_param(new char*[model_files.size()]);
  for (int i = 0; i < model_files.size(); i++) {
    model_files_param[i] = (char*) model_files[i].data();
  }
  if (NvAFX_SetStringList(handle, NVAFX_PARAM_MODEL_PATH, (const char**)model_files_param.get(),
                          model_files.size())
      != NVAFX_STATUS_SUCCESS) {
    std::cerr << "NvAFX_SetString() failed" << std::endl;
    return false;
  }

  if (NvAFX_SetU32(handle, NVAFX_PARAM_NUM_STREAMS, num_streams) != NVAFX_STATUS_SUCCESS) {
    std::cerr << "NvAFX_SetU32(NVAFX_PARAM_NUM_STREAMS) failed" << std::endl;
    return false;
  }

  unsigned int list_size = 0;
  std::unique_ptr<unsigned int[]> supported_list = nullptr;
  auto ret = NvAFX_GetU32List(handle, NVAFX_PARAM_SUPPORTED_NUM_SAMPLES_PER_FRAME,
                              supported_list.get(), &list_size);
  if (ret != NVAFX_STATUS_OUTPUT_BUFFER_TOO_SMALL) {
    std::cerr << "NvAFX_GetU32List(NVAFX_PARAM_SUPPORTED_NUM_SAMPLES_PER_FRAME) failed."
              << std::endl;
    return false;
  }
  supported_list.reset(new unsigned int[list_size]);
  if (NvAFX_GetU32List(handle, NVAFX_PARAM_SUPPORTED_NUM_SAMPLES_PER_FRAME,
                       supported_list.get(), &list_size) != NVAFX_STATUS_SUCCESS) {
    std::cerr << "NvAFX_GetU32List(NVAFX_PARAM_SUPPORTED_NUM_SAMPLES_PER_FRAME) failed."
    << std::endl;
    return false;
  }
  // If config specifies certain value check if it is in supported list
  if (config_reader.IsConfigValueAvailable(kConfigFrameSize)) {
    auto pointer = std::find(supported_list.get(), supported_list.get() + list_size,
                             num_input_samples_per_frame);
    if (pointer == supported_list.get() + list_size) {
      std::ostringstream oss;
      std::string separator("");
      oss << "'";
      for (unsigned i = 0; i < list_size; ++i) {
        oss << separator << (supported_list[i] * 1000) / input_sample_rate_;
        separator = ", ";
      }
      oss << "'.";
      std::cerr << "Supplied value for " << kConfigFrameSize << " is not supported. Supplied value: "
                << config_reader.GetConfigValue(kConfigFrameSize) << "." << " Supported values: "
                << oss.str() << std::endl;
      return false;
    }
  } else {
    num_input_samples_per_frame = supported_list[0];
  }

  status = NvAFX_SetU32(handle, NVAFX_PARAM_NUM_SAMPLES_PER_INPUT_FRAME, num_input_samples_per_frame);
  if (status == NVAFX_STATUS_INVALID_PARAM) {
    // Try previous version
    status = NvAFX_SetU32(handle, NVAFX_PARAM_NUM_SAMPLES_PER_FRAME, num_input_samples_per_frame);
  }
  if (status!= NVAFX_STATUS_SUCCESS) {
    std::cerr << "NvAFX_SetU32(NVAFX_PARAM_NUM_SAMPLES_PER_FRAME) failed" << std::endl;
    return false;
  }

  // Intensity ratio may not be supported for some effects, ignore if status is invalid param
  status = NvAFX_SetFloat(handle, NVAFX_PARAM_INTENSITY_RATIO, intensity_ratio_);
  if (status != NVAFX_STATUS_SUCCESS && status != NVAFX_STATUS_INVALID_PARAM) {
    std::cerr << "NvAFX_SetFloat(NVAFX_PARAM_INTENSITY_RATIO) failed" << std::endl;
  }

  std::cout << "Loading effect" << " ... ";
  if (NvAFX_Load(handle) != NVAFX_STATUS_SUCCESS) {
    std::cerr << "NvAFX_Load() failed" << std::endl;
    return false;
  }
  std::cout << "Done" << std::endl;

  // Try to get output number of samples per frame (valid only with new SDK)
  status = NvAFX_GetU32(handle, NVAFX_PARAM_NUM_SAMPLES_PER_OUTPUT_FRAME,
                        &num_output_samples_per_frame);
  if (status != NVAFX_STATUS_SUCCESS) {
    if (status == NVAFX_STATUS_INVALID_PARAM) {
      // Old SDK, output samples = input samples
      num_output_samples_per_frame = num_input_samples_per_frame;
    } else {
      std::cerr << "NvAFX_GetU32() failed: NVAFX_PARAM_NUM_SAMPLES_PER_OUTPUT_FRAME" << std::endl;
      return false;
    }
  }

  // Try to get output sample rate (valid only with new SDK)
  status = NvAFX_GetU32(handle, NVAFX_PARAM_OUTPUT_SAMPLE_RATE, &output_sample_rate_);
  if (status != NVAFX_STATUS_SUCCESS) {
    if (status == NVAFX_STATUS_INVALID_PARAM) {
      // Old SDK, output sample rate = input sample rate
      output_sample_rate_ = input_sample_rate_;
    } else {
      std::cerr << "NvAFX_GetU32() failed: NVAFX_PARAM_OUTPUT_SAMPLE_RATE" << std::endl;
      return false;
    }
  }

  // Try to get input/output channels, if that fails (old SDK?), try depreciated num_channels
  unsigned num_input_channels, num_output_channels;
  status = NvAFX_GetU32(handle, NVAFX_PARAM_NUM_INPUT_CHANNELS, &num_input_channels);
  if (status  != NVAFX_STATUS_SUCCESS) {
    if (status == NVAFX_STATUS_INVALID_PARAM) {
      if (NvAFX_GetU32(handle, NVAFX_PARAM_NUM_CHANNELS, &num_input_channels) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "NvAFX_GetU32(NVAFX_PARAM_NUM_CHANNELS) failed" << std::endl;
        return false;
      } else {
        num_output_channels = num_input_channels;
      }
    } else {
      std::cerr << "NvAFX_GetU32() failed" << std::endl;
      return false;
    }
  } else if (NvAFX_GetU32(handle, NVAFX_PARAM_NUM_OUTPUT_CHANNELS, &num_output_channels) !=
  NVAFX_STATUS_SUCCESS) {
    std::cerr << "NvAFX_GetU32(NVAFX_PARAM_NUM_OUTPUT_CHANNELS) failed" << std::endl;
    return false;
  }

  output_channels_ = num_output_channels;

  std::cout << "Effect properties:" << std::endl
            << "  Effect             : " << effect_name << std::endl
            << "  Input Channels     : " << num_input_channels << std::endl
            << "  Input Sample Rate  : " << input_sample_rate_ << std::endl
            << "  Output Sample Rate : " << output_sample_rate_ << std::endl
            << "  Samples per frame  : " << num_input_samples_per_frame << std::endl
            << "  Number of streams  : " << num_streams << std::endl;

  size_t max_num_input_samples = 0;
  // read the noisy input wav file data and cache it in RAM.
  for (unsigned i = 0; i < num_streams; ++i) {
    StreamData data;
    data.input_wav_file = input_wav_list[i];

    if (!ReadWavFile(input_wav_list[i], input_sample_rate_, &data.input_wav_samples,
                     &data.num_samples, &data.file_end_offsets, num_input_samples_per_frame)) {
      std::cerr << "Unable to read wav file: " << input_wav_list[i] << std::endl;
      if (errno == EMFILE) {
        std::cerr << "Open file limit reached. Please increase file limit using the ulimit "
                     "utility (for example, \"ulimit -n 20000\"). Please refer to the "
                     "documentation of the ulimit utility for more details." << std::endl;
      }
      return false;
    }

    // max_num_samples should be aligned to 'num_input_samples_per_frame' automatically
    if (max_num_input_samples < data.input_wav_samples->size()) {
      max_num_input_samples = data.input_wav_samples->size();
    }

    if (is_aec) {
      unsigned int num_samples = 0;
      std::vector<int> end_offsets;

      if (!ReadWavFile(input_farend_wav_list[i], input_sample_rate_, &data.input_farend_wav_samples,
                       &num_samples, &end_offsets, num_input_samples_per_frame)) {
        std::cerr << "Unable to read wav file: " << input_farend_wav_list[i] << std::endl;
        if (errno == EMFILE) {
          std::cerr << "Open file limit reached. Please increase file limit using the ulimit "
                       "utility (for example, \"ulimit -n 20000\"). Please refer to the "
                       "documentation of the ulimit utility for more details." << std::endl;
        }
        return false;
      }

      if (num_samples != data.num_samples || end_offsets != data.file_end_offsets) {
        std::cerr << "Input farend file specification does not match nearend file specification."
                     "Farend and Nearend files must have the same number of samples";
        return false;
      }
    }

    data.output_wav_file = output_wav_list[i];
    if (output_channels_ > 1) {
      auto filename_wo_ext = output_wav_list[i].substr(0, output_wav_list[i].find_last_of("."));

      for (int ch = 0; ch < output_channels_; ch++) {
        std::string filename = filename_wo_ext + "_ch" + std::to_string(ch + 1) + ".wav";
        data.wav_write.emplace_back(
            std::unique_ptr<CWaveFileWrite>(new CWaveFileWrite(filename,
                                                               output_sample_rate_,
                                                               1, 32, true)));

      }
    } else {
      data.wav_write.emplace_back(std::unique_ptr<CWaveFileWrite>(new CWaveFileWrite(output_wav_list[i], output_sample_rate_,
                                                                                     1, 32, true)));
    }
    for (int ch = 0; ch < num_output_channels; ch++) {
      if (!data.wav_write[ch]->initFile()) {
        std::cerr << "Unable to open file for writing: "
                  << data.wav_write[ch]->getFileName() << std::endl;
        if (errno == EMFILE) {
          std::cerr << "Open file limit reached. Please increase file limit using the ulimit "
                       "utility (for example, \"ulimit -n 20000\"). Please refer to the "
                       "documentation of the ulimit utility for more details." << std::endl;
        }
        return false;
      }

    }
    data.output_file_num = 0;

    stream_data_.push_back(std::move(data));
#ifndef ENABLE_PERF_DUMP
    std::cout << "Input wav file: " << input_wav_list[i] << std::endl
              << "Total " << data.num_samples << " samples read" << std::endl;
#endif  // ENABLE_PERF_DUMP
  }

  // make all sizes same to ease effect run loop's work
  for (auto& data : stream_data_) {
    data.input_wav_samples->resize(max_num_input_samples);
  }

  float frame_in_secs = static_cast<float>(num_input_samples_per_frame) /
                        static_cast<float>(input_sample_rate_);
  float total_run_time = 0.f;
  float total_audio_duration = 0.f;
  float checkpoint = 0.1f;
  float expected_audio_duration = static_cast<float>(max_num_input_samples) /
                                  static_cast<float>(input_sample_rate_);
  std::vector<float> input_frame(num_input_channels * num_input_samples_per_frame * num_streams);
  std::vector<float> output_frame(num_output_channels * num_output_samples_per_frame * num_streams);

  std::string progress_bar = "[          ] ";
  std::cout << "Processed: " << progress_bar << "0%\r";
  std::cout.flush();

  std::vector<NvAFX_Bool> bitmap(input_wav_list.size(), NVAFX_FALSE);

  // wav data is already padded to align to num_input_samples_per_frame by ReadWavFile()
  for (size_t offset = 0; offset < max_num_input_samples; offset += num_input_samples_per_frame) {
    // prepare input data
    bool should_reset = false;
    for (unsigned i = 0; i < num_streams; ++i) {
      float *in = stream_data_[i].input_wav_samples->data() + offset;
      std::copy(in, in + num_input_samples_per_frame, &input_frame[i * num_input_samples_per_frame]);
      if (is_aec) {
        in = stream_data_[i].input_farend_wav_samples->data() + offset;
        std::copy(in, in + num_input_samples_per_frame, &input_frame[(num_streams + i) * num_input_samples_per_frame]);
      }

      // Check if this file finished, if so, start a new file
      if (stream_data_[i].file_end_offsets.size() &&
          offset == static_cast<unsigned>(stream_data_[i].file_end_offsets[0])) {
        stream_data_[i].file_end_offsets.erase(stream_data_[i].file_end_offsets.begin());
        if (stream_data_[i].file_end_offsets.size()) {
          // Open a new file only if this isn't the last one

          if (output_channels_ > 1) {
            uint32_t output_num = ++stream_data_[i].output_file_num;
            auto filename_wo_ext = output_wav_list[i].substr(0, output_wav_list[i].find_last_of("_"));
            for (int ch = 0; ch < output_channels_; ch++) {
              stream_data_[i].wav_write[ch]->commitFile();
              std::string filename = filename_wo_ext + + "_" + std::to_string(output_num) +
                  "_ch" + std::to_string(ch + 1) + ".wav";
              stream_data_[i].wav_write[ch] = std::unique_ptr<CWaveFileWrite>(new CWaveFileWrite(filename, output_sample_rate_,
                                                                     1, 32, true));
              if (!stream_data_[i].wav_write[ch]->initFile()) {
                std::cerr << "Unable to open file for writing: " << filename;
                return false;
              }
            }
          } else {
            auto ext = output_wav_list[i].find_last_of(".");
            std::string filename = output_wav_list[i].substr(0, ext) + "_" +
                std::to_string(++stream_data_[i].output_file_num) + ".wav";
            stream_data_[i].wav_write[0] = std::unique_ptr<CWaveFileWrite>(new CWaveFileWrite(filename, output_sample_rate_,
                                                                   1, 32, true));

            if (!stream_data_[i].wav_write[0]->initFile()) {
              std::cerr << "Unable to open file for writing: " << filename;
              return false;
            }
          }


          if (reset_list[i]) {
            should_reset = true;
            bitmap[i] = NVAFX_TRUE;
          }
        }
      }
    }

    if (should_reset) {
      if (NvAFX_Reset(handle, bitmap.data(), input_wav_list.size()) != NVAFX_STATUS_SUCCESS) {
        std::cerr << "Reset failed" << std::endl;
        return false;
      }
      memset(bitmap.data(), NVAFX_FALSE, sizeof(NvAFX_Bool) * bitmap.size());
    }

    const float* input[2];
    float* output[2];
    input[0] = input_frame.data();
    if (is_aec) {
      input[1] = &input_frame[input_frame.size() / num_input_channels];
    }

    output[0] = output_frame.data();
    if (num_output_channels > 1) {
      output[1] = output_frame.data() + output_frame.size()/2;
    }

    auto start_tick = std::chrono::high_resolution_clock::now();
    if (NvAFX_Run(handle, input, output, num_input_samples_per_frame, num_input_channels) != NVAFX_STATUS_SUCCESS) {
      std::cerr << "NvAFX_Run() failed" << std::endl;
      return false;
    }

    auto run_end_tick = std::chrono::high_resolution_clock::now();
    total_run_time += (std::chrono::duration<float>(run_end_tick - start_tick)).count();
    total_audio_duration += frame_in_secs;

    if ((total_audio_duration / expected_audio_duration) > checkpoint) {
      progress_bar[checkpoint * 10] = '=';
      std::cout << "Processed: " << progress_bar << checkpoint * 100.f << "% ";
      std::cout << (checkpoint >=1 ? "\n" : "\r");
      std::cout.flush();
      checkpoint += 0.1f;
    }

    for (unsigned i = 0; i < num_streams; ++i) {
      unsigned idx = i * num_output_samples_per_frame;
      std::vector<float*> outputs = {};
      unsigned size_per_channel = output_frame.size()/output_channels_;
      for (int ch = 0; ch < output_channels_; ch++) {
        outputs.push_back(&output_frame[idx + ch * size_per_channel]);
      }
      if (!writeOutputWav(outputs, num_output_samples_per_frame, i)) {
        return false;
      }
    }

    if (real_time_) {
      auto end_tick = std::chrono::high_resolution_clock::now();
      std::chrono::duration<float> elapsed = end_tick - start_tick;
      float sleep_time_secs = frame_in_secs - elapsed.count();
      std::this_thread::sleep_for(std::chrono::milliseconds(static_cast<int>(sleep_time_secs * 1000)));
    }
  }

  std::cout << "Processing time " << std::setprecision(2) << total_run_time << " secs for "
            << total_audio_duration << std::setprecision(2) << " secs audio file (" << total_run_time / total_audio_duration
            << " secs processing time per sec of audio)" << std::endl;
  if (real_time_) {
    std::cout << "Note: App ran in real time mode i.e. simulated the input data rate of a mic" << std::endl
              << "'Processing time' could be less then actual run time" << std::endl;
  }

  for (auto& data : stream_data_) {
    for (int ch = 0; ch < output_channels_; ch++) {
      data.wav_write[ch]->commitFile();
    }
  }

  std::cout << "Output wav files written. " << std::endl;
  if (NvAFX_DestroyEffect(handle) != NVAFX_STATUS_SUCCESS) {
    std::cerr << "NvAFX_Release() failed" << std::endl;
    return false;
  }

  if (NvAFX_UninitializeLogger() != NVAFX_STATUS_SUCCESS) {
    std::cerr << "NvAFX_UninitializeLogger() failed" << std::endl;
    return false;
  }

  return true;
}

void ShowHelpAndExit(const char* bad_option) {
  std::ostringstream oss;
  if (bad_option) {
    oss << "Error parsing \"" << bad_option << "\"" << std::endl;
  }
  std::cout << "Command Line Options:" << std::endl
            << "-c           Config file" << std::endl;

  std::cout << oss.str();
  exit(0);
}

void ParseCommandLine(int argc, char* argv[], std::string* config_file) {
  if (argc == 1) {
    ShowHelpAndExit(nullptr);
  }

  for (int i = 1; i < argc; i++) {
    if (!strcasecmp(argv[i], "-h")) {
      ShowHelpAndExit(nullptr);
    }
    if (!strcasecmp(argv[i], "-c")) {
      if (++i == argc || !config_file->empty()) {
        ShowHelpAndExit("-f");
      }
      config_file->assign(argv[i]);
      continue;
    }

    ShowHelpAndExit(argv[i]);
  }
}

int main(int argc, char *argv[]) {
  std::string config_file;
  ParseCommandLine(argc, argv, &config_file);

  ConfigReader config_reader;
  if (config_reader.Load(config_file) == false) {
    std::cerr << "Config file load failed" << std::endl;
    return -1;
  }

  EffectsDemoApp app;
  if (app.run(config_reader)) { return 0; }
  else { return -1; }
}