// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2009-2015 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_SPARSE_DIAGONAL_PRODUCT_H
#define EIGEN_SPARSE_DIAGONAL_PRODUCT_H

namespace Eigen {

// The product of a diagonal matrix with a sparse matrix can be easily
// implemented using expression template.
// We have two consider very different cases:
// 1 - diag * row-major sparse
//     => each inner vector <=> scalar * sparse vector product
//     => so we can reuse CwiseUnaryOp::InnerIterator
// 2 - diag * col-major sparse
//     => each inner vector <=> densevector * sparse vector cwise product
//     => again, we can reuse specialization of CwiseBinaryOp::InnerIterator
//        for that particular case
// The two other cases are symmetric.

namespace internal {

enum
{
    SDP_AsScalarProduct,
    SDP_AsCwiseProduct
};

template<typename SparseXprType, typename DiagonalCoeffType, int SDP_Tag>
struct sparse_diagonal_product_evaluator;

template<typename Lhs, typename Rhs, int ProductTag>
struct product_evaluator<Product<Lhs, Rhs, DefaultProduct>, ProductTag, DiagonalShape, SparseShape>
    : public sparse_diagonal_product_evaluator<Rhs, typename Lhs::DiagonalVectorType, Rhs::Flags&RowMajorBit?SDP_AsScalarProduct:SDP_AsCwiseProduct>
{
    typedef Product<Lhs, Rhs, DefaultProduct> XprType;
    enum { CoeffReadCost = HugeCost, Flags = Rhs::Flags&RowMajorBit, Alignment = 0 }; // FIXME CoeffReadCost & Flags

    typedef sparse_diagonal_product_evaluator<Rhs, typename Lhs::DiagonalVectorType, Rhs::Flags&RowMajorBit?SDP_AsScalarProduct:SDP_AsCwiseProduct> Base;
    explicit
    product_evaluator(const XprType& xpr) : Base(xpr.rhs(), xpr.lhs().diagonal()) {}
};

template<typename Lhs, typename Rhs, int ProductTag>
struct product_evaluator<Product<Lhs, Rhs, DefaultProduct>, ProductTag, SparseShape, DiagonalShape>
    : public sparse_diagonal_product_evaluator<Lhs, Transpose<const typename Rhs::DiagonalVectorType>, Lhs::Flags&RowMajorBit?SDP_AsCwiseProduct:SDP_AsScalarProduct>
{
    typedef Product<Lhs, Rhs, DefaultProduct> XprType;
    enum { CoeffReadCost = HugeCost, Flags = Lhs::Flags&RowMajorBit, Alignment = 0 }; // FIXME CoeffReadCost & Flags

    typedef sparse_diagonal_product_evaluator<Lhs, Transpose<const typename Rhs::DiagonalVectorType>, Lhs::Flags&RowMajorBit?SDP_AsCwiseProduct:SDP_AsScalarProduct> Base;
    explicit
    product_evaluator(const XprType& xpr) : Base(xpr.lhs(), xpr.rhs().diagonal().transpose()) {}
};

template<typename SparseXprType, typename DiagonalCoeffType>
struct sparse_diagonal_product_evaluator<SparseXprType, DiagonalCoeffType, SDP_AsScalarProduct>
{
  protected:
    typedef typename evaluator<SparseXprType>::InnerIterator SparseXprInnerIterator;
    typedef typename SparseXprType::Scalar Scalar;

  public:
    class InnerIterator : public SparseXprInnerIterator
    {
      public:
        InnerIterator(const sparse_diagonal_product_evaluator &xprEval, Index outer)
            : SparseXprInnerIterator(xprEval.m_sparseXprImpl, outer),
              m_coeff(xprEval.m_diagCoeffImpl.coeff(outer))
        {}

        EIGEN_STRONG_INLINE Scalar
        value() const
        {
            return m_coeff * SparseXprInnerIterator::value();
        }
      protected:
        typename DiagonalCoeffType::Scalar m_coeff;
    };

    sparse_diagonal_product_evaluator(const SparseXprType &sparseXpr, const DiagonalCoeffType &diagCoeff)
        : m_sparseXprImpl(sparseXpr), m_diagCoeffImpl(diagCoeff)
    {}

    Index
    nonZerosEstimate() const
    {
        return m_sparseXprImpl.nonZerosEstimate();
    }

  protected:
    evaluator<SparseXprType> m_sparseXprImpl;
    evaluator<DiagonalCoeffType> m_diagCoeffImpl;
};


template<typename SparseXprType, typename DiagCoeffType>
struct sparse_diagonal_product_evaluator<SparseXprType, DiagCoeffType, SDP_AsCwiseProduct>
{
    typedef typename SparseXprType::Scalar Scalar;
    typedef typename SparseXprType::StorageIndex StorageIndex;

    typedef typename nested_eval<DiagCoeffType,SparseXprType::IsRowMajor ? SparseXprType::RowsAtCompileTime
    : SparseXprType::ColsAtCompileTime>::type DiagCoeffNested;

    class InnerIterator
    {
        typedef typename evaluator<SparseXprType>::InnerIterator SparseXprIter;
      public:
        InnerIterator(const sparse_diagonal_product_evaluator &xprEval, Index outer)
            : m_sparseIter(xprEval.m_sparseXprEval, outer), m_diagCoeffNested(xprEval.m_diagCoeffNested)
        {}

        inline Scalar
        value() const
        {
            return m_sparseIter.value() * m_diagCoeffNested.coeff(index());
        }
        inline StorageIndex
        index() const
        {
            return m_sparseIter.index();
        }
        inline Index
        outer() const
        {
            return m_sparseIter.outer();
        }
        inline Index
        col() const
        {
            return SparseXprType::IsRowMajor ? m_sparseIter.index() : m_sparseIter.outer();
        }
        inline Index
        row() const
        {
            return SparseXprType::IsRowMajor ? m_sparseIter.outer() : m_sparseIter.index();
        }

        EIGEN_STRONG_INLINE InnerIterator&
        operator++()
        {
            ++m_sparseIter;
            return *this;
        }
        inline
        operator bool() const
        {
            return m_sparseIter;
        }

      protected:
        SparseXprIter m_sparseIter;
        DiagCoeffNested m_diagCoeffNested;
    };

    sparse_diagonal_product_evaluator(const SparseXprType &sparseXpr, const DiagCoeffType &diagCoeff)
        : m_sparseXprEval(sparseXpr), m_diagCoeffNested(diagCoeff)
    {}

    Index
    nonZerosEstimate() const
    {
        return m_sparseXprEval.nonZerosEstimate();
    }

  protected:
    evaluator<SparseXprType> m_sparseXprEval;
    DiagCoeffNested m_diagCoeffNested;
};

} // end namespace internal

} // end namespace Eigen

#endif // EIGEN_SPARSE_DIAGONAL_PRODUCT_H
