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

#include <memory>
#include <functional>
#include <string>
#include <unordered_map>
#include <vector>

#include "pypreproc.h"

namespace py = pybind11;

bool TestPyPreproc(const std::unordered_map<std::string, std::string> &params) {
  auto model = std::make_shared<edk::ModelLoader>("data/test_model.cambricon", "subnet0");
  std::vector<float*> inputs;
  for (uint32_t i = 0; i < model->InputNum(); ++i) {
    inputs.push_back(new float[model->InputShape(i).DataCount()]);
  }
  auto free_inputs = [&inputs] () {
    for (auto ptr : inputs) delete[] ptr;
  };
  cnstream::PyPreproc pypreproc;
  if (!pypreproc.Init(params)) {
    free_inputs();
    return false;
  }
  bool ret = 0 == pypreproc.Execute(inputs, model, nullptr);
  free_inputs();
  return ret;
}

void PreprocTestWrapper(py::module &m) {  // NOLINT
  m.def("cpptest_pypreproc", &TestPyPreproc);
}