#define Py_LIMITED_API
#include <Python.h>
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include <numpy/arrayobject.h>

#include <vector>
#include <algorithm>


static PyObject* quantize_quantize(PyObject *self, PyObject *args);

static PyMethodDef quantize_methods[] = {
    {"quantize", quantize_quantize, METH_VARARGS, "Add to numbers."},
    {nullptr, nullptr, 0, nullptr} /* Sentinel */
};

static PyModuleDef quantize_module = {
    PyModuleDef_HEAD_INIT,
    "quantize",         /* name of module */
    "Simple module",    /* module documentation */
    -1,                 /* irrelevant */
    quantize_methods
};

// extern "C" PyObject* PyInit_quantize()
PyMODINIT_FUNC PyInit_quantize()
{
    return PyModule_Create(&quantize_module);
}

template<class T>
static std::vector<T> calculate_boundaries(T *array, size_t size, size_t quantized_quantity) {
    std::vector<T> input(size);
    for (size_t i = 0; i < size; ++i) {
        input[i] = array[i];
    }
    std::sort(input.begin(), input.end());
    std::vector<T> boundaries(quantized_quantity);
    for (size_t i = 0; i < quantized_quantity; ++i) {
        boundaries[i] = input[i*size / quantized_quantity];
    }
    return boundaries;
}

template<class T>
static std::vector<int64_t> calculate_quantized(T*array, size_t size, const std::vector<T>& boundaries) {
    std::vector<int64_t> quantized(size);
    for (size_t i = 0; i < size; ++i) {
        auto it = std::lower_bound(boundaries.begin(), boundaries.end(), array[i]);
        quantized[i] = boundaries.begin() - it;
    }
    return quantized;
}

template<class T>
static PyObject * convert_to_py_object(const std::vector<T>& boundaries,
                                       int type,
                                       const std::vector<int64_t>& quantized) {
    auto boundaries_size = (npy_intp)boundaries.size();
    auto quantize_size   = (npy_intp)quantized.size();

    auto boundaries_data = new T [boundaries.size()];
    auto quantize_data   = new int64_t [quantized.size()];

    std::memcpy(boundaries_data, &boundaries[0], boundaries.size() * sizeof(T));
    std::memcpy(quantize_data,   &quantized[0],  quantized.size() * sizeof(int64_t));

    auto py_boundaries = (PyArrayObject *) PyArray_SimpleNewFromData(1,
                                                                     &boundaries_size,
                                                                     type,
                                                                     boundaries_data);
    auto py_quantized  = (PyArrayObject *) PyArray_SimpleNewFromData(1,
                                                                     &quantize_size,
                                                                     NPY_INT64,
                                                                     quantize_data);
    return PyTuple_Pack(2, py_boundaries, py_quantized);
}

static PyObject* quantize_quantize(PyObject *self, PyObject *args)
{
    if (!self || !args) {
        PyErr_SetString(PyExc_RuntimeError, "bad self or args\n");
        return PyErr_Occurred();
    }

    PyArrayObject *input_arr = nullptr;
    int arg2 = 0;

    if (!PyArg_ParseTuple(args, "O!i", &PyArray_Type, &input_arr, &arg2)) {
        PyErr_SetString(PyExc_RuntimeError, "bad type or number of args\n");
        return PyErr_Occurred();
    }

    if (PyArray_NDIM(input_arr) != 1) {
        PyErr_SetString(PyExc_RuntimeError, "bad dimension of array\n");
        return PyErr_Occurred();
    }

    int type = PyArray_TYPE(input_arr);
    auto size = (size_t)PyArray_SHAPE(input_arr)[0];
    auto quantize_quantity = (size_t)arg2;

    switch (type) {
        case NPY_FLOAT:
        case NPY_DOUBLE: {
            auto array = (double *) PyArray_DATA(input_arr);
            const auto boundaries = calculate_boundaries<double>(array, size, quantize_quantity);
            const auto quantized  = calculate_quantized<double>( array, size, boundaries);
            return convert_to_py_object<double>(boundaries, NPY_DOUBLE, quantized);
        }
        default:
            PyErr_SetString(PyExc_RuntimeError, "unknown type elements in array\n");
            return PyErr_Occurred();
    }
}