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

#include <opencv2/opencv.hpp>

#include <iostream>

namespace py = pybind11;

// ============================================================================
int determine_cv_type(pybind11::dtype const& dtype)
{
    switch (dtype.kind()) {
    case 'u':
        switch (dtype.itemsize()) {
        case 1: return CV_8U;
        case 2: return CV_16U;
        default:
            throw std::invalid_argument("Unsupported unsigned integer size.");
        }
    case 'i':
        switch (dtype.itemsize()) {
        case 1: return CV_8S;
        case 2: return CV_16S;
        case 4: return CV_32S;
        default:
            throw std::invalid_argument("Unsupported signed integer size.");
        }
    case 'f':
        switch (dtype.itemsize()) {
        case 4: return CV_32F;
        case 8: return CV_64F;
        default:
            throw std::invalid_argument("Unsupported float size.");
        }
    default:
        throw std::invalid_argument("Unsupported array dtype.");
    }
}
// ----------------------------------------------------------------------------
cv::Mat nparray_to_mat(py::array& arr)
{
    auto const depth = determine_cv_type(arr.dtype());

    switch (arr.ndim()) {
    case 1:
        return cv::Mat(
            static_cast<int>(arr.shape(0)) // Rows
            , 1 // Columns
            , CV_MAKETYPE(depth, 1) // Data type
            , arr.mutable_data()
        );
    case 2:
        return cv::Mat(
            static_cast<int>(arr.shape(0)) // Rows
            , static_cast<int>(arr.shape(1)) // Columns
            , CV_MAKETYPE(depth, 1) // Data type
            , arr.mutable_data()
        );
    case 3:
        if (arr.shape(2) > CV_CN_MAX) {
            std::invalid_argument("Too many array channels.");
        }
        return cv::Mat(
            static_cast<int>(arr.shape(0)) // Rows
            , static_cast<int>(arr.shape(1)) // Columns
            , CV_MAKETYPE(depth, static_cast<int>(arr.shape(2))) // Data type
            , arr.mutable_data()
        );
    default:
        throw std::invalid_argument("Only 1D, 2D and 3D arrays supported.");
    }
}
// ============================================================================
py::dtype determine_np_dtype(int depth)
{
    switch (depth) {
    case CV_8U: return py::dtype::of<uint8_t>();
    case CV_8S: return py::dtype::of<int8_t>();
    case CV_16U: return py::dtype::of<uint16_t>();
    case CV_16S: return py::dtype::of<int16_t>();
    case CV_32S: return py::dtype::of<int32_t>();
    case CV_32F: return py::dtype::of<float>();
    case CV_64F: return py::dtype::of<double>();
    default:
        throw std::invalid_argument("Unsupported data type.");
    }
}
// ----------------------------------------------------------------------------
py::capsule make_capsule(cv::Mat& m)
{
    return py::capsule(new cv::Mat(m)
        , [](void *v) { delete reinterpret_cast<cv::Mat*>(v); }
        );
}
// ----------------------------------------------------------------------------
py::array mat_to_nparray(cv::Mat& m)
{
    if (!m.isContinuous()) {
        throw std::invalid_argument("Only continuous Mats supported.");
    }

    std::vector<std::size_t> shape;

    if (m.channels() == 1) {
        shape = {
            static_cast<size_t>(m.rows)
            , static_cast<size_t>(m.cols)
        };
    } else {
        shape = {
            static_cast<size_t>(m.rows)
            , static_cast<size_t>(m.cols)
            , static_cast<size_t>(m.channels())
        };
    }

    return py::array(determine_np_dtype(m.depth())
        , shape
        , m.data
        , make_capsule(m));
}
// ============================================================================
PYBIND11_EMBEDDED_MODULE(test_module, m)
{
    m.doc() = "Test module";
}
// ============================================================================
int main()
{
    // Start the interpreter and keep it alive
    py::scoped_interpreter guard{};

    try {
        auto locals = py::dict{};

        py::exec(R"(
            import numpy as np

            def test_cpp_to_py(arr):
                return (1,2,3)
        )");

        auto test_cpp_to_py = py::globals()["test_cpp_to_py"];


        for (int i = 0; i < 10; i++) {
            int64 t0 = cv::getTickCount();

            cv::Mat img(cv::Mat::zeros(1024, 1024, CV_8UC3) + cv::Scalar(1, 1, 1));

            int64 t1 = cv::getTickCount();

            auto result = test_cpp_to_py(mat_to_nparray(img));

            int64 t2 = cv::getTickCount();

            double delta0 = (t1 - t0) / cv::getTickFrequency() * 1000;
            double delta1 = (t2 - t1) / cv::getTickFrequency() * 1000;

            std::cout << "* " << delta0 << " ms | " << delta1 << " ms" << std::endl;
        }
    } catch (py::error_already_set& e) {
        std::cerr << e.what() << "\n";
    }

    return 0;
}
// ============================================================================