/***************************************************************************
* Copyright (c) Wolf Vollprecht, Johan Mabille and Sylvain Corlay          *
* Copyright (c) QuantStack                                                 *
*                                                                          *
* Distributed under the terms of the BSD 3-Clause License.                 *
*                                                                          *
* The full license is in the file LICENSE, distributed with this software. *
****************************************************************************/

#ifndef PY_ARRAY_HPP
#define PY_ARRAY_HPP

#include <algorithm>
#include <cstddef>
#include <vector>

#include "xtensor/xbuffer_adaptor.hpp"
#include "xtensor/xiterator.hpp"
#include "xtensor/xsemantic.hpp"

#include "pyarray_backstrides.hpp"
#include "pycontainer.hpp"
#include "pystrides_adaptor.hpp"
#include "xtensor_type_caster_base.hpp"

namespace xt
{
    template <class T, layout_type L = layout_type::dynamic>
    class pyarray;
}

namespace pybind11
{
    namespace detail
    {
#ifdef PYBIND11_DESCR // The macro is removed from pybind11 since 2.3
        template <class T, xt::layout_type L>
        struct handle_type_name<xt::pyarray<T, L>>
        {
            static PYBIND11_DESCR name()
            {
                return _("numpy.ndarray[") + npy_format_descriptor<T>::name() + _("]");
            }
        };
#endif

        template <typename T, xt::layout_type L>
        struct pyobject_caster<xt::pyarray<T, L>>
        {
            using type = xt::pyarray<T, L>;

            bool load(handle src, bool convert)
            {
                if (!convert)
                {
                    if (!xt::detail::check_array<T>(src))
                    {
                        return false;
                    }
                }
                value = type::ensure(src);
                return static_cast<bool>(value);
            }

            static handle cast(const handle& src, return_value_policy, handle)
            {
                return src.inc_ref();
            }

#ifdef PYBIND11_DESCR // The macro is removed from pybind11 since 2.3
            PYBIND11_TYPE_CASTER(type, handle_type_name<type>::name());
#else
            PYBIND11_TYPE_CASTER(type, _("numpy.ndarray[") + npy_format_descriptor<T>::name + _("]"));
#endif
        };

        // Type caster for casting ndarray to xexpression<pyarray>
        template <typename T, xt::layout_type L>
        struct type_caster<xt::xexpression<xt::pyarray<T, L>>> : pyobject_caster<xt::pyarray<T, L>>
        {
            using Type = xt::xexpression<xt::pyarray<T, L>>;

            operator Type&()
            {
                return this->value;
            }

            operator const Type&()
            {
                return this->value;
            }
        };

        // Type caster for casting xarray to ndarray
        template <class T, xt::layout_type L>
        struct type_caster<xt::xarray<T, L>> : xtensor_type_caster_base<xt::xarray<T, L>>
        {
        };
    }
}

namespace xt
{
    template <class T, layout_type L>
    struct xiterable_inner_types<pyarray<T, L>>
        : xcontainer_iterable_types<pyarray<T, L>>
    {
    };

    template <class T, layout_type L>
    struct xcontainer_inner_types<pyarray<T, L>>
    {
        using storage_type = xbuffer_adaptor<T*>;
        using reference = typename storage_type::reference;
        using const_reference = typename storage_type::const_reference;
        using size_type = typename storage_type::size_type;
        using shape_type = std::vector<typename storage_type::size_type>;
        using strides_type = std::vector<typename storage_type::difference_type>;
        using backstrides_type = pyarray_backstrides<pyarray<T, L>>;
        using inner_shape_type = xbuffer_adaptor<std::size_t*>;
        using inner_strides_type = pystrides_adaptor<sizeof(T)>;
        using inner_backstrides_type = backstrides_type;
        using temporary_type = pyarray<T, L>;
        static constexpr layout_type layout = L;
    };

    /**
     * @class pyarray
     * @brief Multidimensional container providing the xtensor container semantics to a numpy array.
     *
     * pyarray is similar to the xarray container in that it has a dynamic dimensionality.
     * Reshapes of a pyarray container are reflected in the underlying numpy array.
     *
     * @tparam T The type of the element stored in the pyarray.
     * @tparam L Static layout of the pyarray
     *
     * @sa pytensor
     */
    template <class T, layout_type L>
    class pyarray : public pycontainer<pyarray<T, L>>,
                    public xcontainer_semantic<pyarray<T, L>>
    {
    public:

        using self_type = pyarray<T, L>;
        using semantic_base = xcontainer_semantic<self_type>;
        using base_type = pycontainer<self_type>;
        using storage_type = typename base_type::storage_type;
        using value_type = typename base_type::value_type;
        using reference = typename base_type::reference;
        using const_reference = typename base_type::const_reference;
        using pointer = typename base_type::pointer;
        using size_type = typename base_type::size_type;
        using difference_type = typename base_type::difference_type;
        using shape_type = typename base_type::shape_type;
        using strides_type = typename base_type::strides_type;
        using backstrides_type = typename base_type::backstrides_type;
        using inner_shape_type = typename base_type::inner_shape_type;
        using inner_strides_type = typename base_type::inner_strides_type;
        using inner_backstrides_type = typename base_type::inner_backstrides_type;

        pyarray();
        pyarray(const value_type& t);
        pyarray(nested_initializer_list_t<T, 1> t);
        pyarray(nested_initializer_list_t<T, 2> t);
        pyarray(nested_initializer_list_t<T, 3> t);
        pyarray(nested_initializer_list_t<T, 4> t);
        pyarray(nested_initializer_list_t<T, 5> t);

        pyarray(pybind11::handle h, pybind11::object::borrowed_t);
        pyarray(pybind11::handle h, pybind11::object::stolen_t);
        pyarray(const pybind11::object& o);

        explicit pyarray(const shape_type& shape, layout_type l = layout_type::row_major);
        explicit pyarray(const shape_type& shape, const_reference value, layout_type l = layout_type::row_major);
        explicit pyarray(const shape_type& shape, const strides_type& strides, const_reference value);
        explicit pyarray(const shape_type& shape, const strides_type& strides);

        template <class S = shape_type>
        static pyarray from_shape(S&& s);

        pyarray(const self_type& rhs);
        self_type& operator=(const self_type& rhs);

        pyarray(self_type&&) = default;
        self_type& operator=(self_type&& e) = default;

        template <class E>
        pyarray(const xexpression<E>& e);

        template <class E>
        self_type& operator=(const xexpression<E>& e);

        using base_type::begin;
        using base_type::end;

        static self_type ensure(pybind11::handle h);
        static bool check_(pybind11::handle h);

#if PYBIND11_VERSION_MAJOR == 2 && PYBIND11_VERSION_MINOR >= 3
        // Prevent ambiguous overload resolution for operators defined for
        // both xt::xcontainer_semantic and pybind11::object.
        using semantic_base::operator+=;
        using semantic_base::operator-=;
        using semantic_base::operator*=;
        using semantic_base::operator/=;
        using semantic_base::operator|=;
        using semantic_base::operator&=;
        using semantic_base::operator^=;
        // using semantic_base::operator<<=;
        // using semantic_base::operator>>=;
#endif

    private:

        inner_shape_type m_shape;
        inner_strides_type m_strides;
        mutable inner_backstrides_type m_backstrides;
        storage_type m_storage;

        void init_array(const shape_type& shape, const strides_type& strides);
        void init_from_python();

        const inner_shape_type& shape_impl() const noexcept;
        const inner_strides_type& strides_impl() const noexcept;
        const inner_backstrides_type& backstrides_impl() const noexcept;

        storage_type& storage_impl() noexcept;
        const storage_type& storage_impl() const noexcept;

        friend class xcontainer<pyarray<T, L>>;
        friend class pycontainer<pyarray<T, L>>;
    };

    /**************************
     * pyarray implementation *
     **************************/

    /**
     * @name Constructors
     */
    //@{
    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray()
        : base_type()
    {
        // TODO: avoid allocation
        shape_type shape = xtl::make_sequence<shape_type>(0, size_type(1));
        strides_type strides = xtl::make_sequence<strides_type>(0, size_type(0));
        init_array(shape, strides);
        detail::default_initialize(m_storage);
    }

    /**
     * Allocates a pyarray with nested initializer lists.
     */
    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(const value_type& t)
        : base_type()
    {
        base_type::resize(xt::shape<shape_type>(t), layout_type::row_major);
        nested_copy(m_storage.begin(), t);
    }

    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(nested_initializer_list_t<T, 1> t)
        : base_type()
    {
        base_type::resize(xt::shape<shape_type>(t), layout_type::row_major);
        nested_copy(m_storage.begin(), t);
    }

    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(nested_initializer_list_t<T, 2> t)
        : base_type()
    {
        base_type::resize(xt::shape<shape_type>(t), layout_type::row_major);
        nested_copy(m_storage.begin(), t);
    }

    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(nested_initializer_list_t<T, 3> t)
        : base_type()
    {
        base_type::resize(xt::shape<shape_type>(t), layout_type::row_major);
        nested_copy(m_storage.begin(), t);
    }

    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(nested_initializer_list_t<T, 4> t)
        : base_type()
    {
        base_type::resize(xt::shape<shape_type>(t), layout_type::row_major);
        nested_copy(m_storage.begin(), t);
    }

    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(nested_initializer_list_t<T, 5> t)
        : base_type()
    {
        base_type::resize(xt::shape<shape_type>(t), layout_type::row_major);
        nested_copy(m_storage.begin(), t);
    }

    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(pybind11::handle h, pybind11::object::borrowed_t b)
        : base_type(h, b)
    {
        init_from_python();
    }

    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(pybind11::handle h, pybind11::object::stolen_t s)
        : base_type(h, s)
    {
        init_from_python();
    }

    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(const pybind11::object& o)
        : base_type(o)
    {
        init_from_python();
    }

    /**
     * Allocates an uninitialized pyarray with the specified shape and
     * layout.
     * @param shape the shape of the pyarray
     * @param l the layout of the pyarray
     */
    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(const shape_type& shape, layout_type l)
        : base_type()
    {
        strides_type strides(shape.size());
        compute_strides(shape, l, strides);
        init_array(shape, strides);
    }

    /**
     * Allocates a pyarray with the specified shape and layout. Elements
     * are initialized to the specified value.
     * @param shape the shape of the pyarray
     * @param value the value of the elements
     * @param l the layout of the pyarray
     */
    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(const shape_type& shape, const_reference value, layout_type l)
        : base_type()
    {
        strides_type strides(shape.size());
        compute_strides(shape, l, strides);
        init_array(shape, strides);
        std::fill(m_storage.begin(), m_storage.end(), value);
    }

    /**
     * Allocates an uninitialized pyarray with the specified shape and strides.
     * Elements are initialized to the specified value.
     * @param shape the shape of the pyarray
     * @param strides the strides of the pyarray
     * @param value the value of the elements
     */
    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(const shape_type& shape, const strides_type& strides, const_reference value)
        : base_type()
    {
        init_array(shape, strides);
        std::fill(m_storage.begin(), m_storage.end(), value);
    }

    /**
     * Allocates an uninitialized pyarray with the specified shape and strides.
     * @param shape the shape of the pyarray
     * @param strides the strides of the pyarray
     */
    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(const shape_type& shape, const strides_type& strides)
        : base_type()
    {
        init_array(shape, strides);
    }

    /**
     * Allocates and returns an pyarray with the specified shape.
     * @param shape the shape of the pyarray
     */
    template <class T, layout_type L>
    template <class S>
    inline pyarray<T, L> pyarray<T, L>::from_shape(S&& shape)
    {
        auto shp = xtl::forward_sequence<shape_type, S>(shape);
        return self_type(shp);
    }
    //@}

    /**
     * @name Copy semantic
     */
    //@{
    /**
     * The copy constructor.
     */
    template <class T, layout_type L>
    inline pyarray<T, L>::pyarray(const self_type& rhs)
        : base_type(), semantic_base(rhs)
    {
        auto tmp = pybind11::reinterpret_steal<pybind11::object>(
            PyArray_NewLikeArray(rhs.python_array(), NPY_KEEPORDER, nullptr, 1));

        if (!tmp)
        {
            throw std::runtime_error("NumPy: unable to create ndarray");
        }

        this->m_ptr = tmp.release().ptr();
        init_from_python();
        std::copy(rhs.storage().cbegin(), rhs.storage().cend(), this->storage().begin());
    }

    /**
     * The assignment operator.
     */
    template <class T, layout_type L>
    inline auto pyarray<T, L>::operator=(const self_type& rhs) -> self_type&
    {
        self_type tmp(rhs);
        *this = std::move(tmp);
        return *this;
    }

    //@}

    /**
     * @name Extended copy semantic
     */
    //@{
    /**
     * The extended copy constructor.
     */
    template <class T, layout_type L>
    template <class E>
    inline pyarray<T, L>::pyarray(const xexpression<E>& e)
        : base_type()
    {
        // TODO: prevent intermediary shape allocation
        shape_type shape = xtl::forward_sequence<shape_type, decltype(e.derived_cast().shape())>(e.derived_cast().shape());
        strides_type strides = xtl::make_sequence<strides_type>(shape.size(), size_type(0));
        compute_strides(shape, layout_type::row_major, strides);
        init_array(shape, strides);
        semantic_base::assign(e);
    }

    /**
     * The extended assignment operator.
     */
    template <class T, layout_type L>
    template <class E>
    inline auto pyarray<T, L>::operator=(const xexpression<E>& e) -> self_type&
    {
        return semantic_base::operator=(e);
    }
    //@}

    template <class T, layout_type L>
    inline auto pyarray<T, L>::ensure(pybind11::handle h) -> self_type
    {
        return base_type::ensure(h);
    }

    template <class T, layout_type L>
    inline bool pyarray<T, L>::check_(pybind11::handle h)
    {
        return base_type::check_(h);
    }

    template <class T, layout_type L>
    inline void pyarray<T, L>::init_array(const shape_type& shape, const strides_type& strides)
    {
        strides_type adapted_strides(strides);

        std::transform(strides.begin(), strides.end(), adapted_strides.begin(),
                       [](auto v) { return sizeof(T) * v; });

        int flags = NPY_ARRAY_ALIGNED;
        if (!std::is_const<T>::value)
        {
            flags |= NPY_ARRAY_WRITEABLE;
        }

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

        npy_intp* shape_data = reinterpret_cast<npy_intp*>(const_cast<size_type*>(shape.data()));
        npy_intp* strides_data = reinterpret_cast<npy_intp*>(adapted_strides.data());

        auto tmp = pybind11::reinterpret_steal<pybind11::object>(
            PyArray_NewFromDescr(&PyArray_Type, (PyArray_Descr*) dtype.release().ptr(), static_cast<int>(shape.size()), shape_data, strides_data,
                        nullptr, flags, nullptr));

        if (!tmp)
        {
            throw std::runtime_error("NumPy: unable to create ndarray");
        }

        this->m_ptr = tmp.release().ptr();
        init_from_python();
    }

    template <class T, layout_type L>
    inline void pyarray<T, L>::init_from_python()
    {
        if (!static_cast<bool>(*this))
        {
            return;
        }
        
        m_shape = inner_shape_type(reinterpret_cast<size_type*>(PyArray_SHAPE(this->python_array())),
                                   static_cast<size_type>(PyArray_NDIM(this->python_array())));
        m_strides = inner_strides_type(reinterpret_cast<difference_type*>(PyArray_STRIDES(this->python_array())),
                                       static_cast<size_type>(PyArray_NDIM(this->python_array())));

        if (L != layout_type::dynamic && !do_strides_match(m_shape, m_strides, L, 1))
        {
            throw std::runtime_error("NumPy: passing container with bad strides for layout (is it a view?).");
        }

        m_backstrides = backstrides_type(*this);
        m_storage = storage_type(reinterpret_cast<pointer>(PyArray_DATA(this->python_array())),
                                 this->get_buffer_size());
    }

    template <class T, layout_type L>
    inline auto pyarray<T, L>::shape_impl() const noexcept -> const inner_shape_type&
    {
        return m_shape;
    }

    template <class T, layout_type L>
    inline auto pyarray<T, L>::strides_impl() const noexcept -> const inner_strides_type&
    {
        return m_strides;
    }

    template <class T, layout_type L>
    inline auto pyarray<T, L>::backstrides_impl() const noexcept -> const inner_backstrides_type&
    {
        // m_backstrides wraps the numpy array backstrides, which is a raw pointer.
        // The address of the raw pointer stored in the wrapper would be invalidated when the pyarray is copied.
        // Hence, we build a new backstrides object (cheap wrapper around the underlying pointer) upon access.
        m_backstrides = backstrides_type(*this);
        return m_backstrides;
    }

    template <class T, layout_type L>
    inline auto pyarray<T, L>::storage_impl() noexcept -> storage_type&
    {
        return m_storage;
    }

    template <class T, layout_type L>
    inline auto pyarray<T, L>::storage_impl() const noexcept -> const storage_type&
    {
        return m_storage;
    }
}

#endif
