#include "calligraphy_process.h"

#define PY_SSIZE_T_CLEAN
#include <Python.h>

#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include <numpy/ndarrayobject.h>

class CalligraphyProcess::Impl {
public:
  Impl() {
    Py_Initialize();
    if (!Py_IsInitialized()) {
      throw std::runtime_error("failed to initialize python environment");
    }

    if (PyArray_ImportNumPyAPI()) {
      Py_Finalize();
      throw std::runtime_error("failed to import numpy from python");
    }

    mModule = PyImport_ImportModule("calligraphy_process");
    if (!mModule) {
      Py_Finalize();
      throw std::runtime_error("failed to load calligraphy_process module");
    }

    mFindCalligraphers = PyObject_GetAttrString(mModule, "find_calligraphers");
    assert(mFindCalligraphers);
  }

  ~Impl() {
    Py_XDECREF(mModule);
    Py_XDECREF(mFindCalligraphers);
    Py_Finalize();
  }

  void findCalligraphers(const QImage &image, QImage &imageOut, QList<int> &indices) {
    PyObject *imageArray = imageToArray(image);
    PyObject *result = PyObject_CallOneArg(mFindCalligraphers, imageArray);
    Py_XDECREF(imageArray);

    imageOut = arrayToImage(PyTuple_GetItem(result, 0));
    indices = listToVectorInt(PyTuple_GetItem(result, 1));
    Py_XDECREF(result);
  }

private:
  PyObject *mModule = nullptr;
  PyObject *mFindCalligraphers = nullptr;

  PyObject *imageToArray(const QImage &image) {
    npy_intp dims[] = {image.height(), image.width(), 4};
    void *data = const_cast<uchar *>(image.bits());
    PyObject *array = PyArray_SimpleNewFromData(3, dims, NPY_UINT8, data);
    return array;
  }

  QImage arrayToImage(PyObject *object) {
    PyArrayObject *array = reinterpret_cast<PyArrayObject *>(object);
    PyArrayObject *arrayCont = PyArray_GETCONTIGUOUS(array);

    npy_intp *dims = PyArray_DIMS(arrayCont);
    int width = static_cast<int>(dims[1]);
    int height = static_cast<int>(dims[0]);
    QImage image{width, height, QImage::Format_RGB32};

    memcpy(image.bits(), PyArray_DATA(arrayCont), image.sizeInBytes());
    Py_XDECREF(arrayCont);
    return image;
  }

  QList<int> listToVectorInt(PyObject *list) {
    QList<int> result(PyList_Size(list));
    for (int i = 0; i < result.size(); i++) {
      PyArg_Parse(PyList_GET_ITEM(list, i), "i", &result[i]);
    }
    return result;
  }
};

CalligraphyProcess::CalligraphyProcess() {
  mImpl = std::make_unique<Impl>();
}

CalligraphyProcess::~CalligraphyProcess() {}

void CalligraphyProcess::findCalligraphers(const QImage &image, QImage &imageOut, QList<int> &indices) {
  mImpl->findCalligraphers(image, imageOut, indices);
}
