#ifndef TORCHAUDIO_REGISTER_H
#define TORCHAUDIO_REGISTER_H

#include <torchaudio/csrc/sox_effects.h>
#include <torchaudio/csrc/sox_io.h>
#include <torchaudio/csrc/sox_utils.h>

TORCH_LIBRARY(torchaudio, m) {
  //////////////////////////////////////////////////////////////////////////////
  // sox_utils.h
  //////////////////////////////////////////////////////////////////////////////
  m.class_<torchaudio::sox_utils::TensorSignal>("TensorSignal")
      .def(torch::init<torch::Tensor, int64_t, bool>())
      .def("get_tensor", &torchaudio::sox_utils::TensorSignal::getTensor)
      .def(
          "get_sample_rate",
          &torchaudio::sox_utils::TensorSignal::getSampleRate)
      .def(
          "get_channels_first",
          &torchaudio::sox_utils::TensorSignal::getChannelsFirst);

  m.def("torchaudio::sox_utils_set_seed", &torchaudio::sox_utils::set_seed);
  m.def(
      "torchaudio::sox_utils_set_verbosity",
      &torchaudio::sox_utils::set_verbosity);
  m.def(
      "torchaudio::sox_utils_set_use_threads",
      &torchaudio::sox_utils::set_use_threads);
  m.def(
      "torchaudio::sox_utils_set_buffer_size",
      &torchaudio::sox_utils::set_buffer_size);
  m.def(
      "torchaudio::sox_utils_list_effects",
      &torchaudio::sox_utils::list_effects);
  m.def(
      "torchaudio::sox_utils_list_read_formats",
      &torchaudio::sox_utils::list_read_formats);
  m.def(
      "torchaudio::sox_utils_list_write_formats",
      &torchaudio::sox_utils::list_write_formats);

  //////////////////////////////////////////////////////////////////////////////
  // sox_io.h
  //////////////////////////////////////////////////////////////////////////////
  m.class_<torchaudio::sox_io::SignalInfo>("SignalInfo")
      .def("get_sample_rate", &torchaudio::sox_io::SignalInfo::getSampleRate)
      .def("get_num_channels", &torchaudio::sox_io::SignalInfo::getNumChannels)
      .def("get_num_frames", &torchaudio::sox_io::SignalInfo::getNumFrames);

  m.def("torchaudio::sox_io_get_info", &torchaudio::sox_io::get_info);
  m.def(
      "torchaudio::sox_io_load_audio_file",
      &torchaudio::sox_io::load_audio_file);
  m.def(
      "torchaudio::sox_io_save_audio_file",
      &torchaudio::sox_io::save_audio_file);

  //////////////////////////////////////////////////////////////////////////////
  // sox_effects.h
  //////////////////////////////////////////////////////////////////////////////
  m.def(
      "torchaudio::sox_effects_initialize_sox_effects",
      &torchaudio::sox_effects::initialize_sox_effects);
  m.def(
      "torchaudio::sox_effects_shutdown_sox_effects",
      &torchaudio::sox_effects::shutdown_sox_effects);
  m.def(
      "torchaudio::sox_effects_apply_effects_tensor",
      &torchaudio::sox_effects::apply_effects_tensor);
  m.def(
      "torchaudio::sox_effects_apply_effects_file",
      &torchaudio::sox_effects::apply_effects_file);
}
#endif
