#include <algorithm>
#include <assert.h>
#include <cmath>
#include <cstdlib>
#include <iostream>
#include <iomanip>
#include <string>
#include <vector>

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

class Matrix
{
private:
  static const size_t MAX_DIM = 6;
  std::vector<double> m_data;             // data container
  size_t              m_ndim=0;           // actual number of dimensions
  size_t              m_size=0;           // total number of entries == data.size() == prod(shape)
  size_t              m_shape[MAX_DIM];   // number of entries in each dimensions
  size_t              m_strides[MAX_DIM]; // stride length for each index

public:

  // constructor
  Matrix(){};                                     // empty
  Matrix(const std::vector<size_t> &shape);       // allocate, don't initialize

  // index operators: access plain storage
  double&       operator[](size_t i);
  const double& operator[](size_t i) const;

  // index operators: access using Matrix indices
  double&       operator()(size_t a);
  const double& operator()(size_t a) const;
  double&       operator()(size_t a, size_t b);
  const double& operator()(size_t a, size_t b) const;
  double&       operator()(size_t a, size_t b, size_t c);
  const double& operator()(size_t a, size_t b, size_t c) const;
  double&       operator()(size_t a, size_t b, size_t c, size_t d);
  const double& operator()(size_t a, size_t b, size_t c, size_t d) const;
  double&       operator()(size_t a, size_t b, size_t c, size_t d, size_t e);
  const double& operator()(size_t a, size_t b, size_t c, size_t d, size_t e) const;
  double&       operator()(size_t a, size_t b, size_t c, size_t d, size_t e, size_t f);
  const double& operator()(size_t a, size_t b, size_t c, size_t d, size_t e, size_t f) const;

  // pointer to data
  double*       data();
  const double* data() const;

  // iterator to first and last entry
  auto begin();
  auto begin() const;
  auto end();
  auto end() const;

  // get dimensions
  size_t size() const;
  size_t shape(size_t i) const;
  std::vector<size_t> shape() const;

  // initialization
  void setZero();

  // basic algebra
  Matrix sum(size_t axis) const;
};

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

inline std::vector<size_t> del(const std::vector<size_t> &A, size_t idx);

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

inline Matrix::Matrix(const std::vector<size_t> &shape)
{
  assert( shape.size()  > 0       );
  assert( shape.size() <= MAX_DIM );

  for ( size_t i = 0 ; i < MAX_DIM ; ++i )
  {
    m_shape  [i] = 1;
    m_strides[i] = 1;
  }

  m_ndim = shape.size();
  m_size = 1;

  for ( size_t i = 0 ; i < m_ndim ; ++i )
  {
    m_shape[i] = shape[i];
    m_size    *= shape[i];
  }

  for ( size_t i = 0 ; i < m_ndim ; ++i )
    for ( size_t j = i+1 ; j < m_ndim ; ++j )
      m_strides[i] *= m_shape[j];

  m_data.resize(m_size);
}

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

inline double& Matrix::operator[](size_t i)
{
  return m_data[i];
}

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

inline const double& Matrix::operator[](size_t i) const
{
  return m_data[i];
}

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

inline double& Matrix::operator()(size_t a)
{
  return m_data[a*m_strides[0]];
}

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

inline const double& Matrix::operator()(size_t a) const
{
  return m_data[a*m_strides[0]];
}

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

inline double& Matrix::operator()(size_t a, size_t b)
{
  return m_data[a*m_strides[0]+b*m_strides[1]];
}

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

inline const double& Matrix::operator()(size_t a, size_t b) const
{
  return m_data[a*m_strides[0]+b*m_strides[1]];
}

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

inline double& Matrix::operator()(size_t a, size_t b, size_t c)
{
  return m_data[a*m_strides[0]+b*m_strides[1]+c*m_strides[2]];
}

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

inline const double& Matrix::operator()(size_t a, size_t b, size_t c) const
{
  return m_data[a*m_strides[0]+b*m_strides[1]+c*m_strides[2]];
}

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

inline double& Matrix::operator()(size_t a, size_t b, size_t c, size_t d)
{
  return m_data[a*m_strides[0]+b*m_strides[1]+c*m_strides[2]+d*m_strides[3]];
}

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

inline const double& Matrix::operator()(size_t a, size_t b, size_t c, size_t d) const
{
  return m_data[a*m_strides[0]+b*m_strides[1]+c*m_strides[2]+d*m_strides[3]];
}

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

inline double& Matrix::operator()(size_t a, size_t b, size_t c, size_t d, size_t e)
{
  return m_data[a*m_strides[0]+b*m_strides[1]+c*m_strides[2]+d*m_strides[3]+e*m_strides[4]];
}

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

inline const double& Matrix::operator()(size_t a, size_t b, size_t c, size_t d, size_t e) const
{
  return m_data[a*m_strides[0]+b*m_strides[1]+c*m_strides[2]+d*m_strides[3]+e*m_strides[4]];
}

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

inline double& Matrix::operator()(size_t a, size_t b, size_t c, size_t d, size_t e, size_t f)
{
  return m_data[a*m_strides[0]+b*m_strides[1]+c*m_strides[2]+d*m_strides[3]+e*m_strides[4]+f*m_strides[5]];
}

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

inline const double& Matrix::operator()(size_t a, size_t b, size_t c, size_t d, size_t e, size_t f) const
{
  return m_data[a*m_strides[0]+b*m_strides[1]+c*m_strides[2]+d*m_strides[3]+e*m_strides[4]+f*m_strides[5]];
}

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

inline double* Matrix::data()
{
  return m_data.data();
}

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

inline const double* Matrix::data() const
{
  return m_data.data();
}

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

inline auto Matrix::begin()
{
  return m_data.begin();
}

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

inline auto Matrix::begin() const
{
  return m_data.begin();
}

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

inline auto Matrix::end()
{
  return m_data.end();
}

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

inline auto Matrix::end() const
{
  return m_data.end();
}

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

inline size_t Matrix::shape(size_t i) const
{
  return m_shape[i];
}

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

inline std::vector<size_t> Matrix::shape() const
{
  std::vector<size_t> out(m_ndim);

  for ( size_t i = 0 ; i < m_ndim ; ++i ) out[i] = m_shape[i];

  return out;
}

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

inline size_t Matrix::size() const
{
  return m_size;
}

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

inline void Matrix::setZero()
{
  for ( size_t i = 0 ; i < m_size ; ++i ) m_data[i] = static_cast<double>(0);
}

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

inline Matrix Matrix::sum(size_t axis) const
{
  Matrix out(del(this->shape(),axis));

  out.setZero();

  // How many elements to advance after each reduction
  size_t step_axis = m_strides[MAX_DIM-1];

  if ( axis == MAX_DIM-1)
    step_axis = m_strides[MAX_DIM-2];

  // Position of the first element of the current reduction
  size_t offset = 0;

  for ( auto &v : out )
  {
    // Current reduced element
    size_t offset_i = offset;

    for ( size_t i = 0 ; i < m_shape[axis] ; ++i )
    {
      // - reduce
      v += *(m_data.data() + offset_i);
      // - advance to next element
      offset_i += m_strides[axis];
    }

    offset += step_axis;
  }

  return out;
}

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

inline std::vector<size_t> del(const std::vector<size_t> &A, size_t idx)
{
  assert( idx < A.size() );

  std::vector<size_t> B = A;

  B.erase(B.begin()+idx, B.begin()+idx+1);

  return B;
}

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

int main()
{
  Matrix A({6,11,16,3});

  for ( size_t i = 0 ; i < A.size() ; ++i )
    A[i] = static_cast<double>(i) + 0.1;

  Matrix B = A.sum(1);

  Matrix b({6,16,3});

  b.setZero();

  for ( size_t i = 0 ; i < A.shape(0) ; i++ )
    for ( size_t j = 0 ; j < A.shape(1) ; j++ )
      for ( size_t k = 0 ; k < A.shape(2) ; k++ )
        for ( size_t l = 0 ; l < A.shape(3) ; l++ )
          b(i,k,l) += A(i,j,k,l);

  for ( size_t i = 0 ; i < b.size() ; ++i )
    if ( std::abs( b[i] - B[i] ) > 1.e-12 )
      std::cout << b[i] << ", " << B[i] << std::endl;

  return 0;
}

