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

#pragma once

#include <utility>  // std::forward

#include "Scan.h"
#include "CudaScanKernel.h"

#include <TNL/Assert.h>
#include <TNL/Containers/Array.h>
#include <TNL/Containers/StaticArray.h>
#include <TNL/Algorithms/reduce.h>

namespace TNL::Algorithms::detail {

template< ScanType Type, ScanPhaseType PhaseType >
template< typename InputArray, typename OutputArray, typename Reduction >
typename OutputArray::ValueType
Scan< Devices::Sequential, Type, PhaseType >::perform( const InputArray& input,
                                                       OutputArray& output,
                                                       typename InputArray::IndexType begin,
                                                       typename InputArray::IndexType end,
                                                       typename OutputArray::IndexType outputBegin,
                                                       Reduction&& reduction,
                                                       typename OutputArray::ValueType identity )
{
   using ValueType = typename OutputArray::ValueType;

   // simple sequential algorithm - not split into phases
   ValueType aux = identity;
   if( Type == ScanType::Inclusive ) {
      for( ; begin < end; begin++, outputBegin++ )
         output[ outputBegin ] = aux = reduction( aux, input[ begin ] );
   }
   else  // Exclusive scan
   {
      for( ; begin < end; begin++, outputBegin++ ) {
         const ValueType x = input[ begin ];
         output[ outputBegin ] = aux;
         aux = reduction( aux, x );
      }
   }
   // return the last value of inclusive scan (reduction of the whole input)
   return aux;
}

template< ScanType Type, ScanPhaseType PhaseType >
template< typename InputArray, typename OutputArray, typename Reduction >
auto
Scan< Devices::Sequential, Type, PhaseType >::performFirstPhase( const InputArray& input,
                                                                 OutputArray& output,
                                                                 typename InputArray::IndexType begin,
                                                                 typename InputArray::IndexType end,
                                                                 typename OutputArray::IndexType outputBegin,
                                                                 Reduction&& reduction,
                                                                 typename OutputArray::ValueType identity )
{
   if( end <= begin ) {
      Containers::Array< typename OutputArray::ValueType, Devices::Sequential > block_results( 1 );
      block_results.setValue( identity );
      return block_results;
   }

   if constexpr( PhaseType == ScanPhaseType::WriteInFirstPhase ) {
      // artificial second phase - pre-scan the block
      Containers::Array< typename OutputArray::ValueType, Devices::Sequential > block_results( 2 );
      block_results[ 0 ] = identity;
      block_results[ 1 ] = perform( input, output, begin, end, outputBegin, reduction, identity );
      return block_results;
   }
   else  // ScanPhaseType::WriteInSecondPhase
   {
      // artificial first phase - only reduce the block
      Containers::Array< typename OutputArray::ValueType, Devices::Sequential > block_results( 2 );
      block_results[ 0 ] = identity;
      block_results[ 1 ] = reduce< Devices::Sequential >( begin, end, input, reduction, identity );
      return block_results;
   }
}

template< ScanType Type, ScanPhaseType PhaseType >
template< typename InputArray, typename OutputArray, typename BlockShifts, typename Reduction >
void
Scan< Devices::Sequential, Type, PhaseType >::performSecondPhase( const InputArray& input,
                                                                  OutputArray& output,
                                                                  const BlockShifts& blockShifts,
                                                                  typename InputArray::IndexType begin,
                                                                  typename InputArray::IndexType end,
                                                                  typename OutputArray::IndexType outputBegin,
                                                                  Reduction&& reduction,
                                                                  typename OutputArray::ValueType identity,
                                                                  typename OutputArray::ValueType shift )
{
   if constexpr( PhaseType == ScanPhaseType::WriteInFirstPhase ) {
      // artificial second phase - uniform shift of a pre-scanned block
      shift = reduction( shift, blockShifts[ 0 ] );
      typename InputArray::IndexType outputEnd = outputBegin + end - begin;
      for( typename InputArray::IndexType i = outputBegin; i < outputEnd; i++ )
         output[ i ] = reduction( output[ i ], shift );
   }
   else {  // ScanPhaseType::WriteInSecondPhase
      // artificial second phase - only one block, use the shift as the initial value
      perform( input, output, begin, end, outputBegin, reduction, reduction( shift, blockShifts[ 0 ] ) );
   }
}

template< ScanType Type, ScanPhaseType PhaseType >
template< typename InputArray, typename OutputArray, typename Reduction >
void
Scan< Devices::Host, Type, PhaseType >::perform( const InputArray& input,
                                                 OutputArray& output,
                                                 typename InputArray::IndexType begin,
                                                 typename InputArray::IndexType end,
                                                 typename OutputArray::IndexType outputBegin,
                                                 Reduction&& reduction,
                                                 typename OutputArray::ValueType identity )
{
#ifdef HAVE_OPENMP
   using ValueType = typename OutputArray::ValueType;
   using IndexType = typename InputArray::IndexType;

   if( end <= begin )
      return;

   const IndexType size = end - begin;
   const IndexType max_threads = Devices::Host::getMaxThreadsCount();
   const IndexType block_size = TNL::max( 1024, TNL::roundUpDivision( size, max_threads ) );
   const IndexType blocks = TNL::roundUpDivision( size, block_size );

   if( Devices::Host::isOMPEnabled() && blocks >= 2 ) {
      const int threads = TNL::min( blocks, Devices::Host::getMaxThreadsCount() );
      Containers::Array< ValueType > block_results( blocks + 1 );

      #pragma omp parallel num_threads( threads )
      {
         const int block_idx = omp_get_thread_num();
         const IndexType block_offset = block_idx * block_size;
         const IndexType block_begin = begin + block_offset;
         const IndexType block_end = TNL::min( block_begin + block_size, end );
         const IndexType block_output_begin = outputBegin + block_offset;

         if constexpr( PhaseType == ScanPhaseType::WriteInFirstPhase ) {
            // step 1: pre-scan the block and save the result of the block reduction
            block_results[ block_idx ] = Scan< Devices::Sequential, Type >::perform(
               input, output, block_begin, block_end, block_output_begin, reduction, identity );

            #pragma omp barrier

            // step 2: scan the block results
            #pragma omp single
            {
               Scan< Devices::Sequential, ScanType::Exclusive >::perform(
                  block_results, block_results, 0, blocks + 1, 0, reduction, identity );
            }

            // step 3: uniform shift of the pre-scanned block
            const ValueType block_shift = block_results[ block_idx ];
            const IndexType block_output_end = block_output_begin + block_end - block_begin;
            for( IndexType i = block_output_begin; i < block_output_end; i++ )
               output[ i ] = reduction( output[ i ], block_shift );
         }
         else  // ScanPhaseType::WriteInSecondPhase
         {
            // step 1: per-block reductions, write the result into the buffer
            block_results[ block_idx ] = reduce< Devices::Sequential >( block_begin, block_end, input, reduction, identity );

            #pragma omp barrier

            // step 2: scan the block results
            #pragma omp single
            {
               Scan< Devices::Sequential, ScanType::Exclusive >::perform(
                  block_results, block_results, 0, blocks + 1, 0, reduction, identity );
            }

            // step 3: per-block scan using the block results as initial values
            Scan< Devices::Sequential, Type >::perform(
               input, output, block_begin, block_end, block_output_begin, reduction, block_results[ block_idx ] );
         }
      }
   }
   else
#endif
   {
      Scan< Devices::Sequential, Type >::perform( input, output, begin, end, outputBegin, reduction, identity );
   }
}

template< ScanType Type, ScanPhaseType PhaseType >
template< typename InputArray, typename OutputArray, typename Reduction >
auto
Scan< Devices::Host, Type, PhaseType >::performFirstPhase( const InputArray& input,
                                                           OutputArray& output,
                                                           typename InputArray::IndexType begin,
                                                           typename InputArray::IndexType end,
                                                           typename OutputArray::IndexType outputBegin,
                                                           Reduction&& reduction,
                                                           typename OutputArray::ValueType identity )
{
#ifdef HAVE_OPENMP
   using ValueType = typename OutputArray::ValueType;
   using IndexType = typename InputArray::IndexType;

   if( end <= begin ) {
      Containers::Array< ValueType, Devices::Sequential > block_results( 1 );
      block_results.setValue( identity );
      return block_results;
   }

   const IndexType size = end - begin;
   const IndexType max_threads = Devices::Host::getMaxThreadsCount();
   const IndexType block_size = TNL::max( 1024, TNL::roundUpDivision( size, max_threads ) );
   const IndexType blocks = TNL::roundUpDivision( size, block_size );

   if( Devices::Host::isOMPEnabled() && blocks >= 2 ) {
      const int threads = TNL::min( blocks, Devices::Host::getMaxThreadsCount() );
      Containers::Array< ValueType, Devices::Sequential > block_results( blocks + 1 );

      #pragma omp parallel num_threads( threads )
      {
         const int block_idx = omp_get_thread_num();
         const IndexType block_offset = block_idx * block_size;
         const IndexType block_begin = begin + block_offset;
         const IndexType block_end = TNL::min( block_begin + block_size, end );
         const IndexType block_output_begin = outputBegin + block_offset;

         if constexpr( PhaseType == ScanPhaseType::WriteInFirstPhase )
            // pre-scan the block, write the result of the block reduction into the buffer
            block_results[ block_idx ] = Scan< Devices::Sequential, Type >::perform(
               input, output, block_begin, block_end, block_output_begin, reduction, identity );
         else  // ScanPhaseType::WriteInSecondPhase
            // upsweep: per-block reductions, write the result into the buffer
            block_results[ block_idx ] = reduce< Devices::Sequential >( block_begin, block_end, input, reduction, identity );
      }

      // spine step: scan the block results
      Scan< Devices::Sequential, ScanType::Exclusive >::perform(
         block_results, block_results, 0, blocks + 1, 0, reduction, identity );

      // block_results now contains shift values for each block - to be used in the second phase
      return block_results;
   }
   else
#endif
   {
      return Scan< Devices::Sequential, Type >::performFirstPhase(
         input, output, begin, end, outputBegin, reduction, identity );
   }
}

template< ScanType Type, ScanPhaseType PhaseType >
template< typename InputArray, typename OutputArray, typename BlockShifts, typename Reduction >
void
Scan< Devices::Host, Type, PhaseType >::performSecondPhase( const InputArray& input,
                                                            OutputArray& output,
                                                            const BlockShifts& blockShifts,
                                                            typename InputArray::IndexType begin,
                                                            typename InputArray::IndexType end,
                                                            typename OutputArray::IndexType outputBegin,
                                                            Reduction&& reduction,
                                                            typename OutputArray::ValueType identity,
                                                            typename OutputArray::ValueType shift )
{
#ifdef HAVE_OPENMP
   using ValueType = typename OutputArray::ValueType;
   using IndexType = typename InputArray::IndexType;

   if( end <= begin )
      return;

   const IndexType size = end - begin;
   const IndexType max_threads = Devices::Host::getMaxThreadsCount();
   const IndexType block_size = TNL::max( 1024, TNL::roundUpDivision( size, max_threads ) );
   const IndexType blocks = TNL::roundUpDivision( size, block_size );

   if( Devices::Host::isOMPEnabled() && blocks >= 2 ) {
      const int threads = TNL::min( blocks, Devices::Host::getMaxThreadsCount() );
      #pragma omp parallel num_threads( threads )
      {
         const int block_idx = omp_get_thread_num();
         const IndexType block_offset = block_idx * block_size;
         const IndexType block_begin = begin + block_offset;
         const IndexType block_end = TNL::min( block_begin + block_size, end );
         const IndexType block_output_begin = outputBegin + block_offset;

         const ValueType block_shift = reduction( shift, blockShifts[ block_idx ] );

         if constexpr( PhaseType == ScanPhaseType::WriteInFirstPhase ) {
            // uniform shift of a pre-scanned block
            const IndexType block_output_end = block_output_begin + block_end - block_begin;
            for( IndexType i = block_output_begin; i < block_output_end; i++ )
               output[ i ] = reduction( output[ i ], block_shift );
         }
         else {  // ScanPhaseType::WriteInSecondPhase
            // downsweep: per-block scan using the block results as initial values
            Scan< Devices::Sequential, Type >::perform(
               input, output, block_begin, block_end, block_output_begin, reduction, block_shift );
         }
      }
   }
   else
#endif
   {
      Scan< Devices::Sequential, Type >::performSecondPhase(
         input, output, blockShifts, begin, end, outputBegin, reduction, identity, shift );
   }
}

template< ScanType Type, ScanPhaseType PhaseType >
template< typename InputArray, typename OutputArray, typename Reduction >
void
Scan< Devices::Cuda, Type, PhaseType >::perform( const InputArray& input,
                                                 OutputArray& output,
                                                 typename InputArray::IndexType begin,
                                                 typename InputArray::IndexType end,
                                                 typename OutputArray::IndexType outputBegin,
                                                 Reduction&& reduction,
                                                 typename OutputArray::ValueType identity )
{
   if( end <= begin )
      return;

   detail::CudaScanKernelLauncher< Type, PhaseType, typename OutputArray::ValueType >::perform(
      input, output, begin, end, outputBegin, std::forward< Reduction >( reduction ), identity );
}

template< ScanType Type, ScanPhaseType PhaseType >
template< typename InputArray, typename OutputArray, typename Reduction >
auto
Scan< Devices::Cuda, Type, PhaseType >::performFirstPhase( const InputArray& input,
                                                           OutputArray& output,
                                                           typename InputArray::IndexType begin,
                                                           typename InputArray::IndexType end,
                                                           typename OutputArray::IndexType outputBegin,
                                                           Reduction&& reduction,
                                                           typename OutputArray::ValueType identity )
{
   if( end <= begin ) {
      Containers::Array< typename OutputArray::ValueType, Devices::Cuda > block_results( 1 );
      block_results.setValue( identity );
      return block_results;
   }

   return detail::CudaScanKernelLauncher< Type, PhaseType, typename OutputArray::ValueType >::performFirstPhase(
      input, output, begin, end, outputBegin, std::forward< Reduction >( reduction ), identity );
}

template< ScanType Type, ScanPhaseType PhaseType >
template< typename InputArray, typename OutputArray, typename BlockShifts, typename Reduction >
void
Scan< Devices::Cuda, Type, PhaseType >::performSecondPhase( const InputArray& input,
                                                            OutputArray& output,
                                                            const BlockShifts& blockShifts,
                                                            typename InputArray::IndexType begin,
                                                            typename InputArray::IndexType end,
                                                            typename OutputArray::IndexType outputBegin,
                                                            Reduction&& reduction,
                                                            typename OutputArray::ValueType identity,
                                                            typename OutputArray::ValueType shift )
{
   if( end <= begin )
      return;

   detail::CudaScanKernelLauncher< Type, PhaseType, typename OutputArray::ValueType >::performSecondPhase(
      input, output, blockShifts, begin, end, outputBegin, std::forward< Reduction >( reduction ), identity, shift );
}

}  // namespace TNL::Algorithms::detail
