// This file contains bindings for the stream functions

// It is expected, that the #include statemnts can raise errors in your IDE.
// The needed files for those imports are only provided during the build
// process.

extern "C" {
#include <topotoolbox.h>
}

#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>

namespace py = pybind11;

// Make wrap_funcname() function available as stream_funcname() to be used by
// by functions in the pytopotoolbox package

void wrap_streamquad_trapz_f32(py::array_t<float> integral,
                               py::array_t<float> integrand,
                               py::array_t<ptrdiff_t> source,
                               py::array_t<ptrdiff_t> target,
                               py::array_t<float> weight) {

  float *integral_ptr = integral.mutable_data();
  float *integrand_ptr = integrand.mutable_data();
  ptrdiff_t *source_ptr = source.mutable_data();
  ptrdiff_t *target_ptr = target.mutable_data();
  float *weight_ptr = weight.mutable_data();

  ptrdiff_t edge_count = source.size();

  streamquad_trapz_f32(integral_ptr, integrand_ptr, source_ptr, target_ptr,
                       weight_ptr, edge_count);
}

void wrap_streamquad_trapz_f64(py::array_t<double> integral,
                               py::array_t<double> integrand,
                               py::array_t<ptrdiff_t> source,
                               py::array_t<ptrdiff_t> target,
                               py::array_t<float> weight) {
  double *integral_ptr = integral.mutable_data();
  double *integrand_ptr = integrand.mutable_data();
  ptrdiff_t *source_ptr = source.mutable_data();
  ptrdiff_t *target_ptr = target.mutable_data();
  float *weight_ptr = weight.mutable_data();

  ptrdiff_t edge_count = source.size();

  streamquad_trapz_f64(integral_ptr, integrand_ptr, source_ptr, target_ptr,
                       weight_ptr, edge_count);
}

void wrap_traverse_up_u32_or_and(py::array_t<uint32_t> output,
                                 py::array_t<uint32_t> input,
                                 py::array_t<ptrdiff_t> source,
                                 py::array_t<ptrdiff_t> target) {

  uint32_t *output_ptr = output.mutable_data();
  uint32_t *input_ptr = input.mutable_data();
  ptrdiff_t *source_ptr = source.mutable_data();
  ptrdiff_t *target_ptr = target.mutable_data();

  ptrdiff_t edge_count = source.size();

  traverse_up_u32_or_and(output_ptr, input_ptr, source_ptr,
                         target_ptr, edge_count);
}

void wrap_traverse_down_u32_or_and(py::array_t<uint32_t> output,
                                   py::array_t<uint32_t> input,
                                   py::array_t<ptrdiff_t> source,
                                   py::array_t<ptrdiff_t> target) {

  uint32_t *output_ptr = output.mutable_data();
  uint32_t *input_ptr = input.mutable_data();
  ptrdiff_t *source_ptr = source.mutable_data();
  ptrdiff_t *target_ptr = target.mutable_data();

  ptrdiff_t edge_count = source.size();

  traverse_down_u32_or_and(output_ptr, input_ptr, source_ptr,
                           target_ptr, edge_count);
}

void wrap_traverse_down_f32_max_add(py::array_t<float> output,
                                    py::array_t<float> input,
                                    py::array_t<ptrdiff_t> source,
                                    py::array_t<ptrdiff_t> target) {

  float *output_ptr = output.mutable_data();
  float *input_ptr = input.mutable_data();
  ptrdiff_t *source_ptr = source.mutable_data();
  ptrdiff_t *target_ptr = target.mutable_data();

  ptrdiff_t edge_count = source.size();

  traverse_down_f32_max_add(output_ptr, input_ptr, source_ptr,
                            target_ptr, edge_count);
}

void wrap_traverse_up_f32_max_add(py::array_t<float> output,
                                  py::array_t<float> input,
                                  py::array_t<ptrdiff_t> source,
                                  py::array_t<ptrdiff_t> target)
{

  float *output_ptr = output.mutable_data();
  float *input_ptr = input.mutable_data();
  ptrdiff_t *source_ptr = source.mutable_data();
  ptrdiff_t *target_ptr = target.mutable_data();

  ptrdiff_t edge_count = source.size();

  traverse_up_f32_max_add(output_ptr, input_ptr, source_ptr,
                          target_ptr, edge_count);
}

void wrap_traverse_down_f32_max_mul(py::array_t<float> output,
                                    py::array_t<float> input,
                                    py::array_t<ptrdiff_t> source,
                                    py::array_t<ptrdiff_t> target) {
  traverse_down_f32_max_mul(output.mutable_data(), input.mutable_data(),
                            source.mutable_data(), target.mutable_data(),
                            source.size());
}

void wrap_traverse_down_f32_max_mul_arg(py::array_t<float> output,
                                        py::array_t<int64_t> idx,
                                        py::array_t<float> input,
                                        py::array_t<ptrdiff_t> source,
                                        py::array_t<ptrdiff_t> target) {
  traverse_down_f32_max_mul_arg(output.mutable_data(), idx.mutable_data(),
                                input.mutable_data(),
                                source.mutable_data(), target.mutable_data(),
                                source.size());
}

void wrap_traverse_up_f32_max_mul(py::array_t<float> output,
                                  py::array_t<float> input,
                                  py::array_t<ptrdiff_t> source,
                                  py::array_t<ptrdiff_t> target) {
  traverse_up_f32_max_mul(output.mutable_data(), input.mutable_data(),
                          source.mutable_data(), target.mutable_data(),
                          source.size());
}

void wrap_traverse_up_f32_max_mul_arg(py::array_t<float> output,
                                      py::array_t<int64_t> idx,
                                      py::array_t<float> input,
                                      py::array_t<ptrdiff_t> source,
                                      py::array_t<ptrdiff_t> target) {
  traverse_up_f32_max_mul_arg(output.mutable_data(), idx.mutable_data(),
                              input.mutable_data(),
                              source.mutable_data(), target.mutable_data(),
                              source.size());
}

void wrap_traverse_down_f32_min_add(py::array_t<float> output,
                                    py::array_t<float> input,
                                    py::array_t<ptrdiff_t> source,
                                    py::array_t<ptrdiff_t> target) {

  float *output_ptr = output.mutable_data();
  float *input_ptr = input.mutable_data();
  ptrdiff_t *source_ptr = source.mutable_data();
  ptrdiff_t *target_ptr = target.mutable_data();

  ptrdiff_t edge_count = source.size();

  traverse_down_f32_min_add(output_ptr, input_ptr, source_ptr,
                            target_ptr, edge_count);
}

void wrap_traverse_down_f32_add_mul(py::array_t<float> output,
                                    py::array_t<float> input,
                                    py::array_t<ptrdiff_t> source,
                                    py::array_t<ptrdiff_t> target) {
  float *output_ptr = output.mutable_data();
  float *input_ptr = input.mutable_data();
  ptrdiff_t *source_ptr = source.mutable_data();
  ptrdiff_t *target_ptr = target.mutable_data();

  ptrdiff_t edge_count = source.size();

  traverse_down_f32_add_mul(output_ptr, input_ptr, source_ptr,
                            target_ptr, edge_count);
}

void wrap_traverse_down_f32_strahler(py::array_t<float> output,
                                     py::array_t<float> input,
                                     py::array_t<ptrdiff_t> source,
                                     py::array_t<ptrdiff_t> target) {
  float *output_ptr = output.mutable_data();
  float *input_ptr = input.mutable_data();
  ptrdiff_t *source_ptr = source.mutable_data();
  ptrdiff_t *target_ptr = target.mutable_data();

  ptrdiff_t edge_count = source.size();

  traverse_down_f32_strahler(output_ptr, input_ptr, source_ptr,
                             target_ptr, edge_count);
}

void wrap_edgelist_degree(py::array_t<uint8_t> indegree,
                          py::array_t<uint8_t> outdegree,
                          py::array_t<ptrdiff_t> source,
                          py::array_t<ptrdiff_t> target) {
  uint8_t *indegree_ptr = indegree.mutable_data();
  uint8_t *outdegree_ptr = outdegree.mutable_data();
  ptrdiff_t *source_ptr = source.mutable_data();
  ptrdiff_t *target_ptr = target.mutable_data();

  ptrdiff_t node_count = indegree.size();
  ptrdiff_t edge_count = source.size();

  edgelist_degree(indegree_ptr, outdegree_ptr, source_ptr,
                  target_ptr, node_count, edge_count);
}

void wrap_propagatevaluesupstream_f32(py::array_t<float> data,
                                      py::array_t<ptrdiff_t> source,
                                      py::array_t<ptrdiff_t> target) {

  propagatevaluesupstream_f32(data.mutable_data(),
                              source.mutable_data(),
                              target.mutable_data(),
                              source.size());
}

void wrap_propagatevaluesupstream_f64(py::array_t<double> data,
                                      py::array_t<ptrdiff_t> source,
                                      py::array_t<ptrdiff_t> target) {

  propagatevaluesupstream_f64(data.mutable_data(),
                              source.mutable_data(),
                              target.mutable_data(),
                              source.size());
}

void wrap_propagatevaluesupstream_u8(py::array_t<uint8_t> data,
                                     py::array_t<ptrdiff_t> source,
                                     py::array_t<ptrdiff_t> target) {

  propagatevaluesupstream_u8(data.mutable_data(),
                             source.mutable_data(),
                             target.mutable_data(),
                             source.size());
}

void wrap_propagatevaluesupstream_u32(py::array_t<uint32_t> data,
                                      py::array_t<ptrdiff_t> source,
                                      py::array_t<ptrdiff_t> target) {

  propagatevaluesupstream_u32(data.mutable_data(),
                              source.mutable_data(),
                              target.mutable_data(),
                              source.size());
}

void wrap_propagatevaluesupstream_u64(py::array_t<uint64_t> data,
                                      py::array_t<ptrdiff_t> source,
                                      py::array_t<ptrdiff_t> target) {

  propagatevaluesupstream_u64(data.mutable_data(),
                              source.mutable_data(),
                              target.mutable_data(),
                              source.size());
}

void wrap_propagatevaluesupstream_i8(py::array_t<int8_t> data,
                                     py::array_t<ptrdiff_t> source,
                                     py::array_t<ptrdiff_t> target) {

  propagatevaluesupstream_i8(data.mutable_data(),
                             source.mutable_data(),
                             target.mutable_data(),
                             source.size());
}

void wrap_propagatevaluesupstream_i32(py::array_t<int32_t> data,
                                      py::array_t<ptrdiff_t> source,
                                      py::array_t<ptrdiff_t> target) {

  propagatevaluesupstream_i32(data.mutable_data(),
                              source.mutable_data(),
                              target.mutable_data(),
                              source.size());
}

void wrap_propagatevaluesupstream_i64(py::array_t<int64_t> data,
                                      py::array_t<ptrdiff_t> source,
                                      py::array_t<ptrdiff_t> target) {

  propagatevaluesupstream_i64(data.mutable_data(),
                              source.mutable_data(),
                              target.mutable_data(),
                              source.size());
}

void wrap_lowerenv(py::array_t<float> elevation,
                   py::array_t<uint8_t> knickpoints,
                   py::array_t<float> distance,
                   py::array_t<ptrdiff_t> ix,
                   py::array_t<uint8_t> onenvelope,
                   py::array_t<ptrdiff_t> source,
                   py::array_t<ptrdiff_t> target) {

  ptrdiff_t node_count = elevation.size();
  ptrdiff_t edge_count = source.size();

  lowerenv(elevation.mutable_data(), knickpoints.mutable_data(), distance.mutable_data(),
           ix.mutable_data(), onenvelope.mutable_data(), source.mutable_data(),
           target.mutable_data(), edge_count, node_count);
}

PYBIND11_MODULE(_stream, m) {
  m.def("streamquad_trapz_f32", &wrap_streamquad_trapz_f32);
  m.def("streamquad_trapz_f64", &wrap_streamquad_trapz_f64);
  m.def("traverse_up_u32_or_and", &wrap_traverse_up_u32_or_and);
  m.def("traverse_down_u32_or_and", &wrap_traverse_down_u32_or_and);
  m.def("traverse_down_f32_max_add", &wrap_traverse_down_f32_max_add);
  m.def("traverse_up_f32_max_add", &wrap_traverse_up_f32_max_add);
  m.def("traverse_down_f32_max_mul", &wrap_traverse_down_f32_max_mul);
  m.def("traverse_down_f32_max_mul_arg", &wrap_traverse_down_f32_max_mul_arg);
  m.def("traverse_up_f32_max_mul", &wrap_traverse_up_f32_max_mul);
  m.def("traverse_up_f32_max_mul_arg", &wrap_traverse_up_f32_max_mul_arg);
  m.def("traverse_down_f32_min_add", &wrap_traverse_down_f32_min_add);
  m.def("traverse_down_f32_add_mul", &wrap_traverse_down_f32_add_mul);
  m.def("traverse_down_f32_strahler", &wrap_traverse_down_f32_strahler);
  m.def("edgelist_degree", &wrap_edgelist_degree);
  m.def("propagatevaluesupstream_f32", &wrap_propagatevaluesupstream_f32);
  m.def("propagatevaluesupstream_f64", &wrap_propagatevaluesupstream_f64);
  m.def("propagatevaluesupstream_u8", &wrap_propagatevaluesupstream_u8);
  m.def("propagatevaluesupstream_u32", &wrap_propagatevaluesupstream_u32);
  m.def("propagatevaluesupstream_u64", &wrap_propagatevaluesupstream_u64);
  m.def("propagatevaluesupstream_i8", &wrap_propagatevaluesupstream_i8);
  m.def("propagatevaluesupstream_i32", &wrap_propagatevaluesupstream_i32);
  m.def("propagatevaluesupstream_i64", &wrap_propagatevaluesupstream_i64);
  m.def("lowerenv", &wrap_lowerenv);
}
