// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2012 Désiré Nuentsa-Wakam <desire.nuentsa_wakam@inria.fr>
// Copyright (C) 2012 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_SPARSELU_SUPERNODAL_MATRIX_H
#define EIGEN_SPARSELU_SUPERNODAL_MATRIX_H

namespace Eigen {
namespace internal {

/** \ingroup SparseLU_Module
 * \brief a class to manipulate the L supernodal factor from the SparseLU factorization
 *
 * This class  contain the data to easily store
 * and manipulate the supernodes during the factorization and solution phase of Sparse LU.
 * Only the lower triangular matrix has supernodes.
 *
 * NOTE : This class corresponds to the SCformat structure in SuperLU
 *
 */
/* TODO
 * InnerIterator as for sparsematrix
 * SuperInnerIterator to iterate through all supernodes
 * Function for triangular solve
 */
template <typename _Scalar, typename _StorageIndex>
class MappedSuperNodalMatrix
{
  public:
    typedef _Scalar Scalar;
    typedef _StorageIndex StorageIndex;
    typedef Matrix<StorageIndex,Dynamic,1> IndexVector;
    typedef Matrix<Scalar,Dynamic,1> ScalarVector;
  public:
    MappedSuperNodalMatrix()
    {

    }
    MappedSuperNodalMatrix(Index m, Index n,  ScalarVector& nzval, IndexVector& nzval_colptr, IndexVector& rowind,
                           IndexVector& rowind_colptr, IndexVector& col_to_sup, IndexVector& sup_to_col )
    {
        setInfos(m, n, nzval, nzval_colptr, rowind, rowind_colptr, col_to_sup, sup_to_col);
    }

    ~MappedSuperNodalMatrix()
    {

    }
    /**
     * Set appropriate pointers for the lower triangular supernodal matrix
     * These infos are available at the end of the numerical factorization
     * FIXME This class will be modified such that it can be use in the course
     * of the factorization.
     */
    void
    setInfos(Index m, Index n, ScalarVector& nzval, IndexVector& nzval_colptr, IndexVector& rowind,
             IndexVector& rowind_colptr, IndexVector& col_to_sup, IndexVector& sup_to_col )
    {
        m_row = m;
        m_col = n;
        m_nzval = nzval.data();
        m_nzval_colptr = nzval_colptr.data();
        m_rowind = rowind.data();
        m_rowind_colptr = rowind_colptr.data();
        m_nsuper = col_to_sup(n);
        m_col_to_sup = col_to_sup.data();
        m_sup_to_col = sup_to_col.data();
    }

    /**
     * Number of rows
     */
    Index
    rows()
    {
        return m_row;
    }

    /**
     * Number of columns
     */
    Index
    cols()
    {
        return m_col;
    }

    /**
     * Return the array of nonzero values packed by column
     *
     * The size is nnz
     */
    Scalar*
    valuePtr()
    {
        return m_nzval;
    }

    const Scalar*
    valuePtr() const
    {
        return m_nzval;
    }
    /**
     * Return the pointers to the beginning of each column in \ref valuePtr()
     */
    StorageIndex*
    colIndexPtr()
    {
        return m_nzval_colptr;
    }

    const StorageIndex*
    colIndexPtr() const
    {
        return m_nzval_colptr;
    }

    /**
     * Return the array of compressed row indices of all supernodes
     */
    StorageIndex*
    rowIndex()
    {
        return m_rowind;
    }

    const StorageIndex*
    rowIndex() const
    {
        return m_rowind;
    }

    /**
     * Return the location in \em rowvaluePtr() which starts each column
     */
    StorageIndex*
    rowIndexPtr()
    {
        return m_rowind_colptr;
    }

    const StorageIndex*
    rowIndexPtr() const
    {
        return m_rowind_colptr;
    }

    /**
     * Return the array of column-to-supernode mapping
     */
    StorageIndex*
    colToSup()
    {
        return m_col_to_sup;
    }

    const StorageIndex*
    colToSup() const
    {
        return m_col_to_sup;
    }
    /**
     * Return the array of supernode-to-column mapping
     */
    StorageIndex*
    supToCol()
    {
        return m_sup_to_col;
    }

    const StorageIndex*
    supToCol() const
    {
        return m_sup_to_col;
    }

    /**
     * Return the number of supernodes
     */
    Index
    nsuper() const
    {
        return m_nsuper;
    }

    class InnerIterator;
    template<typename Dest>
    void
    solveInPlace( MatrixBase<Dest>&X) const;




  protected:
    Index m_row; // Number of rows
    Index m_col; // Number of columns
    Index m_nsuper; // Number of supernodes
    Scalar* m_nzval; //array of nonzero values packed by column
    StorageIndex* m_nzval_colptr; //nzval_colptr[j] Stores the location in nzval[] which starts column j
    StorageIndex* m_rowind; // Array of compressed row indices of rectangular supernodes
    StorageIndex* m_rowind_colptr; //rowind_colptr[j] stores the location in rowind[] which starts column j
    StorageIndex* m_col_to_sup; // col_to_sup[j] is the supernode number to which column j belongs
    StorageIndex* m_sup_to_col; //sup_to_col[s] points to the starting column of the s-th supernode

  private :
};

/**
  * \brief InnerIterator class to iterate over nonzero values of the current column in the supernodal matrix L
  *
  */
template<typename Scalar, typename StorageIndex>
class MappedSuperNodalMatrix<Scalar,StorageIndex>::InnerIterator
{
  public:
    InnerIterator(const MappedSuperNodalMatrix& mat, Index outer)
        : m_matrix(mat),
          m_outer(outer),
          m_supno(mat.colToSup()[outer]),
          m_idval(mat.colIndexPtr()[outer]),
          m_startidval(m_idval),
          m_endidval(mat.colIndexPtr()[outer+1]),
          m_idrow(mat.rowIndexPtr()[mat.supToCol()[mat.colToSup()[outer]]]),
          m_endidrow(mat.rowIndexPtr()[mat.supToCol()[mat.colToSup()[outer]]+1])
    {}
    inline InnerIterator&
    operator++()
    {
        m_idval++;
        m_idrow++;
        return *this;
    }
    inline Scalar
    value() const
    {
        return m_matrix.valuePtr()[m_idval];
    }

    inline Scalar&
    valueRef()
    {
        return const_cast<Scalar&>(m_matrix.valuePtr()[m_idval]);
    }

    inline Index
    index() const
    {
        return m_matrix.rowIndex()[m_idrow];
    }
    inline Index
    row() const
    {
        return index();
    }
    inline Index
    col() const
    {
        return m_outer;
    }

    inline Index
    supIndex() const
    {
        return m_supno;
    }

    inline
    operator bool() const
    {
        return ( (m_idval < m_endidval) && (m_idval >= m_startidval)
                 && (m_idrow < m_endidrow) );
    }

  protected:
    const MappedSuperNodalMatrix& m_matrix; // Supernodal lower triangular matrix
    const Index m_outer;                    // Current column
    const Index m_supno;                    // Current SuperNode number
    Index m_idval;                          // Index to browse the values in the current column
    const Index m_startidval;               // Start of the column value
    const Index m_endidval;                 // End of the column value
    Index m_idrow;                          // Index to browse the row indices
    Index m_endidrow;                       // End index of row indices of the current column
};

/**
 * \brief Solve with the supernode triangular matrix
 *
 */
template<typename Scalar, typename Index_>
template<typename Dest>
void
MappedSuperNodalMatrix<Scalar,Index_>::solveInPlace( MatrixBase<Dest>&X) const
{
    /* Explicit type conversion as the Index type of MatrixBase<Dest> may be wider than Index */
    //    eigen_assert(X.rows() <= NumTraits<Index>::highest());
    //    eigen_assert(X.cols() <= NumTraits<Index>::highest());
    Index n    = int(X.rows());
    Index nrhs = Index(X.cols());
    const Scalar * Lval = valuePtr();                 // Nonzero values
    Matrix<Scalar,Dynamic,Dest::ColsAtCompileTime, ColMajor> work(n, nrhs);     // working vector
    work.setZero();

    for (Index k = 0; k <= nsuper(); k ++)
    {
        Index fsupc = supToCol()[k];                    // First column of the current supernode
        Index istart = rowIndexPtr()[fsupc];            // Pointer index to the subscript of the current column
        Index nsupr = rowIndexPtr()[fsupc+1] - istart;  // Number of rows in the current supernode
        Index nsupc = supToCol()[k+1] - fsupc;          // Number of columns in the current supernode
        Index nrow = nsupr - nsupc;                     // Number of rows in the non-diagonal part of the supernode
        Index irow;                                     //Current index row

        if (nsupc == 1 )
        {
            for (Index j = 0; j < nrhs; j++)
            {
                InnerIterator it(*this, fsupc);
                ++it; // Skip the diagonal element

                for (; it; ++it)
                {
                    irow = it.row();
                    X(irow, j) -= X(fsupc, j) * it.value();
                }
            }
        }

        else
        {
            // The supernode has more than one column
            Index luptr = colIndexPtr()[fsupc];
            Index lda = colIndexPtr()[fsupc+1] - luptr;

            // Triangular solve
            Map<const Matrix<Scalar,Dynamic,Dynamic, ColMajor>, 0, OuterStride<> > A( &(Lval[luptr]), nsupc, nsupc, OuterStride<>(lda) );
            Map< Matrix<Scalar,Dynamic,Dest::ColsAtCompileTime, ColMajor>, 0, OuterStride<> > U (&(X(fsupc,0)), nsupc, nrhs, OuterStride<>(n) );
            U = A.template triangularView<UnitLower>().solve(U);

            // Matrix-vector product
            new (&A) Map<const Matrix<Scalar,Dynamic,Dynamic, ColMajor>, 0, OuterStride<> > ( &(Lval[luptr+nsupc]), nrow, nsupc, OuterStride<>(lda) );
            work.topRows(nrow).noalias() = A * U;

            //Begin Scatter
            for (Index j = 0; j < nrhs; j++)
            {
                Index iptr = istart + nsupc;

                for (Index i = 0; i < nrow; i++)
                {
                    irow = rowIndex()[iptr];
                    X(irow, j) -= work(i, j); // Scatter operation
                    work(i, j) = Scalar(0);
                    iptr++;
                }
            }
        }
    }
}

} // end namespace internal

} // end namespace Eigen

#endif // EIGEN_SPARSELU_MATRIX_H
