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

#pragma once

#include <utility>  // std::pair, std::forward

#include <TNL/Functional.h>  // extension of STL function objects for reduction
#include <TNL/Algorithms/detail/Reduction.h>
#include <TNL/Containers/Expressions/TypeTraits.h>  // RemoveET

namespace TNL::Algorithms {

/**
 * \brief \e reduce implements [(parallel) reduction](https://en.wikipedia.org/wiki/Reduce_(parallel_pattern))
 * for vectors and arrays.
 *
 * Reduction can be used for operations having one or more vectors (or arrays)
 * elements as input and returning one number (or element) as output. Some
 * examples of such operations can be vectors/arrays comparison, vector norm,
 * scalar product of two vectors or computing minimum or maximum. If one needs
 * to know even the position of the smallest or the largest element, the
 * function \ref reduceWithArgument can be used.
 *
 * \tparam Device is a type of the device where the reduction will be performed.
 *         It can be on of the following: \ref TNL::Devices::Sequential,
 *         \ref TNL::Devices::Host, \ref TNL::Devices::Cuda.
 * \tparam Index is a type for indexing.
 * \tparam Result is a type of the reduction result.
 * \tparam Fetch is a lambda function for fetching the input data.
 * \tparam Reduction is a function object performing the reduction.
 *
 * \param begin defines range [begin, end) of indexes which will be used for the reduction.
 * \param end defines range [begin, end) of indexes which will be used for the reduction.
 * \param fetch is a lambda function fetching the input data.
 * \param reduction is a function object defining the reduction operation.
 *                  This can be a user-defined lambda function or an instance of
 *                  some \ref ReductionFunctionObjects.
 * \param identity is the [identity element](https://en.wikipedia.org/wiki/Identity_element)
 *                 for the reduction operation, i.e. element which does not
 *                 change the result of the reduction.
 * \return result of the reduction
 *
 * The `fetch` lambda function takes one argument which is index of the element to be fetched:
 *
 * ```
 * auto fetch = [=] __cuda_callable__ ( Index i ) { return ... };
 * ```
 *
 * The `reduction` lambda function takes two variables which are supposed to be reduced:
 *
 * ```
 * auto reduction = [] __cuda_callable__ ( const Result& a, const Result& b ) { return ... };
 * ```
 *
 * \par Example
 *
 * \include ReductionAndScan/SumExample.cpp
 *
 * \par Output
 *
 * \include SumExample.out
 */
template< typename Device, typename Index, typename Result, typename Fetch, typename Reduction >
Result
reduce( Index begin, Index end, Fetch&& fetch, Reduction&& reduction, const Result& identity )
{
   return detail::Reduction< Device >::reduce(
      begin, end, std::forward< Fetch >( fetch ), std::forward< Reduction >( reduction ), identity );
}

/**
 * \brief Variant of \ref reduce with a reduction function object instead of lambda function.
 *
 * \tparam Device is a type of the device where the reduction will be performed.
 *         It can be on of the following: \ref TNL::Devices::Sequential,
 *         \ref TNL::Devices::Host, \ref TNL::Devices::Cuda.
 * \tparam Index is a type for indexing.
 * \tparam Fetch is a lambda function for fetching the input data.
 * \tparam Reduction is a function object performing the reduction.
 *
 * \param begin defines range [begin, end) of indexes which will be used for the reduction.
 * \param end defines range [begin, end) of indexes which will be used for the reduction.
 * \param fetch is a lambda function fetching the input data.
 * \param reduction is a function object defining the reduction operation,
 *                  it must be an instance of some \ref ReductionFunctionObjects.
 * \return result of the reduction
 *
 * The `fetch` lambda function takes one argument which is index of the element to be fetched:
 *
 * ```
 * auto fetch = [=] __cuda_callable__ ( Index i ) { return ... };
 * ```
 *
 * \par Example
 *
 * \include ReductionAndScan/SumExampleWithFunctional.cpp
 *
 * \par Output
 *
 * \include SumExampleWithFunctional.out
 */
template< typename Device, typename Index, typename Fetch, typename Reduction = TNL::Plus >
auto
reduce( Index begin, Index end, Fetch&& fetch, Reduction&& reduction = TNL::Plus{} )
{
   using Result = Containers::Expressions::RemoveET< decltype( reduction( fetch( 0 ), fetch( 0 ) ) ) >;
   return reduce< Device >(
      begin,
      end,
      std::forward< Fetch >( fetch ),
      std::forward< Reduction >( reduction ),
      reduction.template getIdentity< Result >() );
}

/**
 * \brief Variant of \ref reduce for arrays, views and compatible objects.
 *
 * The referenced \ref reduce function is called with:
 *
 * - `Device`, which is `typename Array::DeviceType` by default, as the `Device` type,
 * - `0` as the beginning of the interval for reduction,
 * - `array.getSize()` as the end of the interval for reduction,
 * - `array.getConstView()` as the `fetch` functor,
 * - `reduction` as the reduction operation,
 * - and `identity` as the identity element of the reduction.
 *
 * \par Example
 *
 * \include Algorithms/reduceArrayExample.cpp
 *
 * \par Output
 *
 * \include reduceArrayExample.out
 */
template< typename Array, typename Device = typename Array::DeviceType, typename Reduction, typename Result >
auto
reduce( const Array& array, Reduction&& reduction, Result identity )
{
   return reduce< Device >(
      static_cast< typename Array::IndexType >( 0 ),
      array.getSize(),
      array.getConstView(),
      std::forward< Reduction >( reduction ),
      identity );
}

/**
 * \brief Variant of \ref reduce for arrays, views and compatible objects.
 *
 * The referenced \ref reduce function is called with:
 *
 * - `Device`, which is `typename Array::DeviceType` by default, as the `Device` type,
 * - `0` as the beginning of the interval for reduction,
 * - `array.getSize()` as the end of the interval for reduction,
 * - `array.getConstView()` as the `fetch` functor,
 * - `reduction` as the reduction operation, it must be an instance of some \ref ReductionFunctionObjects
 * - and the identity element obtained from the reduction function object.
 *
 * \par Example
 *
 * \include Algorithms/reduceArrayExample.cpp
 *
 * \par Output
 *
 * \include reduceArrayExample.out
 */
template< typename Array, typename Device = typename Array::DeviceType, typename Reduction = TNL::Plus >
auto
reduce( const Array& array, Reduction&& reduction = TNL::Plus{} )
{
   using Result = Containers::Expressions::RemoveET< decltype( reduction( array( 0 ), array( 0 ) ) ) >;
   return reduce< Array, Device >( array, std::forward< Reduction >( reduction ), reduction.template getIdentity< Result >() );
}

/**
 * \brief Variant of \ref reduce returning also the position of the element of interest.
 *
 * For example, in case of computing minimal or maximal element in array/vector,
 * the position of the element having given value can be obtained. This method
 * is, however, more flexible.
 *
 * \tparam Device is a type of the device where the reduction will be performed.
 *         It can be on of the following: \ref TNL::Devices::Sequential,
 *         \ref TNL::Devices::Host, \ref TNL::Devices::Cuda.
 * \tparam Index is a type for indexing.
 * \tparam Result is a type of the reduction result.
 * \tparam Reduction is a function object performing the reduction.
 * \tparam Fetch is a lambda function for fetching the input data.
 *
 * \param begin defines range [begin, end) of indexes which will be used for the reduction.
 * \param end defines range [begin, end) of indexes which will be used for the reduction.
 * \param fetch is a lambda function fetching the input data.
 * \param reduction is an instance of lambda function or function object defining the reduction operation.
 *                  This can be a user-defined lambda function or an instance of some \ref ReductionFunctionObjectsWithArgument.
 * \param identity is the [identity element](https://en.wikipedia.org/wiki/Identity_element) for the reduction
 *                 operation, i.e. element which does not change the result of the reduction.
 * \return result of the reduction in a form of `std::pair<Index, Result>` structure, where
 *         `pair.first` is the element position and `pair.second` is the reduction result.
 *
 * The `fetch` lambda function takes one argument which is index of the element to be fetched:
 *
 * ```
 * auto fetch = [=] __cuda_callable__ ( Index i ) { return ... };
 * ```
 *
 * The `reduction` lambda function takes two variables which are supposed to be reduced:
 *
 * ```
 * auto reduction = [] __cuda_callable__ ( Result& a, const Result& b, Index& aIdx, const Index& bIdx ) { return ... };
 * ```
 *
 * \par Example
 *
 * \include ReductionAndScan/ReductionWithArgument.cpp
 *
 * \par Output
 *
 * \include ReductionWithArgument.out
 */
template< typename Device, typename Index, typename Result, typename Fetch, typename Reduction >
std::pair< Result, Index >
reduceWithArgument( Index begin, Index end, Fetch&& fetch, Reduction&& reduction, const Result& identity )
{
   return detail::Reduction< Device >::reduceWithArgument(
      begin, end, std::forward< Fetch >( fetch ), std::forward< Reduction >( reduction ), identity );
}

/**
 * \brief Variant of \ref reduceWithArgument with reduction function object instead of lambda function.
 *
 * \tparam Device is a type of the device where the reduction will be performed.
 *         It can be on of the following: \ref TNL::Devices::Sequential,
 *         \ref TNL::Devices::Host, \ref TNL::Devices::Cuda.
 * \tparam Index is a type for indexing.
 * \tparam Result is a type of the reduction result.
 * \tparam Reduction is a function object performing the reduction.
 * \tparam Fetch is a lambda function for fetching the input data.
 *
 * \param begin defines range [begin, end) of indexes which will be used for the reduction.
 * \param end defines range [begin, end) of indexes which will be used for the reduction.
 * \param fetch is a lambda function fetching the input data.
 * \param reduction is an instance of lambda function or a function object defining the reduction operation.
 *                  This must be an instance of some \ref ReductionFunctionObjectsWithArgument.
 * \return result of the reduction in a form of `std::pair<Index, Result>` structure, where
 *         `pair.first` is the element position and `pair.second` is the reduction result.
 *
 * The `fetch` lambda function takes one argument which is index of the element to be fetched:
 *
 * ```
 * auto fetch = [=] __cuda_callable__ ( Index i ) { return ... };
 * ```
 *
 * \par Example
 *
 * \include ReductionAndScan/ReductionWithArgumentWithFunctional.cpp
 *
 * \par Output
 *
 * \include ReductionWithArgumentWithFunctional.out
 */
template< typename Device, typename Index, typename Fetch, typename Reduction >
auto
reduceWithArgument( Index begin, Index end, Fetch&& fetch, Reduction&& reduction )
{
   using Result = Containers::Expressions::RemoveET< decltype( fetch( 0 ) ) >;
   return reduceWithArgument< Device >(
      begin,
      end,
      std::forward< Fetch >( fetch ),
      std::forward< Reduction >( reduction ),
      reduction.template getIdentity< Result >() );
}

/**
 * \brief Variant of \ref reduceWithArgument for arrays, views and compatible objects.
 *
 * The referenced \ref reduceWithArgument function is called with:
 *
 * - `Device`, which is `typename Array::DeviceType` by default, as the `Device` type,
 * - `0` as the beginning of the interval for reduction,
 * - `array.getSize()` as the end of the interval for reduction,
 * - `array.getConstView()` as the `fetch` functor,
 * - `reduction` as the reduction operation,
 * - and `identity` as the identity element of the reduction.
 *
 * \par Example
 *
 * \include Algorithms/reduceWithArgumentArrayExample.cpp
 *
 * \par Output
 *
 * \include reduceWithArgumentArrayExample.out
 */
template< typename Array, typename Device = typename Array::DeviceType, typename Reduction, typename Result >
auto
reduceWithArgument( const Array& array, Reduction&& reduction, Result identity )
{
   return reduceWithArgument< Device >(
      static_cast< typename Array::IndexType >( 0 ),
      array.getSize(),
      array.getConstView(),
      std::forward< Reduction >( reduction ),
      identity );
}

/**
 * \brief Variant of \ref reduceWithArgument for arrays, views and compatible objects.
 *
 * \e Reduction can be one of \ref TNL::MinWithArg, \ref TNL::MaxWithArg.
 *
 * The referenced \ref reduceWithArgument function is called with:
 *
 * - `Device`, which is `typename Array::DeviceType` by default, as the `Device` type,
 * - `0` as the beginning of the interval for reduction,
 * - `array.getSize()` as the end of the interval for reduction,
 * - `array.getConstView()` as the `fetch` functor,
 * - `reduction` as the reduction operation, it must be an instance of some \ref ReductionFunctionObjectsWithArgument
 * - and the identity element obtained from the reduction function object.
 *
 * \par Example
 *
 * \include Algorithms/reduceWithArgumentArrayExample.cpp
 *
 * \par Output
 *
 * \include reduceWithArgumentArrayExample.out
 */
template< typename Array, typename Device = typename Array::DeviceType, typename Reduction >
auto
reduceWithArgument( const Array& array, Reduction&& reduction )
{
   using Result = Containers::Expressions::RemoveET< decltype( array( 0 ) ) >;
   return reduceWithArgument< Array, Device >(
      array, std::forward< Reduction >( reduction ), reduction.template getIdentity< Result >() );
}

}  // namespace TNL::Algorithms
