// https://www.boost.org/doc/libs/1_92_0/libs/python/doc/html/numpy/index.html

#include <boost/python/numpy.hpp>

namespace p = boost::python;
namespace np = boost::python::numpy;

// Return a 1-D array [0, 1, ..., n-1] built on the C++ side
np::ndarray arange(int n)
{
    p::tuple shape = p::make_tuple(n);
    np::dtype dt = np::dtype::get_builtin<double>();
    np::ndarray a = np::zeros(shape, dt);

    for (int i = 0; i < n; ++i)
    {
        a[i] = static_cast<double>(i);
    }

    return a;
}

// Sum the elements of an array handed over from Python
double total(const np::ndarray &a)
{
    if (a.get_dtype() != np::dtype::get_builtin<double>())
    {
        PyErr_SetString(PyExc_TypeError, "expected an array of doubles");
        p::throw_error_already_set();
    }

    double sum = 0;
    const int n = a.shape(0);
    for (int i = 0; i < n; ++i)
    {
        sum += p::extract<double>(a[i]);
    }

    return sum;
}

BOOST_PYTHON_MODULE(libdemo1)
{
    np::initialize();
    p::def("arange", arange);
    p::def("total", total);
}
