#include <boost/python.hpp>
#include <numpy/ndarrayobject.h>
namespace bp = boost::python;
void reference_contiguous_array(PyObject* in,
        PyArrayObject* in_con,
        double* ptr, int& count)
{
    in_con = PyArray_GETCONTIGUOUS((PyArrayObject*)in);
    ptr = (double*)PyArray_DATA(in_con);
    int num_dim = PyArray_NDIM(in_con);
    npy_intp* pdim = PyArray_DIMS(in_con);
    count = 1;
    for (int i = 0; i < num_dim; i++)
    {
        count *= pdim[i];
    }
}
void dereference(PyObject* o)
{
    Py_DECREF(o);
}
PyObject* entry_square_matrix(PyObject* input_matrix)
{
    // get the input array
    double* ptr;
    int count;
    PyArrayObject* input_contigous_array;
    reference_contiguous_array(input_matrix, input_contigous_array, ptr, count);
 
    // create the output array 
    npy_intp dst_dim[1];
    dst_dim[0] = count;    
    PyObject* out_matrix = PyArray_SimpleNew(1, dst_dim, NPY_FLOAT64);
    double* ptr_out;
    PyArrayObject* output_contigous_array;
    reference_contiguous_array(out_matrix, output_contigous_array, ptr_out, count);
    for (int i = 0; i < count; i++)
    {
        ptr_out[i] = ptr[i] * ptr[i];
    }
    dereference((PyObject*)input_contigous_array);
    dereference((PyObject*)output_contigous_array);
    return out_matrix;
}
BOOST_PYTHON_MODULE(_func)
{
    import_array();
    bp::def(square_matrix, entry_square_matrix);
}