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

#pragma once

#include <memory>  // std::unique_ptr

//#define CUDA_REDUCTION_PROFILING

#include <TNL/Assert.h>
#include <TNL/Algorithms/Reduction2D.h>
#include <TNL/Algorithms/copy.h>
#include <TNL/Algorithms/detail/CudaReduction2DKernel.h>

#ifdef CUDA_REDUCTION_PROFILING
   #include <TNL/Timer.h>
   #include <iostream>
#endif

namespace TNL::Algorithms {

template< typename Result, typename Fetch, typename Reduction, typename Index, typename Output >
void constexpr Reduction2D< Devices::Sequential >::reduce(
   Result identity,
   Fetch fetch,
   Reduction reduction,
   Index size,
   int n,
   Output result )
{
   TNL_ASSERT_GT( size, 0, "The size of datasets must be positive." );
   TNL_ASSERT_GT( n, 0, "The number of datasets must be positive." );

   constexpr int block_size = 128;
   const int blocks = size / block_size;

#if defined( __CUDA_ARCH__ )
   for( int k = 0; k < n; k++ )
      result( k ) = identity;

   for( int b = 0; b < blocks; b++ ) {
      const Index offset = b * block_size;
      for( int k = 0; k < n; k++ ) {
         for( int i = 0; i < block_size; i++ )
            result( k ) = reduction( result( k ), fetch( offset + i, k ) );
      }
   }

   for( int k = 0; k < n; k++ ) {
      for( Index i = blocks * block_size; i < size; i++ )
         result( k ) = reduction( result( k ), fetch( i, k ) );
   }
#else
   if( blocks > 1 ) {
      // initialize array for unrolled results
      // (it is accessed as a row-major matrix with n rows and 4 columns)
      std::unique_ptr< Result[] > r{ new Result[ n * 4 ] };
      for( int k = 0; k < n * 4; k++ )
         r[ k ] = identity;

      // main reduction (explicitly unrolled loop)
      for( int b = 0; b < blocks; b++ ) {
         const Index offset = b * block_size;
         for( int k = 0; k < n; k++ ) {
            Result* _r = r.get() + 4 * k;
            for( int i = 0; i < block_size; i += 4 ) {
               _r[ 0 ] = reduction( _r[ 0 ], fetch( offset + i, k ) );
               _r[ 1 ] = reduction( _r[ 1 ], fetch( offset + i + 1, k ) );
               _r[ 2 ] = reduction( _r[ 2 ], fetch( offset + i + 2, k ) );
               _r[ 3 ] = reduction( _r[ 3 ], fetch( offset + i + 3, k ) );
            }
         }
      }

      // reduction of the last, incomplete block (not unrolled)
      for( int k = 0; k < n; k++ ) {
         Result* _r = r.get() + 4 * k;
         for( Index i = blocks * block_size; i < size; i++ )
            _r[ 0 ] = reduction( _r[ 0 ], fetch( i, k ) );
      }

      // reduction of unrolled results
      for( int k = 0; k < n; k++ ) {
         Result* _r = r.get() + 4 * k;
         _r[ 0 ] = reduction( _r[ 0 ], _r[ 1 ] );
         _r[ 0 ] = reduction( _r[ 0 ], _r[ 2 ] );
         _r[ 0 ] = reduction( _r[ 0 ], _r[ 3 ] );

         // copy the result into the output parameter
         result( k ) = _r[ 0 ];
      }
   }
   else {
      for( int k = 0; k < n; k++ )
         result( k ) = identity;

      for( int b = 0; b < blocks; b++ ) {
         const Index offset = b * block_size;
         for( int k = 0; k < n; k++ ) {
            for( int i = 0; i < block_size; i++ )
               result( k ) = reduction( result( k ), fetch( offset + i, k ) );
         }
      }

      for( int k = 0; k < n; k++ ) {
         for( Index i = blocks * block_size; i < size; i++ )
            result( k ) = reduction( result( k ), fetch( i, k ) );
      }
   }
#endif
}

template< typename Result, typename Fetch, typename Reduction, typename Index, typename Output >
void
Reduction2D< Devices::Host >::reduce( Result identity, Fetch fetch, Reduction reduction, Index size, int n, Output result )
{
   if( size < 0 )
      throw std::invalid_argument( "Reduction2D: The size of datasets must be non-negative." );
   if( n < 0 )
      throw std::invalid_argument( "Reduction2D: The number of datasets must be non-negative." );

#ifdef HAVE_OPENMP
   constexpr int block_size = 128;
   const int blocks = size / block_size;

   if( Devices::Host::isOMPEnabled() && blocks >= 2 ) {
      const int threads = TNL::min( blocks, Devices::Host::getMaxThreadsCount() );
      #pragma omp parallel num_threads(threads)
      {
         // first thread initializes the result array
         #pragma omp single nowait
         {
            for( int k = 0; k < n; k++ )
               result( k ) = identity;
         }

         // initialize array for thread-local results
         // (it is accessed as a row-major matrix with n rows and 4 columns)
         std::unique_ptr< Result[] > r{ new Result[ n * 4 ] };
         for( int k = 0; k < n * 4; k++ )
            r[ k ] = identity;

         #pragma omp for nowait
         for( int b = 0; b < blocks; b++ ) {
            const Index offset = b * block_size;
            for( int k = 0; k < n; k++ ) {
               Result* _r = r.get() + 4 * k;
               for( int i = 0; i < block_size; i += 4 ) {
                  _r[ 0 ] = reduction( _r[ 0 ], fetch( offset + i, k ) );
                  _r[ 1 ] = reduction( _r[ 1 ], fetch( offset + i + 1, k ) );
                  _r[ 2 ] = reduction( _r[ 2 ], fetch( offset + i + 2, k ) );
                  _r[ 3 ] = reduction( _r[ 3 ], fetch( offset + i + 3, k ) );
               }
            }
         }

         // the first thread that reaches here processes the last, incomplete block
         #pragma omp single nowait
         {
            for( int k = 0; k < n; k++ ) {
               Result* _r = r.get() + 4 * k;
               for( Index i = blocks * block_size; i < size; i++ )
                  _r[ 0 ] = reduction( _r[ 0 ], fetch( i, k ) );
            }
         }

         // local reduction of unrolled results
         for( int k = 0; k < n; k++ ) {
            Result* _r = r.get() + 4 * k;
            _r[ 0 ] = reduction( _r[ 0 ], _r[ 1 ] );
            _r[ 0 ] = reduction( _r[ 0 ], _r[ 2 ] );
            _r[ 0 ] = reduction( _r[ 0 ], _r[ 3 ] );
         }

         // inter-thread reduction of local results
         #pragma omp critical
         {
            for( int k = 0; k < n; k++ )
               result( k ) = reduction( result( k ), r[ 4 * k ] );
         }
      }
   }
   else
#endif
   {
      Reduction2D< Devices::Sequential >::reduce( identity, fetch, reduction, size, n, result );
   }
}

template< typename Result, typename Fetch, typename Reduction, typename Index, typename Output >
void
Reduction2D< Devices::Cuda >::reduce( Result identity, Fetch fetch, Reduction reduction, Index size, int n, Output hostResult )
{
   if( size < 0 )
      throw std::invalid_argument( "Reduction2D: The size of datasets must be non-negative." );
   if( n < 0 )
      throw std::invalid_argument( "Reduction2D: The number of datasets must be non-negative." );

#ifdef CUDA_REDUCTION_PROFILING
   Timer timer;
   timer.reset();
   timer.start();
#endif

   // start the reduction on the GPU
   Result* deviceAux1 = nullptr;
   const int reducedSize = detail::CudaReduction2DKernelLauncher( identity, fetch, reduction, size, n, deviceAux1 );

#ifdef CUDA_REDUCTION_PROFILING
   timer.stop();
   std::cout << "   Reduction2D of " << n << " datasets on GPU to size " << reducedSize << " took " << timer.getRealTime()
             << " sec.\n";
   timer.reset();
   timer.start();
#endif

   // transfer the reduced data from device to host
   std::unique_ptr< Result[] > resultArray{ new Result[ n * reducedSize ] };
   copy< void, Devices::Cuda >( resultArray.get(), deviceAux1, n * reducedSize );

#ifdef CUDA_REDUCTION_PROFILING
   timer.stop();
   std::cout << "   Transferring data to CPU took " << timer.getRealTime() << " sec.\n";
   timer.reset();
   timer.start();
#endif

   // finish the reduction on the host
   auto fetchFinish = [ & ]( int i, int k )
   {
      return resultArray[ i + k * reducedSize ];
   };
   Reduction2D< Devices::Sequential >::reduce( identity, fetchFinish, reduction, reducedSize, n, hostResult );

#ifdef CUDA_REDUCTION_PROFILING
   timer.stop();
   std::cout << "   Reduction2D of small data set on CPU took " << timer.getRealTime() << " sec.\n";
#endif
}

}  // namespace TNL::Algorithms
