/***************************************************************************
                          DistributedMatrix.h  -  description
                             -------------------
    begin                : Sep 10, 2018
    copyright            : (C) 2018 by Tomas Oberhuber et al.
    email                : tomas.oberhuber@fjfi.cvut.cz
 ***************************************************************************/

/* See Copyright Notice in tnl/Copyright */

// Implemented by: Jakub Klinkovský

#pragma once

#include "DistributedMatrix.h"

namespace TNL {
namespace Matrices {
namespace Legacy {

template< typename Matrix >
DistributedMatrix< Matrix >::
DistributedMatrix( LocalRangeType localRowRange, IndexType rows, IndexType columns, MPI_Comm communicator )
{
   setDistribution( localRowRange, rows, columns, communicator );
}

template< typename Matrix >
void
DistributedMatrix< Matrix >::
setDistribution( LocalRangeType localRowRange, IndexType rows, IndexType columns, MPI_Comm communicator )
{
   this->localRowRange = localRowRange;
   this->rows = rows;
   this->communicator = communicator;
   if( communicator != MPI_COMM_NULL )
      localMatrix.setDimensions( localRowRange.getSize(), columns );

   spmv.reset();
}

template< typename Matrix >
const Containers::Subrange< typename Matrix::IndexType >&
DistributedMatrix< Matrix >::
getLocalRowRange() const
{
   return localRowRange;
}

template< typename Matrix >
MPI_Comm
DistributedMatrix< Matrix >::
getCommunicator() const
{
   return communicator;
}

template< typename Matrix >
const Matrix&
DistributedMatrix< Matrix >::
getLocalMatrix() const
{
   return localMatrix;
}

template< typename Matrix >
Matrix&
DistributedMatrix< Matrix >::
getLocalMatrix()
{
   return localMatrix;
}


/*
 * Some common Matrix methods follow below.
 */

template< typename Matrix >
DistributedMatrix< Matrix >&
DistributedMatrix< Matrix >::
operator=( const DistributedMatrix& matrix )
{
   setLike( matrix );
   localMatrix = matrix.getLocalMatrix();
   return *this;
}

template< typename Matrix >
   template< typename MatrixT >
DistributedMatrix< Matrix >&
DistributedMatrix< Matrix >::
operator=( const MatrixT& matrix )
{
   setLike( matrix );
   localMatrix = matrix.getLocalMatrix();
   return *this;
}

template< typename Matrix >
   template< typename MatrixT >
void
DistributedMatrix< Matrix >::
setLike( const MatrixT& matrix )
{
   localRowRange = matrix.getLocalRowRange();
   rows = matrix.getRows();
   communicator = matrix.getCommunicator();
   localMatrix.setLike( matrix.getLocalMatrix() );

   spmv.reset();
}

template< typename Matrix >
void
DistributedMatrix< Matrix >::
reset()
{
   localRowRange.reset();
   rows = 0;
   communicator = MPI_COMM_NULL;
   localMatrix.reset();

   spmv.reset();
}

template< typename Matrix >
typename Matrix::IndexType
DistributedMatrix< Matrix >::
getRows() const
{
   return rows;
}

template< typename Matrix >
typename Matrix::IndexType
DistributedMatrix< Matrix >::
getColumns() const
{
   return localMatrix.getColumns();
}

template< typename Matrix >
   template< typename RowCapacitiesVector >
void
DistributedMatrix< Matrix >::
setRowCapacities( const RowCapacitiesVector& rowCapacities )
{
   TNL_ASSERT_EQ( rowCapacities.getSize(), getRows(), "row lengths vector has wrong size" );
   TNL_ASSERT_EQ( rowCapacities.getLocalRange(), getLocalRowRange(), "row lengths vector has wrong distribution" );
   TNL_ASSERT_EQ( rowCapacities.getCommunicator(), getCommunicator(), "row lengths vector has wrong communicator" );

   if( getCommunicator() != MPI_COMM_NULL ) {
      localMatrix.setRowCapacities( rowCapacities.getConstLocalView() );

      spmv.reset();
   }
}

template< typename Matrix >
   template< typename Vector >
void
DistributedMatrix< Matrix >::
getCompressedRowLengths( Vector& rowLengths ) const
{
   if( getCommunicator() != MPI_COMM_NULL ) {
      rowLengths.setDistribution( getLocalRowRange(), 0, getRows(), getCommunicator() );
      auto localRowLengths = rowLengths.getLocalView();
      localMatrix.getCompressedRowLengths( localRowLengths );
   }
}

template< typename Matrix >
typename Matrix::IndexType
DistributedMatrix< Matrix >::
getRowCapacity( IndexType row ) const
{
   const IndexType localRow = localRowRange.getLocalIndex( row );
   return localMatrix.getRowCapacity( localRow );
}

template< typename Matrix >
void
DistributedMatrix< Matrix >::
setElement( IndexType row,
            IndexType column,
            RealType value )
{
   const IndexType localRow = localRowRange.getLocalIndex( row );
   localMatrix.setElement( localRow, column, value );
}

template< typename Matrix >
typename Matrix::RealType
DistributedMatrix< Matrix >::
getElement( IndexType row,
            IndexType column ) const
{
   const IndexType localRow = localRowRange.getLocalIndex( row );
   return localMatrix.getElement( localRow, column );
}

template< typename Matrix >
typename Matrix::RealType
DistributedMatrix< Matrix >::
getElementFast( IndexType row,
                IndexType column ) const
{
   const IndexType localRow = localRowRange.getLocalIndex( row );
   return localMatrix.getElementFast( localRow, column );
}

template< typename Matrix >
typename DistributedMatrix< Matrix >::MatrixRow
DistributedMatrix< Matrix >::
getRow( IndexType row )
{
   const IndexType localRow = localRowRange.getLocalIndex( row );
   return localMatrix.getRow( localRow );
}

template< typename Matrix >
typename DistributedMatrix< Matrix >::ConstMatrixRow
DistributedMatrix< Matrix >::
getRow( IndexType row ) const
{
   const IndexType localRow = localRowRange.getLocalIndex( row );
   return localMatrix.getRow( localRow );
}

template< typename Matrix >
   template< typename InVector,
             typename OutVector >
typename std::enable_if< ! HasGetCommunicatorMethod< InVector >::value >::type
DistributedMatrix< Matrix >::
vectorProduct( const InVector& inVector,
               OutVector& outVector ) const
{
   TNL_ASSERT_EQ( inVector.getSize(), getColumns(), "input vector has wrong size" );
   TNL_ASSERT_EQ( outVector.getSize(), getRows(), "output vector has wrong size" );
   TNL_ASSERT_EQ( outVector.getLocalRange(), getLocalRowRange(), "output vector has wrong distribution" );
   TNL_ASSERT_EQ( outVector.getCommunicator(), getCommunicator(), "output vector has wrong communicator" );

   auto outView = outVector.getLocalView();
   localMatrix.vectorProduct( inVector, outView );
}

template< typename Matrix >
void
DistributedMatrix< Matrix >::
updateVectorProductCommunicationPattern()
{
   if( getCommunicator() == MPI_COMM_NULL )
      return;
   spmv.updateCommunicationPattern( getLocalMatrix(), getCommunicator() );
}

template< typename Matrix >
   template< typename InVector,
             typename OutVector >
typename std::enable_if< HasGetCommunicatorMethod< InVector >::value >::type
DistributedMatrix< Matrix >::
vectorProduct( const InVector& inVector,
               OutVector& outVector ) const
{
   TNL_ASSERT_EQ( inVector.getLocalRange(), getLocalRowRange(), "input vector has wrong distribution" );
   TNL_ASSERT_EQ( inVector.getCommunicator(), getCommunicator(), "input vector has wrong communicator" );
   TNL_ASSERT_EQ( outVector.getSize(), getRows(), "output vector has wrong size" );
   TNL_ASSERT_EQ( outVector.getLocalRange(), getLocalRowRange(), "output vector has wrong distribution" );
   TNL_ASSERT_EQ( outVector.getCommunicator(), getCommunicator(), "output vector has wrong communicator" );

   if( getCommunicator() == MPI_COMM_NULL )
      return;

   if( inVector.getGhosts() == 0 ) {
      // NOTE: this branch is deprecated and kept only due to existing benchmarks
      TNL_ASSERT_EQ( inVector.getSize(), getColumns(), "input vector has wrong size" );
      const_cast< DistributedMatrix* >( this )->spmv.vectorProduct( outVector, localMatrix, localRowRange, inVector, getCommunicator() );
   }
   else {
      TNL_ASSERT_EQ( inVector.getConstLocalViewWithGhosts().getSize(), localMatrix.getColumns(), "the matrix uses non-local and non-ghost column indices" );
      TNL_ASSERT_EQ( inVector.getGhosts(), localMatrix.getColumns() - localMatrix.getRows(), "input vector has wrong ghosts size" );
      TNL_ASSERT_EQ( outVector.getGhosts(), localMatrix.getColumns() - localMatrix.getRows(), "output vector has wrong ghosts size" );
      TNL_ASSERT_EQ( outVector.getConstLocalView().getSize(), localMatrix.getRows(), "number of local matrix rows does not match the output vector local size" );

      inVector.waitForSynchronization();
      const auto inView = inVector.getConstLocalViewWithGhosts();
      auto outView = outVector.getLocalView();
      localMatrix.vectorProduct( inView, outView );
      // TODO: synchronization is not always necessary, e.g. when a preconditioning step follows
//      outVector.startSynchronization();
   }
}

template< typename Matrix >
   template< typename Vector1, typename Vector2 >
bool
DistributedMatrix< Matrix >::
performSORIteration( const Vector1& b,
                     const IndexType row,
                     Vector2& x,
                     const RealType& omega ) const
{
   return getLocalMatrix().performSORIteration( b, row, x, omega );
}

} // namespace Legacy
} // namespace Matrices
} // namespace TNL
