/***************************************************************************
                          ChunkedEllpackView.hpp -  description
                             -------------------
    begin                : Mar 21, 2020
    copyright            : (C) 2020 by Tomas Oberhuber
    email                : tomas.oberhuber@fjfi.cvut.cz
 ***************************************************************************/

/* See Copyright Notice in tnl/Copyright */

#pragma once

#include <TNL/Containers/Vector.h>
#include <TNL/Algorithms/ParallelFor.h>
#include <TNL/Containers/Segments/ChunkedEllpackView.h>
//#include <TNL/Containers/Segments/details/ChunkedEllpack.h>

namespace TNL {
   namespace Containers {
      namespace Segments {

#ifdef HAVE_CUDA
template< typename View,
          typename Index,
          typename Fetch,
          typename Reduction,
          typename ResultKeeper,
          typename Real,
          typename... Args >
__global__
void ChunkedEllpackSegmentsReductionKernel( View chunkedEllpack,
                                            Index gridIdx,
                                            Index first,
                                            Index last,
                                            Fetch fetch,
                                            Reduction reduction,
                                            ResultKeeper keeper,
                                            Real zero,
                                            Args... args )
{
   chunkedEllpack.segmentsReductionKernel( gridIdx, first, last, fetch, reduction, keeper, zero, args... );
}
#endif


template< typename Device,
          typename Index,
          bool RowMajorOrder >
__cuda_callable__
ChunkedEllpackView< Device, Index, RowMajorOrder >::
ChunkedEllpackView( const IndexType size,
                    const IndexType storageSize,
                    const IndexType chunksInSlice,
                    const IndexType desiredChunkSize,
                    const OffsetsView& rowToChunkMapping,
                    const OffsetsView& rowToSliceMapping,
                    const OffsetsView& chunksToSegmentsMapping,
                    const OffsetsView& rowPointers,
                    const ChunkedEllpackSliceInfoContainerView& slices,
                    const IndexType numberOfSlices )
: size( size ),
  storageSize( storageSize ),
  numberOfSlices( numberOfSlices ),
  chunksInSlice( chunksInSlice ),
  desiredChunkSize( desiredChunkSize ),
  rowToSliceMapping( rowToSliceMapping ),
  rowToChunkMapping( rowToChunkMapping ),
  chunksToSegmentsMapping( chunksToSegmentsMapping ),
  rowPointers( rowPointers ),
  slices( slices )
{
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
__cuda_callable__
ChunkedEllpackView< Device, Index, RowMajorOrder >::
ChunkedEllpackView( const IndexType size,
                    const IndexType storageSize,
                    const IndexType chunksInSlice,
                    const IndexType desiredChunkSize,
                    const OffsetsView&& rowToChunkMapping,
                    const OffsetsView&& rowToSliceMapping,
                    const OffsetsView&& chunksToSegmentsMapping,
                    const OffsetsView&& rowPointers,
                    const ChunkedEllpackSliceInfoContainerView&& slices,
                    const IndexType numberOfSlices )
: size( size ),
  storageSize( storageSize ),
  numberOfSlices( numberOfSlices ),
  chunksInSlice( chunksInSlice ),
  desiredChunkSize( desiredChunkSize ),
  rowToSliceMapping( std::move( rowToSliceMapping ) ),
  rowToChunkMapping( std::move( rowToChunkMapping ) ),
  chunksToSegmentsMapping( std::move( chunksToSegmentsMapping ) ),
  rowPointers( std::move( rowPointers ) ),
  slices( std::move( slices ) )
{
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
__cuda_callable__
ChunkedEllpackView< Device, Index, RowMajorOrder >::
ChunkedEllpackView( const ChunkedEllpackView& chunked_ellpack_view )
: size( chunked_ellpack_view.size ),
  storageSize( chunked_ellpack_view.storageSize ),
  numberOfSlices( chunked_ellpack_view.numberOfSlices ),
  chunksInSlice( chunked_ellpack_view.chunksInSlice ),
  desiredChunkSize( chunked_ellpack_view.desiredChunkSize ),
  rowToSliceMapping( chunked_ellpack_view.rowToSliceMapping ),
  rowToChunkMapping( chunked_ellpack_view.rowToChunkMapping ),
  chunksToSegmentsMapping( chunked_ellpack_view.chunksToSegmentsMapping ),
  rowPointers( chunked_ellpack_view.rowPointers ),
  slices( chunked_ellpack_view.slices )
{
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
__cuda_callable__
ChunkedEllpackView< Device, Index, RowMajorOrder >::
ChunkedEllpackView( const ChunkedEllpackView&& chunked_ellpack_view )
: size( chunked_ellpack_view.size ),
  storageSize( chunked_ellpack_view.storageSize ),
  numberOfSlices( chunked_ellpack_view.numberOfSlices ),
  chunksInSlice( chunked_ellpack_view.chunksInSlice ),
  desiredChunkSize( chunked_ellpack_view.desiredChunkSize ),
  rowToSliceMapping( std::move( chunked_ellpack_view.rowToSliceMapping ) ),
  rowToChunkMapping( std::move( chunked_ellpack_view.rowToChunkMapping ) ),
  chunksToSegmentsMapping( std::move( chunked_ellpack_view.chunksToSegmentsMapping ) ),
  rowPointers( std::move( chunked_ellpack_view.rowPointers ) ),
  slices( std::move( chunked_ellpack_view.slices ) )
{
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
String
ChunkedEllpackView< Device, Index, RowMajorOrder >::
getSerializationType()
{
   return "ChunkedEllpack< [any_device], " + TNL::getSerializationType< IndexType >() + " >";
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
String
ChunkedEllpackView< Device, Index, RowMajorOrder >::
getSegmentsType()
{
   return "ChunkedEllpack";
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
__cuda_callable__
typename ChunkedEllpackView< Device, Index, RowMajorOrder >::ViewType
ChunkedEllpackView< Device, Index, RowMajorOrder >::
getView()
{
   return ViewType( size, chunksInSlice, desiredChunkSize,
                    rowToChunkMapping.getView(),
                    rowToSliceMapping.getView(),
                    chunksToSegmentsMapping.getView(),
                    rowPointers.getView(),
                    slices.getView(),
                    numberOfSlices );
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
__cuda_callable__
typename ChunkedEllpackView< Device, Index, RowMajorOrder >::ConstViewType
ChunkedEllpackView< Device, Index, RowMajorOrder >::
getConstView() const
{
   return ConstViewType( size, chunksInSlice, desiredChunkSize,
                         rowToChunkMapping.getConstView(),
                         rowToSliceMapping.getConstView(),
                         chunksToSegmentsMapping.getConstView(),
                         rowPointers.getConstView(),
                         slices.getConstView(),
                         numberOfSlices );
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
__cuda_callable__
Index
ChunkedEllpackView< Device, Index, RowMajorOrder >::
getSegmentsCount() const
{
   return this->size;
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
__cuda_callable__
Index
ChunkedEllpackView< Device, Index, RowMajorOrder >::
getSegmentSize( const IndexType segmentIdx ) const
{
   if( std::is_same< DeviceType, Devices::Host >::value )
      return details::ChunkedEllpack< IndexType, DeviceType, RowMajorOrder >::getSegmentSizeDirect(
         rowToSliceMapping,
         slices,
         rowToChunkMapping,
         segmentIdx );
   if( std::is_same< DeviceType, Devices::Cuda >::value )
   {
#ifdef __CUDA_ARCH__
      return details::ChunkedEllpack< IndexType, DeviceType, RowMajorOrder >::getSegmentSizeDirect(
         rowToSliceMapping,
         slices,
         rowToChunkMapping,
         segmentIdx );
#else
      return details::ChunkedEllpack< IndexType, DeviceType, RowMajorOrder >::getSegmentSize(
         rowToSliceMapping,
         slices,
         rowToChunkMapping,
         segmentIdx );
#endif
   }
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
__cuda_callable__
Index
ChunkedEllpackView< Device, Index, RowMajorOrder >::
getSize() const
{
   return this->size;
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
__cuda_callable__
Index
ChunkedEllpackView< Device, Index, RowMajorOrder >::
getStorageSize() const
{
   return this->storageSize;
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
__cuda_callable__
Index
ChunkedEllpackView< Device, Index, RowMajorOrder >::
getGlobalIndex( const Index segmentIdx, const Index localIdx ) const
{
   if( std::is_same< DeviceType, Devices::Host >::value )
      return details::ChunkedEllpack< IndexType, DeviceType, RowMajorOrder >::getGlobalIndexDirect(
         rowToSliceMapping,
         slices,
         rowToChunkMapping,
         chunksInSlice,
         segmentIdx,
         localIdx );
   if( std::is_same< DeviceType, Devices::Cuda >::value )
   {
#ifdef __CUDA_ARCH__
      return details::ChunkedEllpack< IndexType, DeviceType, RowMajorOrder >::getGlobalIndexDirect(
         rowToSliceMapping,
         slices,
         rowToChunkMapping,
         chunksInSlice,
         segmentIdx,
         localIdx );
#else
      return details::ChunkedEllpack< IndexType, DeviceType, RowMajorOrder >::getGlobalIndex(
         rowToSliceMapping,
         slices,
         rowToChunkMapping,
         chunksInSlice,
         segmentIdx,
         localIdx );
#endif
   }
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
__cuda_callable__
auto
ChunkedEllpackView< Device, Index, RowMajorOrder >::
getSegmentView( const IndexType segmentIdx ) const -> SegmentViewType
{
   if( std::is_same< DeviceType, Devices::Host >::value )
      return details::ChunkedEllpack< IndexType, DeviceType, RowMajorOrder >::getSegmentViewDirect(
         rowToSliceMapping,
         slices,
         rowToChunkMapping,
         chunksInSlice,
         segmentIdx );
   if( std::is_same< DeviceType, Devices::Cuda >::value )
   {
#ifdef __CUDA_ARCH__
      return details::ChunkedEllpack< IndexType, DeviceType, RowMajorOrder >::getSegmentViewDirect(
         rowToSliceMapping,
         slices,
         rowToChunkMapping,
         chunksInSlice,
         segmentIdx );
#else
      return details::ChunkedEllpack< IndexType, DeviceType, RowMajorOrder >::getSegmentView(
         rowToSliceMapping,
         slices,
         rowToChunkMapping,
         chunksInSlice,
         segmentIdx );
#endif
   }
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
   template< typename Function, typename... Args >
void
ChunkedEllpackView< Device, Index, RowMajorOrder >::
forSegments( IndexType first, IndexType last, Function& f, Args... args ) const
{
   const IndexType chunksInSlice = this->chunksInSlice;
   auto rowToChunkMapping = this->rowToChunkMapping;
   auto rowToSliceMapping = this->rowToSliceMapping;
   auto slices = this->slices;
   auto work = [=] __cuda_callable__ ( IndexType segmentIdx, Args... args ) mutable {
      const IndexType sliceIdx = rowToSliceMapping[ segmentIdx ];

      IndexType firstChunkOfSegment( 0 );
      if( segmentIdx != slices[ sliceIdx ].firstSegment )
      {
         firstChunkOfSegment = rowToChunkMapping[ segmentIdx - 1 ];
      }

      const IndexType lastChunkOfSegment = rowToChunkMapping[ segmentIdx ];
      const IndexType segmentChunksCount = lastChunkOfSegment - firstChunkOfSegment;
      const IndexType sliceOffset = slices[ sliceIdx ].pointer;
      const IndexType chunkSize = slices[ sliceIdx ].chunkSize;

      const IndexType segmentSize = segmentChunksCount * chunkSize;
      bool compute( true );
      if( RowMajorOrder )
      {
         IndexType begin = sliceOffset + firstChunkOfSegment * chunkSize;
         IndexType end = begin + segmentSize;
         IndexType localIdx( 0 );
         for( IndexType j = begin; j < end && compute; j++ )
            f( segmentIdx, localIdx++, j, compute, args...);
      }
      else
      {
         IndexType localIdx( 0 );
         for( IndexType chunkIdx = 0; chunkIdx < segmentChunksCount; chunkIdx++ )
         {
            IndexType begin = sliceOffset + firstChunkOfSegment + chunkIdx;
            IndexType end = begin + chunksInSlice * chunkSize;
            for( IndexType j = begin; j < end && compute; j += chunksInSlice )
            {
               f( segmentIdx, localIdx++, j, compute, args...);
            }
         }
      }
   };
   Algorithms::ParallelFor< DeviceType >::exec( first, last , work, args... );
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
   template< typename Function, typename... Args >
void
ChunkedEllpackView< Device, Index, RowMajorOrder >::
forAll( Function& f, Args... args ) const
{
   this->forSegments( 0, this->getSegmentsCount(), f, args... );
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
   template< typename Fetch, typename Reduction, typename ResultKeeper, typename Real, typename... Args >
void
ChunkedEllpackView< Device, Index, RowMajorOrder >::
segmentsReduction( IndexType first, IndexType last, Fetch& fetch, Reduction& reduction, ResultKeeper& keeper, const Real& zero, Args... args ) const
{
   using RealType = decltype( fetch( IndexType(), IndexType(), IndexType(), std::declval< bool& >(), args... ) );
   if( std::is_same< DeviceType, Devices::Host >::value )
   {
      //segmentsReductionKernel( 0, first, last, fetch, reduction, keeper, zero, args... );
      //return;
      
      for( IndexType segmentIdx = first; segmentIdx < last; segmentIdx++ )
      {
         const IndexType& sliceIndex = rowToSliceMapping[ segmentIdx ];
         TNL_ASSERT_LE( sliceIndex, this->size, "" );
         IndexType firstChunkOfSegment( 0 );
         if( segmentIdx != slices[ sliceIndex ].firstSegment )
            firstChunkOfSegment = rowToChunkMapping[ segmentIdx - 1 ];

         const IndexType lastChunkOfSegment = rowToChunkMapping[ segmentIdx ];
         const IndexType segmentChunksCount = lastChunkOfSegment - firstChunkOfSegment;
         const IndexType sliceOffset = slices[ sliceIndex ].pointer;
         const IndexType chunkSize = slices[ sliceIndex ].chunkSize;

         const IndexType segmentSize = segmentChunksCount * chunkSize;
         RealType aux( zero );
         bool compute( true );
         if( RowMajorOrder )
         {
            IndexType begin = sliceOffset + firstChunkOfSegment * chunkSize;
            IndexType end = begin + segmentSize;
            IndexType localIdx( 0 );
            for( IndexType j = begin; j < end && compute; j++ )
               reduction( aux, fetch( segmentIdx, localIdx++, j, compute, args...) );
         }
         else
         {
            for( IndexType chunkIdx = 0; chunkIdx < segmentChunksCount; chunkIdx++ )
            {
               IndexType begin = sliceOffset + firstChunkOfSegment + chunkIdx;
               IndexType end = begin + chunksInSlice * chunkSize;
               IndexType localIdx( 0 );
               for( IndexType j = begin; j < end && compute; j += chunksInSlice )
                  reduction( aux, fetch( segmentIdx, localIdx++, j, compute, args...) );
            }
         }
         keeper( segmentIdx, aux );
      }
   }
   if( std::is_same< DeviceType, Devices::Cuda >::value )
   {
#ifdef HAVE_CUDA
      //const IndexType chunksCount = this->numberOfSlices * this->chunksInSlice;
      // TODO: This ignores parameters first and last
      const IndexType cudaBlocks = this->numberOfSlices;
      const IndexType cudaGrids = roundUpDivision( cudaBlocks, Cuda::getMaxGridSize() );
      dim3 cudaBlockSize( this->chunksInSlice ), cudaGridSize;
      const IndexType sharedMemory = cudaBlockSize.x * sizeof( RealType );

      for( IndexType gridIdx = 0; gridIdx < cudaGrids; gridIdx++ )
      {
         if( gridIdx == cudaGrids - 1 )
            cudaGridSize.x = cudaBlocks % Cuda::getMaxGridSize();
         ChunkedEllpackSegmentsReductionKernel< ViewType, IndexType, Fetch, Reduction, ResultKeeper, Real, Args...  >
            <<< cudaGridSize, cudaBlockSize, sharedMemory  >>>
            ( *this, gridIdx, first, last, fetch, reduction, keeper, zero, args... );
      }
#endif
   }
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
   template< typename Fetch, typename Reduction, typename ResultKeeper, typename Real, typename... Args >
void
ChunkedEllpackView< Device, Index, RowMajorOrder >::
allReduction( Fetch& fetch, Reduction& reduction, ResultKeeper& keeper, const Real& zero, Args... args ) const
{
   this->segmentsReduction( 0, this->getSegmentsCount(), fetch, reduction, keeper, zero, args... );
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
ChunkedEllpackView< Device, Index, RowMajorOrder >&
ChunkedEllpackView< Device, Index, RowMajorOrder >::
operator=( const ChunkedEllpackView& view )
{
   this->size = view.size;
   this->storageSize = view.storageSize;
   this->chunksInSlice = view.chunksInSlice;
   this->desiredChunkSize = view.desiredChunkSize;
   this->rowToChunkMapping.bind( view.rowToChunkMapping );
   this->chunksToSegmentsMapping.bind( view.chunksToSegmentsMapping );
   this->rowToSliceMapping.bind( view.rowToSliceMapping );
   this->rowPointers.bind( view.rowPointers );
   this->slices.bind( view.slices );
   this->numberOfSlices = view.numberOfSlices;
   return *this;
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
void
ChunkedEllpackView< Device, Index, RowMajorOrder >::
save( File& file ) const
{
   file.save( &this->size );
   file.save( &this->storageSize );
   file.save( &this->chunksInSlice );
   file.save( &this->desiredChunkSize );
   file << this->rowToChunkMapping
        << this->chunksToSegmentsMapping
        << this->rowToSliceMapping
        << this->rowPointers
        << this->slices;
   file.save( &this->numberOfSlices );
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
void
ChunkedEllpackView< Device, Index, RowMajorOrder >::
load( File& file )
{
   file.load( &this->size );
   file.load( &this->storageSize );
   file.load( &this->chunksInSlice );
   file.load( &this->desiredChunkSize );
   file >> this->rowToChunkMapping
        >> this->chunksToSegmentsMapping
        >> this->rowToSliceMapping
        >> this->rowPointers
        >> this->slices;
   file.load( &this->numberOfSlices );
}

template< typename Device,
          typename Index,
          bool RowMajorOrder >
void
ChunkedEllpackView< Device, Index, RowMajorOrder >::
printStructure( std::ostream& str ) const
{
   //const IndexType numberOfSlices = this->getNumberOfSlices();
   str << "Segments count: " << this->getSize() << std::endl
       << "Slices: " << numberOfSlices << std::endl;
   for( IndexType i = 0; i < numberOfSlices; i++ )
      str << "   Slice " << i
          << " : size = " << this->slices.getElement( i ).size
          << " chunkSize = " << this->slices.getElement( i ).chunkSize
          << " firstSegment = " << this->slices.getElement( i ).firstSegment
          << " pointer = " << this->slices.getElement( i ).pointer << std::endl;
   for( IndexType i = 0; i < this->getSize(); i++ )
      str << "Segment " << i
          << " : slice = " << this->rowToSliceMapping.getElement( i )
          << " chunk = " << this->rowToChunkMapping.getElement( i ) << std::endl;
}

#ifdef HAVE_CUDA
template< typename Device,
          typename Index,
          bool RowMajorOrder >
   template< typename Fetch,
             typename Reduction,
             typename ResultKeeper,
             typename Real,
             typename... Args >
__device__
void
ChunkedEllpackView< Device, Index, RowMajorOrder >::
segmentsReductionKernel( IndexType gridIdx,
                         IndexType first,
                         IndexType last,
                         Fetch fetch,
                         Reduction reduction,
                         ResultKeeper keeper,
                         Real zero,
                         Args... args ) const
{
   using RealType = decltype( fetch( IndexType(), IndexType(), IndexType(), std::declval< bool& >(), args... ) );

   const IndexType firstSlice = rowToChunkMapping[ first ] / chunksInSlice;
   const IndexType lastSlice = rowToChunkMapping[ last - 1 ] / chunksInSlice;
   //for( IndexType sliceIdx = firstSlice; sliceIdx < lastSlice; sliceIdx++ )
   {
      const IndexType sliceIdx = gridIdx * Cuda::getMaxGridSize() + blockIdx.x;
      if( sliceIdx >= lastSlice )
         return;

      RealType* chunksResults = Cuda::getSharedMemory< RealType >();
      //for( IndexType threadIdx = 0; threadIdx < 256; threadIdx++ )
      //{
         __shared__ details::ChunkedEllpackSliceInfo< IndexType > sliceInfo;
         if( threadIdx.x == 0 )
            sliceInfo = this->slices[ sliceIdx ];
         chunksResults[ threadIdx.x ] = zero;
         __syncthreads();

   

         const IndexType sliceOffset = sliceInfo.pointer;
         const IndexType chunkSize = sliceInfo.chunkSize;
         const IndexType chunkIdx = sliceIdx * chunksInSlice + threadIdx.x;
         const IndexType segmentIdx = this->chunksToSegmentsMapping[ chunkIdx ];
         IndexType firstChunkOfSegment( 0 );
         if( segmentIdx != sliceInfo.firstSegment )
            firstChunkOfSegment = rowToChunkMapping[ segmentIdx - 1 ];
         IndexType localIdx = ( threadIdx.x - firstChunkOfSegment ) * chunkSize;
         bool compute( true );
          
         if( RowMajorOrder )
         {
            IndexType begin = sliceOffset + threadIdx.x * chunkSize; // threadIdx.x = chunkIdx within the slice
            IndexType end = begin + chunkSize;
            for( IndexType j = begin; j < end && compute; j++ )
               reduction( chunksResults[ threadIdx.x ], fetch( segmentIdx, localIdx++, j, compute, args...) );
         }
         else
         {
            const IndexType begin = sliceOffset + threadIdx.x; // threadIdx.x = chunkIdx within the slice
            const IndexType end = begin + chunksInSlice * chunkSize;
               for( IndexType j = begin; j < end && compute; j += chunksInSlice )
                  reduction( chunksResults[ threadIdx.x ], fetch( segmentIdx, localIdx++, j, compute, args...) );
         }
         __syncthreads();
      //}

      //details::ChunkedEllpackSliceInfo< IndexType > sliceInfo;
      //for( IndexType threadIdx = 0; threadIdx < 256; threadIdx++ )
      //{
         //if( threadIdx == 0 )
         //   sliceInfo = this->slices[ sliceIdx ];
         if( threadIdx.x < sliceInfo.size )
         {
            const IndexType row = sliceInfo.firstSegment + threadIdx.x;
            IndexType chunkIndex( 0 );
            if( threadIdx.x != 0 )
               chunkIndex = this->rowToChunkMapping[ row - 1 ];
            const IndexType lastChunk = this->rowToChunkMapping[ row ];
            RealType result( zero );
            while( chunkIndex < lastChunk )
               reduction( result,  chunksResults[ chunkIndex++ ] );
            keeper( row, result );
         }
      //} // threadIdx
   } // sliceIdx
}
#endif

      } // namespace Segments
   }  // namespace Containers
} // namespace TNL
