#ifndef PY_CONTAINER_HPP

#define PY_CONTAINER_HPP

#include <cmath>

#include <functional>

#include <numeric>

#include <sstream>

#include "pybind11/complex.h"

#include "pybind11/pybind11.h"

#include "pybind11/numpy.h"

#ifndef FORCE_IMPORT_ARRAY

#define NO_IMPORT_ARRAY

#endif

#ifndef PY_ARRAY_UNIQUE_SYMBOL

#define PY_ARRAY_UNIQUE_SYMBOL xtensor_python_ARRAY_API

#endif

#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION

#include "numpy/arrayobject.h"

#undef copysign

#include <cmath>

#include "xtensor/containers/xcontainer.hpp"

#include "xtl/xsequence.hpp"

namespace xt

{

inline void import_numpy();

/**

template <class D>

class pycontainer : public pybind11::object,

public xcontainer<D>

{

public:

using derived_type = D;

using base_type = xcontainer<D>;

using inner_types = xcontainer_inner_types<D>;

using storage_type = typename inner_types::storage_type;

using value_type = typename storage_type::value_type;

using reference = typename storage_type::reference;

using const_reference = typename storage_type::const_reference;

using pointer = typename storage_type::pointer;

using const_pointer = typename storage_type::const_pointer;

using size_type = typename storage_type::size_type;

using difference_type = typename storage_type::difference_type;

using shape_type = typename inner_types::shape_type;

using strides_type = typename inner_types::strides_type;

using backstrides_type = typename inner_types::backstrides_type;

using inner_shape_type = typename inner_types::inner_shape_type;

using inner_strides_type = typename inner_types::inner_strides_type;

using iterable_base = xcontainer<D>;

using iterator = typename iterable_base::iterator;

using const_iterator = typename iterable_base::const_iterator;

using stepper = typename iterable_base::stepper;

using const_stepper = typename iterable_base::const_stepper;

template <class S = shape_type>

void resize(const S& shape);

template <class S = shape_type>

void resize(const S& shape, layout_type l);

template <class S = shape_type>

void resize(const S& shape, const strides_type& strides);

template <class S = shape_type>

auto& reshape(S&& shape, layout_type layout = base_type::static_layout) &;

layout_type layout() const;

bool is_contiguous() const noexcept;

using base_type::operator();

using base_type::operator[];

using base_type::begin;

using base_type::end;

protected:

pycontainer();

~pycontainer() = default;

pycontainer(pybind11::handle h, borrowed_t);

pycontainer(pybind11::handle h, stolen_t);

pycontainer(const pybind11::object& o);

pycontainer(const pycontainer&) = default;

pycontainer& operator=(const pycontainer&) = default;

pycontainer(pycontainer&&) = default;

pycontainer& operator=(pycontainer&&) = default;

static derived_type ensure(pybind11::handle h);

static bool check_(pybind11::handle h);

static PyObject* raw_array_t(PyObject* ptr);

derived_type& derived_cast();

const derived_type& derived_cast() const;

PyArrayObject* python_array() const;

size_type get_buffer_size() const;

private:

#if (PYBIND11_VERSION_MAJOR == 2 && PYBIND11_VERSION_MINOR >= 3) || PYBIND11_VERSION_MAJOR >= 3

// Prevent ambiguous overload resolution for operators defined for

// both xt::xcontainer and pybind11::object.

using pybind11::object::operator~;

using pybind11::object::operator+;

using pybind11::object::operator-;

using pybind11::object::operator*;

using pybind11::object::operator/;

using pybind11::object::operator|;

using pybind11::object::operator&;

using pybind11::object::operator^;

using pybind11::object::operator<<;

using pybind11::object::operator>>;

#endif

};

namespace detail

{

template <class T, class E = void>

struct numpy_traits;

template <class T>

struct numpy_traits<T, std::enable_if_t<pybind11::detail::satisfies_any_of<T, std::is_arithmetic, xtl::is_complex>::value>>

{

private:

// On Windows 64 bits, NPY_INT != NPY_INT32 and NPY_UINT != NPY_UINT32

// We use the NPY_INT32 and NPY_UINT32 which are consistent with the values

// of NPY_LONG and NPY_ULONG

// On Linux x64, NPY_INT64 != NPY_LONGLONG and NPY_UINT64 != NPY_ULONGLONG,

// we use the values of NPY_INT64 and NPY_UINT64 which are consistent with the

// values of NPY_LONG and NPY_ULONG.

constexpr static const int value_list[15] = {

NPY_BOOL,

NPY_BYTE, NPY_UBYTE, NPY_SHORT, NPY_USHORT,

NPY_INT32, NPY_UINT32, NPY_INT64, NPY_UINT64,

NPY_FLOAT, NPY_DOUBLE, NPY_LONGDOUBLE,

NPY_CFLOAT, NPY_CDOUBLE, NPY_CLONGDOUBLE};

public:

using value_type = std::remove_const_t<T>;

static constexpr int type_num = value_list[pybind11::detail::is_fmt_numeric<value_type>::index];

};

// On Linux x64, NPY_INT64 != NPY_LONGLONG and NPY_UINT64 != NPY_ULONGLONG

// NPY_LONGLONG and NPY_ULONGLONG must be adjusted so the right type is

// selected

template <bool>

struct numpy_enum_adjuster

{

static inline int pyarray_type(PyArrayObject* obj)

{

return PyArray_TYPE(obj);

}

};

template <>

struct numpy_enum_adjuster<true>

{

static inline int pyarray_type(PyArrayObject* obj)

{

int res = PyArray_TYPE(obj);

if(res == NPY_LONGLONG || res == NPY_ULONGLONG)

{

res -= 2;

}

return res;

}

};

inline int pyarray_type(PyArrayObject* obj)

{

return numpy_enum_adjuster<NPY_LONGLONG != NPY_INT64>::pyarray_type(obj);

}

template <class T>

void default_initialize_impl(T& /*storage*/, std::false_type)

{

}

template <class T>

void default_initialize_impl(T& storage, std::true_type)

{

using value_type = typename T::value_type;

storage[0] = value_type{};

}

template <class T>

void default_initialize(T& storage)

{

using value_type = typename T::value_type;

default_initialize_impl(storage, std::is_copy_assignable<value_type>());

}

template <class T>

bool check_array_type(const pybind11::handle& src, std::true_type)

{

int type_num = xt::detail::numpy_traits<T>::type_num;

return xt::detail::pyarray_type(reinterpret_cast<PyArrayObject*>(src.ptr())) == type_num;

}

template <class T>

bool check_array_type(const pybind11::handle& src, std::false_type)

{

return PyArray_EquivTypes((PyArray_Descr*) pybind11::detail::array_proxy(src.ptr())->descr,

(PyArray_Descr*) pybind11::dtype::of<T>().ptr());

}

template <class T>

bool check_array(const pybind11::handle& src)

{

using is_arithmetic_type = std::integral_constant<bool, bool(pybind11::detail::satisfies_any_of<T, std::is_arithmetic, xtl::is_complex>::value)>;

return PyArray_Check(src.ptr()) && check_array_type<T>(src, is_arithmetic_type{});

}

}

/******************************

template <class D>

inline pycontainer<D>::pycontainer()

: pybind11::object()

{

}

template <class D>

inline pycontainer<D>::pycontainer(pybind11::handle h, borrowed_t b)

: pybind11::object(h, b)

{

}

template <class D>

inline pycontainer<D>::pycontainer(pybind11::handle h, stolen_t s)

: pybind11::object(h, s)

{

}

template <class D>

inline pycontainer<D>::pycontainer(const pybind11::object& o)

: pybind11::object(raw_array_t(o.ptr()), pybind11::object::stolen_t{})

{

if (!this->m_ptr)

{

throw pybind11::error_already_set();

}

}

template <class D>

inline auto pycontainer<D>::ensure(pybind11::handle h) -> derived_type

{

auto result = pybind11::reinterpret_steal<derived_type>(raw_array_t(h.ptr()));

if (result.ptr() == nullptr)

{

PyErr_Clear();

}

return result;

}

template <class D>

inline bool pycontainer<D>::check_(pybind11::handle h)

{

return detail::check_array<typename D::value_type>(h);

}

template <class D>

inline PyObject* pycontainer<D>::raw_array_t(PyObject* ptr)

{

if (ptr == nullptr)

{

return nullptr;

}

auto dtype = pybind11::detail::npy_format_descriptor<value_type>::dtype();

auto res = PyArray_FromAny(ptr, (PyArray_Descr *) dtype.release().ptr(), 0, 0,

NPY_ARRAY_ENSUREARRAY | NPY_ARRAY_FORCECAST, nullptr);

return res;

}

template <class D>

inline PyArrayObject* pycontainer<D>::python_array() const

{

return reinterpret_cast<PyArrayObject*>(this->m_ptr);

}

template <class D>

inline auto pycontainer<D>::get_buffer_size() const -> size_type

{

const size_type& (*min)(const size_type&, const size_type&) = std::min<size_type>;

size_type min_stride = this->strides().empty() ? size_type(1) :

std::max(size_type(1), std::accumulate(this->strides().cbegin(),

this->strides().cend(),

std::numeric_limits<size_type>::max(),

min));

return min_stride * static_cast<size_type>(PyArray_SIZE(this->python_array()));

}

template <class D>

inline auto pycontainer<D>::derived_cast() -> derived_type&

{

return *static_cast<derived_type*>(this);

}

template <class D>

inline auto pycontainer<D>::derived_cast() const -> const derived_type&

{

return *static_cast<const derived_type*>(this);

}

namespace detail

{

template <class S>

struct check_dims

{

static bool run(std::size_t)

{

return true;

}

};

template <class T, std::size_t N>

struct check_dims<std::array<T, N>>

{

static bool run(std::size_t new_dim)

{

if(new_dim != N)

{

std::ostringstream err_msg;

err_msg << "Invalid conversion to pycontainer, expecting a container of dimension "

<< N << ", got a container of dimension " << new_dim << ".";

throw std::runtime_error(err_msg.str());

}

return new_dim == N;

}

};

}

/**

template <class D>

template <class S>

inline void pycontainer<D>::resize(const S& shape)

{

if (shape.size() != this->dimension() || !std::equal(std::begin(shape), std::end(shape), std::begin(this->shape())))

{

resize(shape, layout_type::row_major);

}

}

/**

template <class D>

template <class S>

inline void pycontainer<D>::resize(const S& shape, layout_type l)

{

strides_type strides = xtl::make_sequence<strides_type>(shape.size(), size_type(1));

compute_strides(shape, l, strides);

resize(shape, strides);

}

/**

template <class D>

template <class S>

inline void pycontainer<D>::resize(const S& shape, const strides_type& strides)

{

detail::check_dims<shape_type>::run(shape.size());

derived_type tmp(xtl::forward_sequence<shape_type, decltype(shape)>(shape), strides);

*static_cast<derived_type*>(this) = std::move(tmp);

}

template <class D>

template <class S>

inline auto& pycontainer<D>::reshape(S&& shape, layout_type layout) &

{

if (compute_size(shape) != this->size())

{

throw std::runtime_error("Cannot reshape with incorrect number of elements (" + std::to_string(this->size()) + " vs " + std::to_string(compute_size(shape)) + ")");

}

detail::check_dims<shape_type>::run(shape.size());

layout = default_assignable_layout(layout);

NPY_ORDER npy_layout;

if (layout == layout_type::row_major)

{

npy_layout = NPY_CORDER;

}

else if (layout == layout_type::column_major)

{

npy_layout = NPY_FORTRANORDER;

}

else

{

throw std::runtime_error("Cannot reshape with unknown layout_type.");

}

using shape_ptr = typename std::decay_t<S>::pointer;

PyArray_Dims dims = {reinterpret_cast<npy_intp*>(const_cast<shape_ptr>(shape.data())), static_cast<int>(shape.size())};

auto new_ptr = PyArray_Newshape((PyArrayObject*) this->ptr(), &dims, npy_layout);

auto old_ptr = this->ptr();

this->ptr() = new_ptr;

Py_XDECREF(old_ptr);

this->derived_cast().init_from_python();

return *this;

}

/**

template <class D>

inline layout_type pycontainer<D>::layout() const

{

if (PyArray_CHKFLAGS(python_array(), NPY_ARRAY_C_CONTIGUOUS))

{

return layout_type::row_major;

}

else if (PyArray_CHKFLAGS(python_array(), NPY_ARRAY_F_CONTIGUOUS))

{

return layout_type::column_major;

}

else

{

return layout_type::dynamic;

}

}

/**

template <class D>

inline bool pycontainer<D>::is_contiguous() const noexcept

{

if (this->strides().size() == 0)

{

return true;

}

else if (PyArray_CHKFLAGS(python_array(), NPY_ARRAY_C_CONTIGUOUS))

{

return 1 == this->strides().back();

}

else if (PyArray_CHKFLAGS(python_array(), NPY_ARRAY_F_CONTIGUOUS))

{

return 1 == this->strides().front();

}

else

{

return false;

}

}

/**

inline void import_numpy()

{

#ifdef FORCE_IMPORT_ARRAY

if (_import_array() < 0)

{

PyErr_Print();

PyErr_SetString(PyExc_ImportError, "numpy.core.multiarray failed to import");

}

#endif

}

#if defined(__GNUC__) && !defined(__clang__)

namespace workaround

{

// Fixes "undefined symbol" issues

inline void long_long_allocator()

{

std::allocator<long long> a;

std::allocator<unsigned long long> b;

std::allocator<double> c;

std::allocator<std::complex<double>> d;

}

}

#endif

}

#endif