// -*- C++ -*-
//===----------------------------------------------------------------------===//
//
// Copyright (C) Intel Corporation
//
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
// This file incorporates work covered by the following copyright and permission
// notice:
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
//
//===----------------------------------------------------------------------===//

#ifndef _ONEDPL_INTERNAL_OMP_PARALLEL_TRANSFORM_REDUCE_H
#define _ONEDPL_INTERNAL_OMP_PARALLEL_TRANSFORM_REDUCE_H

#include "util.h"

namespace oneapi
{
namespace dpl
{
namespace __omp_backend
{

//------------------------------------------------------------------------
// parallel_transform_reduce
//
// Notation:
//      r(i,j,init) returns reduction of init with reduction over [i,j)
//      u(i) returns f(i,i+1,identity) for a hypothetical left identity element
//      of r c(x,y) combines values x and y that were the result of r or u
//------------------------------------------------------------------------

template <class _RandomAccessIterator, class _UnaryOp, class _Value, class _Combiner, class _Reduction>
_Value
__transform_reduce_body(_RandomAccessIterator __first, _RandomAccessIterator __last, _UnaryOp __unary_op, _Value __init,
                        _Combiner __combiner, _Reduction __reduction)
{
    const std::size_t __num_threads = omp_get_num_threads();
    const std::size_t __size = __last - __first;

    // Serial fallback. The condition is equivalent to __chunk_partitioner returning 1 chunk
    // after reserving the first __num_threads items to initialize the accumulators.
    // The separate check avoids the risk of __first + __num_threads pointing after __last.
    if (__size <= __num_threads + __default_chunk_size)
    {
        return __reduction(__first, __last, std::move(__init));
    }

    auto __policy = oneapi::dpl::__omp_backend::__chunk_partitioner(__first + __num_threads, __last, __num_threads,
                                                                    __default_chunk_size);

    // Here, we cannot use OpenMP UDR because we must store the init value in
    // the combiner and it will be used several times. Although there should be
    // the only one; we manually generate the identity elements for each thread.
    std::vector<_Value> __accums;
    __accums.reserve(__num_threads);

    // initialize accumulators for all threads
    for (std::size_t __i = 0; __i < __num_threads; ++__i)
    {
        __accums.emplace_back(__unary_op(__first + __i));
    }

    // main loop
    _ONEDPL_PRAGMA(omp taskloop shared(__accums))
    for (std::size_t __chunk = 0; __chunk < __policy.__n_chunks; ++__chunk)
    {
        oneapi::dpl::__omp_backend::__process_chunk(
            __policy, __first + __num_threads, __chunk, [&](auto __chunk_first, auto __chunk_last) {
                auto __thread_num = omp_get_thread_num();
                __accums[__thread_num] = __reduction(__chunk_first, __chunk_last, std::move(__accums[__thread_num]));
            });
    }

    // combine by accumulators
    for (std::size_t __i = 0; __i < __num_threads; ++__i)
    {
        __init = __combiner(std::move(__init), std::move(__accums[__i]));
    }

    return std::move(__init);
}

template <class _ExecutionPolicy, class _RandomAccessIterator, class _UnaryOp, class _Value, class _Combiner,
          class _Reduction>
_Value
__parallel_transform_reduce(oneapi::dpl::__internal::__omp_backend_tag, _ExecutionPolicy&&,
                            _RandomAccessIterator __first, _RandomAccessIterator __last, _UnaryOp __unary_op,
                            _Value __init, _Combiner __combiner, _Reduction __reduction)
{
    _Value __result = std::move(__init);
    if (omp_in_parallel())
    {
        // We don't create a nested parallel region in an existing parallel
        // region: just create tasks
        __result = oneapi::dpl::__omp_backend::__transform_reduce_body(__first, __last, __unary_op, std::move(__result),
                                                                       __combiner, __reduction);
    }
    else
    {
        // Create a parallel region, and a single thread will create tasks
        // for the region.
        _ONEDPL_PRAGMA(omp parallel)
        _ONEDPL_PRAGMA(omp single nowait)
        {
            __result = oneapi::dpl::__omp_backend::__transform_reduce_body(
                __first, __last, __unary_op, std::move(__result), __combiner, __reduction);
        }
    }

    return __result;
}

} // namespace __omp_backend
} // namespace dpl
} // namespace oneapi
#endif // _ONEDPL_INTERNAL_OMP_PARALLEL_TRANSFORM_REDUCE_H
