// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2007-2009 Benoit Jacob <jacob.benoit.1@gmail.com>
// Copyright (C) 2009-2010 Gael Guennebaud <gael.guennebaud@inria.fr>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_DIAGONAL_H
#define EIGEN_DIAGONAL_H

#include "./InternalHeaderCheck.h"

namespace Eigen {

/** \class Diagonal
  * \ingroup Core_Module
  *
  * \brief Expression of a diagonal/subdiagonal/superdiagonal in a matrix
  *
  * \tparam MatrixType the type of the object in which we are taking a sub/main/super diagonal
  * \tparam DiagIndex the index of the sub/super diagonal. The default is 0 and it means the main diagonal.
  *              A positive value means a superdiagonal, a negative value means a subdiagonal.
  *              You can also use DynamicIndex so the index can be set at runtime.
  *
  * The matrix is not required to be square.
  *
  * This class represents an expression of the main diagonal, or any sub/super diagonal
  * of a square matrix. It is the return type of MatrixBase::diagonal() and MatrixBase::diagonal(Index) and most of the
  * time this is the only way it is used.
  *
  * \sa MatrixBase::diagonal(), MatrixBase::diagonal(Index)
  */

namespace internal {
template<typename MatrixType, int DiagIndex>
struct traits<Diagonal<MatrixType,DiagIndex> >
 : traits<MatrixType>
{
  typedef typename ref_selector<MatrixType>::type MatrixTypeNested;
  typedef std::remove_reference_t<MatrixTypeNested> MatrixTypeNested_;
  typedef typename MatrixType::StorageKind StorageKind;
  enum {
    RowsAtCompileTime = (int(DiagIndex) == DynamicIndex || int(MatrixType::SizeAtCompileTime) == Dynamic) ? Dynamic
                      : (plain_enum_min(MatrixType::RowsAtCompileTime - plain_enum_max(-DiagIndex, 0),
                                        MatrixType::ColsAtCompileTime - plain_enum_max( DiagIndex, 0))),
    ColsAtCompileTime = 1,
    MaxRowsAtCompileTime = int(MatrixType::MaxSizeAtCompileTime) == Dynamic ? Dynamic
                         : DiagIndex == DynamicIndex ? min_size_prefer_fixed(MatrixType::MaxRowsAtCompileTime,
                                                                             MatrixType::MaxColsAtCompileTime)
                         : (plain_enum_min(MatrixType::MaxRowsAtCompileTime - plain_enum_max(-DiagIndex, 0),
                                           MatrixType::MaxColsAtCompileTime - plain_enum_max( DiagIndex, 0))),
    MaxColsAtCompileTime = 1,
    MaskLvalueBit = is_lvalue<MatrixType>::value ? LvalueBit : 0,
    Flags = (unsigned int)MatrixTypeNested_::Flags & (RowMajorBit | MaskLvalueBit | DirectAccessBit) & ~RowMajorBit, // FIXME DirectAccessBit should not be handled by expressions
    MatrixTypeOuterStride = outer_stride_at_compile_time<MatrixType>::ret,
    InnerStrideAtCompileTime = MatrixTypeOuterStride == Dynamic ? Dynamic : MatrixTypeOuterStride+1,
    OuterStrideAtCompileTime = 0
  };
};
}

template<typename MatrixType, int DiagIndex_> class Diagonal
   : public internal::dense_xpr_base< Diagonal<MatrixType,DiagIndex_> >::type
{
  public:

    enum { DiagIndex = DiagIndex_ };
    typedef typename internal::dense_xpr_base<Diagonal>::type Base;
    EIGEN_DENSE_PUBLIC_INTERFACE(Diagonal)

    EIGEN_DEVICE_FUNC
    explicit inline Diagonal(MatrixType& matrix, Index a_index = DiagIndex) : m_matrix(matrix), m_index(a_index)
    {
      eigen_assert( a_index <= m_matrix.cols() && -a_index <= m_matrix.rows() );
    }

    EIGEN_INHERIT_ASSIGNMENT_OPERATORS(Diagonal)

    EIGEN_DEVICE_FUNC
    inline Index rows() const
    {
      return m_index.value()<0 ? numext::mini<Index>(m_matrix.cols(),m_matrix.rows()+m_index.value())
                               : numext::mini<Index>(m_matrix.rows(),m_matrix.cols()-m_index.value());
    }

    EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR
    inline Index cols() const EIGEN_NOEXCEPT { return 1; }

    EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR
    inline Index innerStride() const EIGEN_NOEXCEPT {
      return m_matrix.outerStride() + 1;
    }

    EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR
    inline Index outerStride() const EIGEN_NOEXCEPT { return 0; }

    typedef std::conditional_t<
              internal::is_lvalue<MatrixType>::value,
              Scalar,
              const Scalar
            > ScalarWithConstIfNotLvalue;

    EIGEN_DEVICE_FUNC
    inline ScalarWithConstIfNotLvalue* data() { return &(m_matrix.coeffRef(rowOffset(), colOffset())); }
    EIGEN_DEVICE_FUNC
    inline const Scalar* data() const { return &(m_matrix.coeffRef(rowOffset(), colOffset())); }

    EIGEN_DEVICE_FUNC
    inline Scalar& coeffRef(Index row, Index)
    {
      EIGEN_STATIC_ASSERT_LVALUE(MatrixType)
      return m_matrix.coeffRef(row+rowOffset(), row+colOffset());
    }

    EIGEN_DEVICE_FUNC
    inline const Scalar& coeffRef(Index row, Index) const
    {
      return m_matrix.coeffRef(row+rowOffset(), row+colOffset());
    }

    EIGEN_DEVICE_FUNC
    inline CoeffReturnType coeff(Index row, Index) const
    {
      return m_matrix.coeff(row+rowOffset(), row+colOffset());
    }

    EIGEN_DEVICE_FUNC
    inline Scalar& coeffRef(Index idx)
    {
      EIGEN_STATIC_ASSERT_LVALUE(MatrixType)
      return m_matrix.coeffRef(idx+rowOffset(), idx+colOffset());
    }

    EIGEN_DEVICE_FUNC
    inline const Scalar& coeffRef(Index idx) const
    {
      return m_matrix.coeffRef(idx+rowOffset(), idx+colOffset());
    }

    EIGEN_DEVICE_FUNC
    inline CoeffReturnType coeff(Index idx) const
    {
      return m_matrix.coeff(idx+rowOffset(), idx+colOffset());
    }

    EIGEN_DEVICE_FUNC
    inline const internal::remove_all_t<typename MatrixType::Nested>&
    nestedExpression() const
    {
      return m_matrix;
    }

    EIGEN_DEVICE_FUNC
    inline Index index() const
    {
      return m_index.value();
    }

  protected:
    typename internal::ref_selector<MatrixType>::non_const_type m_matrix;
    const internal::variable_if_dynamicindex<Index, DiagIndex> m_index;

  private:
    // some compilers may fail to optimize std::max etc in case of compile-time constants...
    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE EIGEN_CONSTEXPR
    Index absDiagIndex() const EIGEN_NOEXCEPT { return m_index.value()>0 ? m_index.value() : -m_index.value(); }
    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE EIGEN_CONSTEXPR
    Index rowOffset() const EIGEN_NOEXCEPT { return m_index.value()>0 ? 0 : -m_index.value(); }
    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE EIGEN_CONSTEXPR
    Index colOffset() const EIGEN_NOEXCEPT { return m_index.value()>0 ? m_index.value() : 0; }
    // trigger a compile-time error if someone try to call packet
    template<int LoadMode> typename MatrixType::PacketReturnType packet(Index) const;
    template<int LoadMode> typename MatrixType::PacketReturnType packet(Index,Index) const;
};

/** \returns an expression of the main diagonal of the matrix \c *this
  *
  * \c *this is not required to be square.
  *
  * Example: \include MatrixBase_diagonal.cpp
  * Output: \verbinclude MatrixBase_diagonal.out
  *
  * \sa class Diagonal */
template<typename Derived>
EIGEN_DEVICE_FUNC inline typename MatrixBase<Derived>::DiagonalReturnType
MatrixBase<Derived>::diagonal()
{
  return DiagonalReturnType(derived());
}

/** This is the const version of diagonal(). */
template<typename Derived>
EIGEN_DEVICE_FUNC inline
const typename MatrixBase<Derived>::ConstDiagonalReturnType
MatrixBase<Derived>::diagonal() const
{
  return ConstDiagonalReturnType(derived());
}

/** \returns an expression of the \a DiagIndex-th sub or super diagonal of the matrix \c *this
  *
  * \c *this is not required to be square.
  *
  * The template parameter \a DiagIndex represent a super diagonal if \a DiagIndex > 0
  * and a sub diagonal otherwise. \a DiagIndex == 0 is equivalent to the main diagonal.
  *
  * Example: \include MatrixBase_diagonal_int.cpp
  * Output: \verbinclude MatrixBase_diagonal_int.out
  *
  * \sa MatrixBase::diagonal(), class Diagonal */
template<typename Derived>
EIGEN_DEVICE_FUNC inline Diagonal<Derived, DynamicIndex>
MatrixBase<Derived>::diagonal(Index index)
{
  return Diagonal<Derived, DynamicIndex>(derived(), index);
}

/** This is the const version of diagonal(Index). */
template<typename Derived>
EIGEN_DEVICE_FUNC inline const Diagonal<const Derived, DynamicIndex>
MatrixBase<Derived>::diagonal(Index index) const
{
  return Diagonal<const Derived, DynamicIndex>(derived(), index);
}

/** \returns an expression of the \a DiagIndex-th sub or super diagonal of the matrix \c *this
  *
  * \c *this is not required to be square.
  *
  * The template parameter \a DiagIndex represent a super diagonal if \a DiagIndex > 0
  * and a sub diagonal otherwise. \a DiagIndex == 0 is equivalent to the main diagonal.
  *
  * Example: \include MatrixBase_diagonal_template_int.cpp
  * Output: \verbinclude MatrixBase_diagonal_template_int.out
  *
  * \sa MatrixBase::diagonal(), class Diagonal */
template<typename Derived>
template<int Index_>
EIGEN_DEVICE_FUNC
inline Diagonal<Derived, Index_>
MatrixBase<Derived>::diagonal()
{
  return Diagonal<Derived, Index_>(derived());
}

/** This is the const version of diagonal<int>(). */
template<typename Derived>
template<int Index_>
EIGEN_DEVICE_FUNC
inline const Diagonal<const Derived, Index_>
MatrixBase<Derived>::diagonal() const
{
  return  Diagonal<const Derived, Index_>(derived());
}

} // end namespace Eigen

#endif // EIGEN_DIAGONAL_H