// -*- C++ -*-
//===-- numeric_impl_hetero.h ---------------------------------------------===//
//
// 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_NUMERIC_IMPL_HETERO_H
#define _ONEDPL_NUMERIC_IMPL_HETERO_H

#include <iterator>
#include "../parallel_backend.h"
#include "../utils.h"
#if _ONEDPL_BACKEND_SYCL
#    include "dpcpp/execution_sycl_defs.h"
#    include "algorithm_impl_hetero.h" // to use __pattern_walk2_brick
#    include "dpcpp/parallel_backend_sycl_utils.h"
#    include "dpcpp/unseq_backend_sycl.h"
#endif

namespace oneapi
{
namespace dpl
{
namespace __internal
{

//------------------------------------------------------------------------
// transform_reduce (version with two binary functions)
//------------------------------------------------------------------------

template <typename _BackendTag, typename _ExecutionPolicy, typename _RandomAccessIterator1,
          typename _RandomAccessIterator2, typename _Tp, typename _BinaryOperation1, typename _BinaryOperation2>
_Tp
__pattern_transform_reduce(__hetero_tag<_BackendTag>, _ExecutionPolicy&& __exec, _RandomAccessIterator1 __first1,
                           _RandomAccessIterator1 __last1, _RandomAccessIterator2 __first2, _Tp __init,
                           _BinaryOperation1 __binary_op1, _BinaryOperation2 __binary_op2)
{
    if (__first1 == __last1)
        return __init;

    using _Functor = unseq_backend::walk_n<_BinaryOperation2>;
    using _RepackedTp = __par_backend_hetero::__repacked_tuple_t<_Tp>;

    auto __n = __last1 - __first1;
    auto __keep1 = oneapi::dpl::__ranges::__get_sycl_range<__par_backend_hetero::access_mode::read>();
    auto __buf1 = __keep1(__first1, __last1);
    auto __keep2 = oneapi::dpl::__ranges::__get_sycl_range<__par_backend_hetero::access_mode::read>();
    auto __buf2 = __keep2(__first2, __first2 + __n);

    return oneapi::dpl::__par_backend_hetero::__parallel_transform_reduce<_RepackedTp,
                                                                          ::std::true_type /*is_commutative*/>(
               _BackendTag{}, ::std::forward<_ExecutionPolicy>(__exec), __binary_op1, _Functor{__binary_op2},
               unseq_backend::__init_value<_RepackedTp>{__init}, // initial value
               __buf1.all_view(), __buf2.all_view())
        .get();
}

//------------------------------------------------------------------------
// transform_reduce (with unary and binary functions)
//------------------------------------------------------------------------

template <typename _BackendTag, typename _ExecutionPolicy, typename _ForwardIterator, typename _Tp,
          typename _BinaryOperation, typename _UnaryOperation>
_Tp
__pattern_transform_reduce(__hetero_tag<_BackendTag>, _ExecutionPolicy&& __exec, _ForwardIterator __first,
                           _ForwardIterator __last, _Tp __init, _BinaryOperation __binary_op,
                           _UnaryOperation __unary_op)
{
    if (__first == __last)
        return __init;

    using _Functor = unseq_backend::walk_n<_UnaryOperation>;
    using _RepackedTp = __par_backend_hetero::__repacked_tuple_t<_Tp>;

    auto __keep = oneapi::dpl::__ranges::__get_sycl_range<__par_backend_hetero::access_mode::read>();
    auto __buf = __keep(__first, __last);

    return oneapi::dpl::__par_backend_hetero::__parallel_transform_reduce<_RepackedTp,
                                                                          ::std::true_type /*is_commutative*/>(
               _BackendTag{}, ::std::forward<_ExecutionPolicy>(__exec), __binary_op, _Functor{__unary_op},
               unseq_backend::__init_value<_RepackedTp>{__init}, // initial value
               __buf.all_view())
        .get();
}

//------------------------------------------------------------------------
// transform_scan
//------------------------------------------------------------------------
template <typename T>
struct ExecutionPolicyWrapper;

#if _ONEDPL_BACKEND_SYCL
template <sycl::access::mode _Mode1, sycl::access::mode _Mode2, typename _T, typename _Allocator>
bool
__iterators_possibly_equal(const sycl_iterator<_Mode1, _T, _Allocator>& __it1,
                           const sycl_iterator<_Mode2, _T, _Allocator>& __it2)
{
    const auto buf1 = __it1.get_buffer();
    const auto buf2 = __it2.get_buffer();

    // If two different sycl iterators belongs to the different sycl buffers, they are different
    if (buf1 != buf2)
        return false;

    // We are unable to compare two sycl_iterator's if one of them is sub_buffer and assume that
    // two different sycl iterators are equal.
    if (buf1.is_sub_buffer() || buf2.is_sub_buffer())
        return true;

    return __it1 == __it2;
}
#endif // _ONEDPL_BACKEND_SYCL

template <typename _BackendTag, typename _ExecutionPolicy, typename _Iterator1, typename _Iterator2,
          typename _UnaryOperation, typename _InitType, typename _BinaryOperation, typename _Inclusive>
_Iterator2
__pattern_transform_scan_base(__hetero_tag<_BackendTag> __tag, _ExecutionPolicy&& __exec, _Iterator1 __first,
                              _Iterator1 __last, _Iterator2 __result, _UnaryOperation __unary_op, _InitType __init,
                              _BinaryOperation __binary_op, _Inclusive)
{
    if (__first == __last)
        return __result;

    const auto __n = __last - __first;

    auto __keep1 = oneapi::dpl::__ranges::__get_sycl_range<__par_backend_hetero::access_mode::read>();
    auto __buf1 = __keep1(__first, __last);

    // This is a temporary workaround for an in-place exclusive scan while the SYCL backend scan pattern is not fixed.
    const bool __is_scan_inplace_exclusive = __n > 1 && !_Inclusive{} && __iterators_possibly_equal(__first, __result);
    if (!__is_scan_inplace_exclusive)
    {
        auto __keep2 = oneapi::dpl::__ranges::__get_sycl_range<__par_backend_hetero::access_mode::write,
                                                               /*_IsNoInitRequested=*/true>();
        auto __buf2 = __keep2(__result, __result + __n);

        oneapi::dpl::__par_backend_hetero::__parallel_transform_scan(
            _BackendTag{}, ::std::forward<_ExecutionPolicy>(__exec), __buf1.all_view(), __buf2.all_view(), __n,
            __unary_op, __init, __binary_op, _Inclusive{})
            .__checked_deferrable_wait();
    }
    else
    {
        assert(__n > 1);
        assert(!_Inclusive{});
        assert(__iterators_possibly_equal(__first, __result));

        using _Type = typename _InitType::__value_type;

        auto __policy =
            __par_backend_hetero::make_wrapped_policy<ExecutionPolicyWrapper>(::std::forward<_ExecutionPolicy>(__exec));

        // Create temporary buffer
        oneapi::dpl::__par_backend_hetero::__buffer<_Type> __tmp_buf(__n);
        auto __first_tmp = __tmp_buf.get();
        auto __last_tmp = __first_tmp + __n;
        auto __keep2 = oneapi::dpl::__ranges::__get_sycl_range<__par_backend_hetero::access_mode::write,
                                                               /*_IsNoInitRequested=*/true>();
        auto __buf2 = __keep2(__first_tmp, __last_tmp);

        // Run main algorithm and save data into temporary buffer
        oneapi::dpl::__par_backend_hetero::__parallel_transform_scan(_BackendTag{}, __policy, __buf1.all_view(),
                                                                     __buf2.all_view(), __n, __unary_op, __init,
                                                                     __binary_op, _Inclusive{})
            .wait();

        // Move data from temporary buffer into results
        oneapi::dpl::__internal::__pattern_walk2_brick(
            __tag, ::std::move(__policy), __first_tmp, __last_tmp, __result,
            oneapi::dpl::__internal::__brick_move<__hetero_tag<_BackendTag>>{});

        //TODO: optimize copy back depending on Iterator, i.e. set_final_data for host iterator/pointer
    }

    return __result + __n;
}
template <typename _BackendTag, typename _ExecutionPolicy, typename _Iterator1, typename _Iterator2,
          typename _UnaryOperation, typename _Type, typename _BinaryOperation, typename _Inclusive>
_Iterator2
__pattern_transform_scan(__hetero_tag<_BackendTag> __tag, _ExecutionPolicy&& __exec, _Iterator1 __first,
                         _Iterator1 __last, _Iterator2 __result, _UnaryOperation __unary_op, _Type __init,
                         _BinaryOperation __binary_op, _Inclusive)
{
    using _RepackedType = __par_backend_hetero::__repacked_tuple_t<_Type>;
    using _InitType = unseq_backend::__init_value<_RepackedType>;

    return __pattern_transform_scan_base(__tag, ::std::forward<_ExecutionPolicy>(__exec), __first, __last, __result,
                                         __unary_op, _InitType{__init}, __binary_op, _Inclusive{});
}

// scan without initial element
template <typename _BackendTag, typename _ExecutionPolicy, typename _Iterator1, typename _Iterator2,
          typename _UnaryOperation, typename _BinaryOperation, typename _Inclusive>
_Iterator2
__pattern_transform_scan(__hetero_tag<_BackendTag> __tag, _ExecutionPolicy&& __exec, _Iterator1 __first,
                         _Iterator1 __last, _Iterator2 __result, _UnaryOperation __unary_op,
                         _BinaryOperation __binary_op, _Inclusive)
{
    using _Type = typename ::std::iterator_traits<_Iterator1>::value_type;
    using _RepackedType = __par_backend_hetero::__repacked_tuple_t<_Type>;
    using _InitType = unseq_backend::__no_init_value<_RepackedType>;

    return __pattern_transform_scan_base(__tag, ::std::forward<_ExecutionPolicy>(__exec), __first, __last, __result,
                                         __unary_op, _InitType{}, __binary_op, _Inclusive{});
}

//------------------------------------------------------------------------
// adjacent_difference
//------------------------------------------------------------------------

// a wrapper for the policy is required to avoid the kernel naming issue
template <typename Name>
struct adjacent_difference_wrapper
{
};

template <typename _BackendTag, typename _ExecutionPolicy, typename _ForwardIterator1, typename _ForwardIterator2,
          typename _BinaryOperation>
_ForwardIterator2
__pattern_adjacent_difference(__hetero_tag<_BackendTag>, _ExecutionPolicy&& __exec, _ForwardIterator1 __first,
                              _ForwardIterator1 __last, _ForwardIterator2 __d_first, _BinaryOperation __op)
{
    auto __n = __last - __first;
    if (__n <= 0)
        return __d_first;

    _ForwardIterator2 __d_last = __d_first + __n;

#if !__SYCL_UNNAMED_LAMBDA__
    // if we have the only element, just copy it according to the specification
    if (__n == 1)
    {
        auto __wrapped_policy = __par_backend_hetero::make_wrapped_policy<adjacent_difference_wrapper>(
            ::std::forward<_ExecutionPolicy>(__exec));

        __internal::__pattern_walk2_brick(__hetero_tag<_BackendTag>{}, __wrapped_policy, __first, __last, __d_first,
                                          __internal::__brick_copy<__hetero_tag<_BackendTag>>{});
    }
    else
#endif
    {
        oneapi::dpl::__internal::__transform_functor<_BinaryOperation, std::true_type> __fn{__op};

        auto __keep1 = oneapi::dpl::__ranges::__get_sycl_range<__par_backend_hetero::access_mode::read>();
        auto __buf1 = __keep1(__first, __last);
        auto __keep2 = oneapi::dpl::__ranges::__get_sycl_range<__par_backend_hetero::access_mode::write,
                                                               /*_IsNoInitRequested=*/true>();
        auto __buf2 = __keep2(__d_first, __d_last);

        using _Function = unseq_backend::walk_adjacent_difference<decltype(__fn)>;

        oneapi::dpl::__par_backend_hetero::__parallel_for(_BackendTag{}, __exec,
                                                          _Function{__fn, static_cast<std::size_t>(__n)}, __n,
                                                          __buf1.all_view(), __buf2.all_view())
            .__checked_deferrable_wait();
    }

    return __d_last;
}

//------------------------------------------------------------------------
// segmented numeric algorithms
//------------------------------------------------------------------------

template <typename _Name>
struct __copy_keys_values_wrapper;

template <typename _BackendTag, typename _ExecutionPolicy, typename _Iterator1, typename _Iterator2,
          typename _Iterator3, typename _Iterator4, typename _BinaryPredicate, typename _BinaryOperator>
typename std::iterator_traits<_Iterator3>::difference_type
__pattern_reduce_by_segment(__hetero_tag<_BackendTag> __tag, _ExecutionPolicy&& __exec, _Iterator1 __keys_first,
                            _Iterator1 __keys_last, _Iterator2 __values_first, _Iterator3 __out_keys_first,
                            _Iterator4 __out_values_first, _BinaryPredicate __binary_pred, _BinaryOperator __binary_op)
{
    const std::size_t __n = std::distance(__keys_first, __keys_last);

    if (__n == 0)
        return 0;

    if (__n == 1)
    {
        __brick_copy<__hetero_tag<_BackendTag>> __copy_op{};

        oneapi::dpl::__internal::__pattern_walk2_n(
            __tag, oneapi::dpl::__par_backend_hetero::make_wrapped_policy<__copy_keys_values_wrapper>(__exec),
            oneapi::dpl::make_zip_iterator(__keys_first, __values_first), 1,
            oneapi::dpl::make_zip_iterator(__out_keys_first, __out_values_first), __copy_op);

        return 1;
    }

    auto __keep_keys = oneapi::dpl::__ranges::__get_sycl_range<__par_backend_hetero::access_mode::read>();
    auto __keys = __keep_keys(__keys_first, __keys_last);
    auto __keep_values = oneapi::dpl::__ranges::__get_sycl_range<__par_backend_hetero::access_mode::read>();
    auto __values = __keep_values(__values_first, __values_first + __n);
    auto __keep_key_outputs = oneapi::dpl::__ranges::__get_sycl_range<__par_backend_hetero::access_mode::read_write>();
    auto __out_keys = __keep_key_outputs(__out_keys_first, __out_keys_first + __n);
    auto __keep_value_outputs =
        oneapi::dpl::__ranges::__get_sycl_range<__par_backend_hetero::access_mode::read_write>();
    auto __out_values = __keep_value_outputs(__out_values_first, __out_values_first + __n);
    return oneapi::dpl::__par_backend_hetero::__parallel_reduce_by_segment(
        _BackendTag{}, std::forward<_ExecutionPolicy>(__exec), __keys.all_view(), __values.all_view(),
        __out_keys.all_view(), __out_values.all_view(), __binary_pred, __binary_op);
}

template <typename _BackendTag, typename _Policy, typename _InputIterator1, typename _InputIterator2,
          typename _OutputIterator, typename _BinaryPredicate, typename _Operator, typename _Inclusive,
          typename _InitType>
_OutputIterator
__pattern_scan_by_segment_impl(__hetero_tag<_BackendTag>, _Policy&& __policy, _InputIterator1 __first1,
                               _InputIterator1 __last1, _InputIterator2 __first2, _OutputIterator __result,
                               _BinaryPredicate __binary_pred, _Operator __binary_op, _Inclusive, _InitType __init)
{
    const auto __n = std::distance(__first1, __last1);

    // Check for empty element ranges
    if (__n <= 0)
        return __result;

    namespace __bknd = oneapi::dpl::__par_backend_hetero;

    auto __keep_keys = oneapi::dpl::__ranges::__get_sycl_range<__bknd::access_mode::read>();
    auto __key_buf = __keep_keys(__first1, __last1);
    auto __keep_values = oneapi::dpl::__ranges::__get_sycl_range<__bknd::access_mode::read>();
    auto __value_buf = __keep_values(__first2, __first2 + __n);
    auto __keep_value_outputs = oneapi::dpl::__ranges::__get_sycl_range<__bknd::access_mode::read_write>();
    auto __value_output_buf = __keep_value_outputs(__result, __result + __n);

    __bknd::__parallel_scan_by_segment<_Inclusive::value>(
        _BackendTag{}, std::forward<_Policy>(__policy), __key_buf.all_view(), __value_buf.all_view(),
        __value_output_buf.all_view(), __binary_pred, __binary_op, __init);
    return __result + __n;
}

template <typename _BackendTag, typename _Policy, typename _InputIterator1, typename _InputIterator2,
          typename _OutputIterator, typename _BinaryPredicate, typename _Operator, typename _Inclusive, typename _T>
_OutputIterator
__pattern_scan_by_segment(__hetero_tag<_BackendTag> __tag, _Policy&& __policy, _InputIterator1 __first1,
                          _InputIterator1 __last1, _InputIterator2 __first2, _OutputIterator __result,
                          _BinaryPredicate __binary_pred, _Operator __binary_op, _Inclusive __is_inclusive, _T __init)
{
    return __pattern_scan_by_segment_impl(__tag, std::forward<_Policy>(__policy), __first1, __last1, __first2, __result,
                                          __binary_pred, __binary_op, __is_inclusive,
                                          unseq_backend::__init_value<_T>{__init});
}

template <typename _BackendTag, typename _Policy, typename _InputIterator1, typename _InputIterator2,
          typename _OutputIterator, typename _BinaryPredicate, typename _Operator, typename _Inclusive>
_OutputIterator
__pattern_scan_by_segment(__hetero_tag<_BackendTag> __tag, _Policy&& __policy, _InputIterator1 __first1,
                          _InputIterator1 __last1, _InputIterator2 __first2, _OutputIterator __result,
                          _BinaryPredicate __binary_pred, _Operator __binary_op, _Inclusive __is_inclusive)
{
    return __pattern_scan_by_segment_impl(__tag, std::forward<_Policy>(__policy), __first1, __last1, __first2, __result,
                                          __binary_pred, __binary_op, __is_inclusive, unseq_backend::__no_init_value{});
}

} // namespace __internal
} // namespace dpl
} // namespace oneapi

#endif // _ONEDPL_NUMERIC_IMPL_HETERO_H
