// SPDX-FileCopyrightText: The Eigen Authors
// SPDX-License-Identifier: MPL-2.0

#include <Eigen/Core>
#include <iostream>

// [functor]
template <class ArgType, class RowIndexType, class ColIndexType>
class indexing_functor {
  const ArgType& m_arg;
  const RowIndexType& m_rowIndices;
  const ColIndexType& m_colIndices;

 public:
  typedef Eigen::Matrix<typename ArgType::Scalar, RowIndexType::SizeAtCompileTime, ColIndexType::SizeAtCompileTime,
                        ArgType::Flags & Eigen::RowMajorBit ? Eigen::RowMajor : Eigen::ColMajor,
                        RowIndexType::MaxSizeAtCompileTime, ColIndexType::MaxSizeAtCompileTime>
      MatrixType;

  indexing_functor(const ArgType& arg, const RowIndexType& row_indices, const ColIndexType& col_indices)
      : m_arg(arg), m_rowIndices(row_indices), m_colIndices(col_indices) {}

  const typename ArgType::Scalar& operator()(Eigen::Index row, Eigen::Index col) const {
    return m_arg(m_rowIndices[row], m_colIndices[col]);
  }
};
// [functor]

// [function]
template <class ArgType, class RowIndexType, class ColIndexType>
Eigen::CwiseNullaryOp<indexing_functor<ArgType, RowIndexType, ColIndexType>,
                      typename indexing_functor<ArgType, RowIndexType, ColIndexType>::MatrixType>
mat_indexing(const Eigen::MatrixBase<ArgType>& arg, const RowIndexType& row_indices, const ColIndexType& col_indices) {
  typedef indexing_functor<ArgType, RowIndexType, ColIndexType> Func;
  typedef typename Func::MatrixType MatrixType;
  return MatrixType::NullaryExpr(row_indices.size(), col_indices.size(), Func(arg.derived(), row_indices, col_indices));
}
// [function]

int main() {
  std::cout << "[main1]\n";
  Eigen::MatrixXi A = Eigen::MatrixXi::Random(4, 4);
  Eigen::Array3i ri(1, 2, 1);
  Eigen::ArrayXi ci(6);
  ci << 3, 2, 1, 0, 0, 2;
  Eigen::MatrixXi B = mat_indexing(A, ri, ci);
  std::cout << "A =" << std::endl;
  std::cout << A << std::endl << std::endl;
  std::cout << "A([" << ri.transpose() << "], [" << ci.transpose() << "]) =" << std::endl;
  std::cout << B << std::endl;
  std::cout << "[main1]\n";

  std::cout << "[main2]\n";
  B = mat_indexing(A, ri + 1, ci);
  std::cout << "A(ri+1,ci) =" << std::endl;
  std::cout << B << std::endl << std::endl;
  B = mat_indexing(A, Eigen::ArrayXi::LinSpaced(13, 0, 12).unaryExpr([](int x) { return x % 4; }),
                   Eigen::ArrayXi::LinSpaced(4, 0, 3));
  std::cout << "A(ArrayXi::LinSpaced(13,0,12).unaryExpr([](int x){return x%4;}), ArrayXi::LinSpaced(4,0,3)) ="
            << std::endl;
  std::cout << B << std::endl << std::endl;
  std::cout << "[main2]\n";
}