#ifndef __ND_ARRAY__
#define __ND_ARRAY__

//================================================
// @file         ndArray.h
// @author       Jonathan Hadida
// @contact      cf https://github.com/Sheljohn/MexArray
//================================================

#include <iostream>
#include <cstdlib>
#include <cstdarg>
#include <cstdio>

#include <exception>
#include <stdexcept>

#include <memory>
#include <algorithm>
#include <type_traits>
#include <initializer_list>

// Protect 1D access (comment for faster access)
#define ND_ARRAY_SAFE_ACCESS
#ifdef ND_ARRAY_SAFE_ACCESS
#define ND_ARRAY_PROTECT(k,n) (k % n)
#else
#define ND_ARRAY_PROTECT(k,n) k
#endif



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



/**
 * Convert nd coordinates to 1d index.
 * Two main variants are provided:
 * - Taking an ARRAY as coordinates (size input by template)
 * - Taking a VA_LIST as a list of coordinate inputs (cf operator() below).
 */
template <unsigned N>
unsigned sub2ind( const unsigned *subs, const unsigned *size, const unsigned *strides )
{
    register unsigned ind = 0;

    for (unsigned i = 0; i < N; ++i)
        ind += ND_ARRAY_PROTECT(subs[i],size[i]) * strides[i];

    return ind;
}

template <unsigned N>
unsigned sub2ind( va_list& vl, const unsigned *size, const unsigned *strides )
{
    register unsigned ind = 0;

    for (unsigned i = 1; i < N; ++i)
        ind += ND_ARRAY_PROTECT(va_arg(vl,unsigned),size[i]) * strides[i];

    va_end(vl); return ind;
}

template <> inline unsigned
sub2ind<0>( const unsigned*, const unsigned*, const unsigned* )
    { return 0; }
template <> inline unsigned
sub2ind<1>( const unsigned *subs, const unsigned*, const unsigned* )
    { return *subs; }
template <> inline unsigned
sub2ind<2>( const unsigned *subs, const unsigned *size, const unsigned* )
    { return (subs[0] + subs[1]*size[0]); }

// ------------------------------------------------------------------------

/**
 * Simple singleton.
 */
template <typename T> struct singleton { static T instance; };
template <typename T> T singleton<T>::instance = T();

// ------------------------------------------------------------------------

/**
 * Dummy deleter functor.
 * This litterally does nothing to the input pointer; it can be used
 * safely with shared pointers for either statically allocated memory
 * (eg fixed-size arrays) or externally managed memory (eg Matlab in/out).
 */
template <typename T>
struct no_delete
{ inline void operator() ( T* ptr ) const {} };



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



/**
 * n-dimensional array.
 * NOTE: T can be CONST (underlying elements non-assignable),
 * or NON-CONST (underlying elements assignable, suitable for
 * owned memory allocation for instance).
 */
template <typename T, unsigned N>
class ndArray
{
public:

    typedef T value_type;
    typedef T* pointer;
    typedef T& reference;

    typedef typename std::add_const<T>::type const_value;
    typedef const_value* const_pointer;
    typedef const_value& const_reference;

    typedef std::shared_ptr<value_type> shared;
    typedef ndArray<T,N> self;



    // Constructors
    ndArray() { reset(); }
    ndArray( pointer ptr, const unsigned *size, bool manage ) { assign(ptr,size,manage); }


    // Copy constructor
    ndArray( const self& other ) { operator=(other); }
    self& operator= ( const self& other );

    // Check pointer validity
    inline bool empty() const { return !((bool) m_data); }
    inline operator bool() const { return m_data; }

    // Print array dimensions
    void info() const;



    // ------------------------------------------------------------------------



    // Clear contents
    void clear();
    void reset();


    // Assign a pointer or another instance
    void assign( pointer ptr, const unsigned *size, bool manage );
    void assign( const self& other ) { operator=(other); }


    // Swap contents with another array
    void swap( self& other );


    // Copy from another array
    template <typename U>
    void copy( const ndArray<U,N>& other );



    // ------------------------------------------------------------------------



    // 1D access
    inline reference operator[] ( unsigned n ) const
        { return data()[ ND_ARRAY_PROTECT(n,m_numel) ]; }

    // ND access
    reference operator() ( const unsigned *subs ) const
        { return data()[ sub2ind<N>(subs, m_size, m_strides) ]; }

    reference operator() ( std::initializer_list<unsigned> subs ) const
        {
#ifdef ND_ARRAY_SAFE_ACCESS
            if ( subs.size() != N )
                throw std::length_error("Invalid coordinates length.");
#endif
            return data()[ sub2ind<N>(subs.begin(), m_size, m_strides) ];
        }

    // Coordinates access
    reference operator() ( unsigned i, ... ) const
        {
            va_list vl; va_start(vl,i);
            return data()[ i+sub2ind<N>(vl, m_size, m_strides) ];
        }


    // Access data directly
    inline const_pointer cdata() const { return m_data.get(); }
    inline pointer data() const { return m_data.get(); }

    // Iterators
    inline const_pointer cbegin() const { return data(); }
    inline const_pointer cend() const { return data() + m_numel; }

    inline pointer begin() const { return data(); }
    inline pointer end() const { return data() + m_numel; }



    // ------------------------------------------------------------------------



    // Dimensions
    inline const unsigned* size() const { return m_size; }
    inline unsigned size( unsigned n ) const { return m_size[ n % N ]; }
    inline const unsigned* strides() const { return m_strides; }
    inline unsigned stride( unsigned n ) const { return m_strides[ n % N ]; }
    inline unsigned numel() const { return m_numel; }
    inline unsigned ndims() const { return N; }



protected:

    void assign_shared( pointer ptr, bool manage );

    unsigned m_numel;
    unsigned m_size[N];
    unsigned m_strides[N];

    shared m_data;
};



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



// Include implementation
#include "ndArray.hpp"

#endif

/**
 * Assignment operator.
 * Performs a shallow copy of the other instance.
 * For deep-copies, see copy() below.
 */
template <typename T, unsigned N>
ndArray<T,N>& ndArray<T,N>::operator=( const self& other )
{
    if ( other.m_data != m_data )
    {
        // Clear current instance first
        clear();

        // Copy data from other
        m_data  = other.m_data;
        m_numel = other.m_numel;
        std::copy_n( other.m_size, N, m_size );
        std::copy_n( other.m_strides, N, m_strides );
    }

    return *this;
}

// ------------------------------------------------------------------------

/**
 * Performs a deep-copy of another instance with possibly
 * different value-type. Deep copies are allowed only if
 * the pointer type is non-const (otherwise no assignment
 * is possible).
 *
 * To perform the copy, a new memory allocation is requested
 * to store as many values as other.m_numel; the current
 * instance takes ownership of this new memory.
 *
 * Note that subsequent shallow copies (see assignment operator)
 * will simply share this ownership (reference counting).
 * Note also that this code might generate warnings because of
 * the value-cast performed on the values of other.
 */
template <typename T, unsigned N>
template <typename U>
void ndArray<T,N>::copy( const ndArray<U,N>& other )
{
    if ( !std::is_const<T>::value )
    {
        // Create new allocation only if necessary
        if ( other.numel() == m_numel )
        {
            // Otherwise simply copy dimensions
            std::copy_n( other.size(), N, m_size );
            std::copy_n( other.strides(), N, m_strides );
        }
        else
            assign( new T[ other.numel() ], other.size(), true );

        // Copy data
        auto dst = begin(); auto src = other.cbegin();
        for ( ;src != other.cend(); ++src, ++dst )  *dst = (T) *src;
    }
    else
        throw std::logic_error("Const values cannot be assigned!");
}

// ------------------------------------------------------------------------

/**
 * Reset shared pointer.
 * This will trigger the deletion of the underlying memory if
 * m_data is unique (m_data.use_count() == 1). Note that if the
 * data was assigned with 'manage' set to false (see below),
 * the deleter(no_delete functor) will NOT release the memory.
 */
template <typename T, unsigned N>
void ndArray<T,N>::clear()
{
    m_data.reset();
}

// ------------------------------------------------------------------------

/**
 * More thorough cleanup. Calls clear() (see above), and sets
 * all the rest to 0.
 */
template <typename T, unsigned N>
void ndArray<T,N>::reset()
{
    clear();
    m_numel = 0;
    std::fill_n( m_size, N, 0 );
    std::fill_n( m_strides, N, 0 );
}

// ------------------------------------------------------------------------

/**
 * Swap contents with another ndArray.
 */
template <typename T, unsigned N>
void ndArray<T,N>::swap( self& other )
{
    std::swap( m_numel, other.m_numel );
    for ( unsigned i = 0; i < N; ++i )
    {
        std::swap( m_size[i], other.m_size[i] );
        std::swap( m_strides[i], other.m_strides[i] );
    }
    m_data.swap( other.m_data );
}

// ------------------------------------------------------------------------

/**
 * Internal method (protected) to assign the shared pointer.
 * Dimensions are assumed to be taken care of by the public
 * assign variants (see below); only the pointer, it's length
 * and the flag 'manage' are required here.
 *
 * 'manage' allows to specify whether or not the shared pointer
 * should release the memory when the last refering instance is
 * destroyed.
 *
 * If true, the default deleter std::default_delete will be
 * assigned to the shared pointer. Note that this deleter releases
 * memory allocated USING NEW ONLY;
 *  DO NOT use malloc/calloc or other C allocation variants.
 *
 * If false, the deleter no_delete is given instead; this will NOT
 * release the memory when the last refering instance is destroyed.
 * Use only with either externally managed (eg Matlab) or static
 * allocations.
 */
template <typename T, unsigned N>
void ndArray<T,N>::assign_shared( pointer ptr, bool manage )
{
    if (manage)
        m_data.reset( ptr );
    else
        m_data.reset( ptr, no_delete<T>() );
}

// ------------------------------------------------------------------------

/**
 * Assign from pointer and size.
 *
 * If manage == true, the internal shared pointer m_data
 * will assume ownership of the memory pointed by ptr, and
 * try to release it using delete[] when the last refering
 * instance gets destroyed.
 * If ptr is dynamically allocated, make sure that the
 * allocation is performed with NEW, and NOT C variants
 * like malloc/calloc/etc.
 *
 * If manage == false, a dummy deleter (no_delete functor) is
 * passed to the shared pointer; nothing happens to the memory
 * pointed by ptr when the last refering instance gets destroyed.
 */
template <typename T, unsigned N>
void ndArray<T,N>::assign( pointer ptr, const unsigned *size, bool manage )
{
    if ( ptr != data() )
    {
        // Compute internal dimensions
        m_numel = 1;
        for ( unsigned i = 0; i < N; ++i )
        {
            m_size[i] = size[i];
            m_numel  *= size[i];
            m_strides[ (i+1) % N ] = m_numel;
        }   m_strides[0] = 1;

        // Set data
        assign_shared( ptr, manage );
    }
}

// ------------------------------------------------------------------------

/**
 * Simply prints information about the dimensions of the
 * n-dimensional array (size & number of elements).
 */
template <typename T, unsigned N>
void ndArray<T,N>::info() const
{
    if ( m_data )
    {
        printf("%u-dimensional array of size (%u", N, m_size[0]);
        for ( unsigned d = 1; d < N; ++d )
            printf(", %u", m_size[d]);
        printf(") = %u elements.n", m_numel);
    }
    else
        printf("Empty %u-dimensional array.n", N);
}