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

#pragma once

#include <TNL/Containers/Array.h>
#include <TNL/Algorithms/detail/CudaScanKernel.h>
#include <TNL/Atomic.h>
#include "QuicksortTask.h"

namespace TNL::Algorithms::Sorting::experimental::detail {

#if defined( __CUDACC__ ) || defined( __HIP__ )

template< typename Value, typename Index, typename Device, typename Compare >
__device__
Value
pickPivot( TNL::Containers::ArrayView< Value, Device, Index > src, const Compare& compare )
{
   if( src.getSize() == 1 )
      return src[ 0 ];

   const Value& a = src[ 0 ];
   const Value& b = src[ src.getSize() / 2 ];
   const Value& c = src[ src.getSize() - 1 ];

   if( compare( a, b ) )  // ..a..b..
   {
      if( compare( b, c ) )  // ..a..b..c
         return b;
      else if( compare( c, a ) )  //..c..a..b..
         return a;
      else  //..a..c..b..
         return c;
   }
   else  //..b..a..
   {
      if( compare( a, c ) )  //..b..a..c
         return a;
      else if( compare( c, b ) )  //..c..b..a..
         return b;
      else  //..b..c..a..
         return c;
   }
}

template< typename Value, typename Index, typename Device, typename Compare >
__device__
Index
pickPivotIdx( TNL::Containers::ArrayView< Value, Device, Index > src, const Compare& compare )
{
   if( src.getSize() <= 1 )
      return 0;

   const Value& a = src[ 0 ];
   const Value& b = src[ src.getSize() / 2 ];
   const Value& c = src[ src.getSize() - 1 ];

   if( compare( a, b ) )  // ..a..b..
   {
      if( compare( b, c ) )  // ..a..b..c
         return src.getSize() / 2;
      else if( compare( c, a ) )  //..c..a..b..
         return 0;
      else  //..a..c..b..
         return src.getSize() - 1;
   }
   else  //..b..a..
   {
      if( compare( a, c ) )  //..b..a..c
         return 0;
      else if( compare( c, b ) )  //..c..b..a..
         return src.getSize() / 2;
      else  //..b..c..a..
         return src.getSize() - 1;
   }
}

template< typename Value, typename Index, typename Compare >
__device__
void
countElements(
   Containers::ArrayView< Value, Devices::Cuda, Index > arr,
   const Compare& compare,
   int& smaller,
   int& bigger,
   const Value& pivot )
{
   for( Index i = threadIdx.x; i < arr.getSize(); i += blockDim.x ) {
      const Value& data = arr[ i ];
      if( compare( data, pivot ) )
         smaller++;
      else if( compare( pivot, data ) )
         bigger++;
   }
}

template< typename Value, typename Index, typename Compare >
__device__
void
copyDataShared(
   Containers::ArrayView< Value, Devices::Cuda, Index > src,
   Containers::ArrayView< Value, Devices::Cuda, Index > dst,
   const Compare& compare,
   Value* sharedMem,
   typename QuicksortTask< Index >::Offset smallerStart,
   typename QuicksortTask< Index >::Offset biggerStart,
   int smallerTotal,
   int biggerTotal,
   int smallerOffset,
   int biggerOffset,  // exclusive prefix sum of elements
   const Value& pivot )
{
   for( Index i = threadIdx.x; i < src.getSize(); i += blockDim.x ) {
      const Value& data = src[ i ];
      if( compare( data, pivot ) )
         sharedMem[ smallerOffset++ ] = data;
      else if( compare( pivot, data ) )
         sharedMem[ smallerTotal + biggerOffset++ ] = data;
   }
   __syncthreads();

   for( int i = threadIdx.x; i < smallerTotal + biggerTotal; i += blockDim.x ) {
      if( i < smallerTotal )
         dst[ smallerStart + i ] = sharedMem[ i ];
      else
         dst[ biggerStart + i - smallerTotal ] = sharedMem[ i ];
   }
}

template< typename Value, typename Index, typename Compare >
__device__
void
copyData(
   Containers::ArrayView< Value, Devices::Cuda, Index > src,
   Containers::ArrayView< Value, Devices::Cuda, Index > dst,
   const Compare& compare,
   typename QuicksortTask< Index >::Offset smallerStart,
   typename QuicksortTask< Index >::Offset biggerStart,
   const Value& pivot )
{
   for( Index i = threadIdx.x; i < src.getSize(); i += blockDim.x ) {
      const Value& data = src[ i ];
      if( compare( data, pivot ) )
         dst[ smallerStart++ ] = data;
      else if( compare( pivot, data ) )
         dst[ biggerStart++ ] = data;
   }
}

template< typename Value, typename Index, typename Compare, bool useShared >
__device__
void
cudaPartition(
   Containers::ArrayView< Value, Devices::Cuda, Index > src,
   Containers::ArrayView< Value, Devices::Cuda, Index > dst,
   const Compare& compare,
   Value* sharedMem,
   const Value& pivot,
   int elementsPerBlock,
   QuicksortTask< Index >& task )
{
   using Offset = typename QuicksortTask< Index >::Offset;

   static __shared__ Offset smallerStart;
   static __shared__ Offset biggerStart;

   Index myBegin = static_cast< Index >( elementsPerBlock ) * ( blockIdx.x - task.firstBlock );
   Index myEnd = TNL::min( myBegin + elementsPerBlock, src.getSize() );

   auto srcView = src.getView( myBegin, myEnd );

   int smaller = 0;
   int bigger = 0;
   countElements( srcView, compare, smaller, bigger, pivot );

   // synchronization is in this function already
   using BlockScan = Algorithms::detail::CudaBlockScan< Algorithms::detail::ScanType::Inclusive, 0, TNL::Plus, int >;
   __shared__ typename BlockScan::Storage storage;
   int smallerPrefSumInc = BlockScan::scan( TNL::Plus{}, 0, smaller, threadIdx.x, storage );
   int biggerPrefSumInc = BlockScan::scan( TNL::Plus{}, 0, bigger, threadIdx.x, storage );

   if( threadIdx.x == blockDim.x - 1 )  // last thread in block has sum of all values
   {
      smallerStart = atomicAdd( &( task.dstBegin ), static_cast< Offset >( smallerPrefSumInc ) );
      biggerStart =
         atomicAdd( &( task.dstEnd ), -static_cast< Offset >( biggerPrefSumInc ) ) - static_cast< Offset >( biggerPrefSumInc );
   }
   __syncthreads();

   if( useShared ) {
      static __shared__ int smallerTotal;
      static __shared__ int biggerTotal;
      if( threadIdx.x == blockDim.x - 1 ) {
         smallerTotal = smallerPrefSumInc;
         biggerTotal = biggerPrefSumInc;
      }
      __syncthreads();

      copyDataShared(
         srcView,
         dst,
         compare,
         sharedMem,
         smallerStart,
         biggerStart,
         smallerTotal,
         biggerTotal,
         smallerPrefSumInc - smaller,
         biggerPrefSumInc - bigger,  // exclusive prefix sum of elements
         pivot );
   }
   else {
      Offset destSmaller = smallerStart + static_cast< Offset >( smallerPrefSumInc - smaller );
      Offset destBigger = biggerStart + static_cast< Offset >( biggerPrefSumInc - bigger );
      copyData( srcView, dst, compare, destSmaller, destBigger, pivot );
   }
}

#endif

}  // namespace TNL::Algorithms::Sorting::experimental::detail
