// SPDX-FileComment: This file is part of TNL - Template Numerical Library (https://tnl-project.org/)
// SPDX-License-Identifier: MIT

#pragma once

#include "CSRAdaptiveKernel.h"

namespace TNL::Algorithms::SegmentsReductionKernels {

template< typename Index, typename Device >
std::string
CSRAdaptiveKernel< Index, Device >::getKernelType()
{
   return ViewType::getKernelType();
}

template< typename Index, typename Device >
template< typename Segments >
void
CSRAdaptiveKernel< Index, Device >::init( const Segments& segments )
{
   const auto& offsets = segments.getOffsets();

   if( max( offsets ) == 0 ) {
      for( int i = 0; i < MaxValueSizeLog(); i++ ) {
         this->blocksArray[ i ].reset();
         this->view.setBlocks( this->blocksArray[ i ], i );
      }
      return;
   }

   this->template initValueSize< 1 >( offsets );
   this->template initValueSize< 2 >( offsets );
   this->template initValueSize< 4 >( offsets );
   this->template initValueSize< 8 >( offsets );
   this->template initValueSize< 16 >( offsets );
   this->template initValueSize< 32 >( offsets );
   for( int i = 0; i < MaxValueSizeLog(); i++ )
      this->view.setBlocks( this->blocksArray[ i ], i );
}

template< typename Index, typename Device >
void
CSRAdaptiveKernel< Index, Device >::reset()
{
   for( int i = 0; i < MaxValueSizeLog(); i++ ) {
      this->blocksArray[ i ].reset();
      this->view.setBlocks( this->blocksArray[ i ], i );
   }
}

template< typename Index, typename Device >
__cuda_callable__
auto
CSRAdaptiveKernel< Index, Device >::getView() -> ViewType
{
   return this->view;
}

template< typename Index, typename Device >
__cuda_callable__
auto
CSRAdaptiveKernel< Index, Device >::getConstView() const -> ConstViewType
{
   return this->view;
}

template< typename Index, typename Device >
template< typename SegmentsView, typename Fetch, typename Reduction, typename ResultKeeper, typename Value >
void
CSRAdaptiveKernel< Index, Device >::reduceSegments(
   const SegmentsView& segments,
   Index begin,
   Index end,
   Fetch& fetch,
   const Reduction& reduction,
   ResultKeeper& keeper,
   const Value& identity ) const
{
   view.reduceSegments( segments, begin, end, fetch, reduction, keeper, identity );
}

template< typename Index, typename Device >
template< typename SegmentsView, typename Fetch, typename Reduction, typename ResultKeeper, typename Value >
void
CSRAdaptiveKernel< Index, Device >::reduceAllSegments(
   const SegmentsView& segments,
   Fetch& fetch,
   const Reduction& reduction,
   ResultKeeper& keeper,
   const Value& identity ) const
{
   view.reduceAllSegments( segments, fetch, reduction, keeper, identity );
}

template< typename Index, typename Device >
template< int SizeOfValue, typename Offsets >
Index
CSRAdaptiveKernel< Index, Device >::findLimit(
   const Index start,
   const Offsets& offsets,
   const Index size,
   Segments::detail::Type& type )
{
   std::size_t sum = 0;
   for( Index current = start; current < size - 1; current++ ) {
      Index elements = offsets[ current + 1 ] - offsets[ current ];
      sum += elements;
      if( sum > Segments::detail::CSRAdaptiveKernelParameters< SizeOfValue >::StreamedSharedElementsPerWarp() ) {
         if( current - start > 0 ) {
            // extra row
            type = Segments::detail::Type::STREAM;
            return current;
         }
         else {
            // one long row
            if( sum <= 2 * Segments::detail::CSRAdaptiveKernelParameters< SizeOfValue >::MaxAdaptiveElementsPerWarp() )
               type = Segments::detail::Type::VECTOR;
            else
               type = Segments::detail::Type::LONG;
            return current + 1;
         }
      }
   }
   type = Segments::detail::Type::STREAM;
   return size - 1;  // return last row pointer
}

template< typename Index, typename Device >
template< int SizeOfValue, typename Offsets >
void
CSRAdaptiveKernel< Index, Device >::initValueSize( const Offsets& offsets )
{
   using HostOffsetsType = Containers::Vector< typename Offsets::IndexType, Devices::Host, typename Offsets::IndexType >;
   HostOffsetsType hostOffsets;
   hostOffsets = offsets;
   const Index rows = offsets.getSize();
   Index start = 0;
   Index nextStart = 0;

   // Fill blocks
   std::vector< Segments::detail::CSRAdaptiveKernelBlockDescriptor< Index > > inBlocks;
   inBlocks.reserve( rows );

   while( nextStart != rows - 1 ) {
      Segments::detail::Type type;
      nextStart = findLimit< SizeOfValue >( start, hostOffsets, rows, type );
      if( type == Segments::detail::Type::LONG ) {
         const Index blocksCount = inBlocks.size();
         const Index warpsPerCudaBlock = Segments::detail::CSRAdaptiveKernelParameters< SizeOfValue >::CudaBlockSize()
                                       / Backend::getWarpSize( Backend::getDevice() );
         Index warpsLeft = roundUpDivision( blocksCount, warpsPerCudaBlock ) * warpsPerCudaBlock - blocksCount;
         if( warpsLeft == 0 )
            warpsLeft = warpsPerCudaBlock;
         for( Index index = 0; index < warpsLeft; index++ )
            inBlocks.emplace_back( start, Segments::detail::Type::LONG, index, warpsLeft );
      }
      else {
         inBlocks.emplace_back( start, type, nextStart, offsets.getElement( nextStart ), offsets.getElement( start ) );
      }
      start = nextStart;
   }
   inBlocks.emplace_back( nextStart );
   this->blocksArray[ getSizeValueLog( SizeOfValue ) ] = inBlocks;
}

}  // namespace TNL::Algorithms::SegmentsReductionKernels
