/*
 *  Copyright (c) Intel Corporation
 *
 *  Copyright 2008-2013 NVIDIA Corporation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#ifndef _ONEDPL_INCLUSIVE_SCAN_BY_SEGMENT_IMPL_H
#define _ONEDPL_INCLUSIVE_SCAN_BY_SEGMENT_IMPL_H

#include "by_segment_extension_defs.h"

#include "../pstl/glue_numeric_impl.h"
#include "../pstl/glue_algorithm_impl.h" // TODO: replace the use of transform with __pattern_walk3

#include "../pstl/parallel_backend.h"
#include "function.h"
#include "../pstl/utils.h"
#include "../pstl/functional_impl.h" // for oneapi::dpl::identity
#include "../pstl/iterator_impl.h"   // make_zip_iterator

namespace oneapi
{
namespace dpl
{
namespace __internal
{

template <typename Name>
class InclusiveScan1;

template <class _Tag, typename Policy, typename InputIterator1, typename InputIterator2, typename OutputIterator,
          typename BinaryPredicate, typename BinaryOperator>
OutputIterator
__pattern_inclusive_scan_by_segment(_Tag, Policy&& policy, InputIterator1 first1, InputIterator1 last1,
                                    InputIterator2 first2, OutputIterator result, BinaryPredicate binary_pred,
                                    BinaryOperator binary_op)
{
    static_assert(__internal::__is_host_dispatch_tag_v<_Tag>);

    const auto n = std::distance(first1, last1);

    // Check for empty and single element ranges
    if (n <= 0)
        return result;
    if (n == 1)
    {
        *result = *first2;
        return result + 1;
    }

    using FlagType = unsigned int;
    using ValueType = typename std::iterator_traits<InputIterator2>::value_type;

    oneapi::dpl::__par_backend::__buffer<FlagType> _mask(n);
    auto mask = _mask.get();

    mask[0] = 1;

    oneapi::dpl::transform(policy, first1, last1 - 1, first1 + 1, _mask.get() + 1,
                           oneapi::dpl::__internal::__not_pred<BinaryPredicate>(binary_pred));

    oneapi::dpl::inclusive_scan(std::forward<Policy>(policy), make_zip_iterator(first2, _mask.get()),
        make_zip_iterator(first2, _mask.get()) + n, make_zip_iterator(result, _mask.get()),
        oneapi::dpl::__internal::__segmented_scan_fun<ValueType, FlagType, BinaryOperator>{binary_op});

    return result + n;
}

#if _ONEDPL_BACKEND_SYCL
template <typename _BackendTag, typename Policy, typename InputIterator1, typename InputIterator2,
          typename OutputIterator, typename BinaryPredicate, typename BinaryOperator>
OutputIterator
__pattern_inclusive_scan_by_segment(__internal::__hetero_tag<_BackendTag> __tag, Policy&& policy, InputIterator1 first1,
                                    InputIterator1 last1, InputIterator2 first2, OutputIterator result,
                                    BinaryPredicate binary_pred, BinaryOperator binary_op)
{
    return __pattern_scan_by_segment(__tag, std::forward<Policy>(policy), first1, last1, first2, result, binary_pred,
                                     binary_op, /*_Inclusive*/ std::true_type{});
}

#endif
} // namespace __internal

template <typename Policy, typename InputIterator1, typename InputIterator2, typename OutputIterator,
          typename BinaryPredicate, typename BinaryOperator>
oneapi::dpl::__internal::__enable_if_execution_policy<Policy, OutputIterator>
inclusive_scan_by_segment(Policy&& policy, InputIterator1 first1, InputIterator1 last1, InputIterator2 first2,
                          OutputIterator result, BinaryPredicate binary_pred, BinaryOperator binary_op)
{
    const auto __dispatch_tag = oneapi::dpl::__internal::__select_backend(policy, first1, first2, result);

    return __internal::__pattern_inclusive_scan_by_segment(__dispatch_tag, std::forward<Policy>(policy), first1, last1,
                                                           first2, result, binary_pred, binary_op);
}

template <typename Policy, typename InputIter1, typename InputIter2, typename OutputIter, typename BinaryPredicate>
oneapi::dpl::__internal::__enable_if_execution_policy<Policy, OutputIter>
inclusive_scan_by_segment(Policy&& policy, InputIter1 first1, InputIter1 last1, InputIter2 first2, OutputIter result,
                          BinaryPredicate binary_pred)
{
    using T = typename std::iterator_traits<InputIter2>::value_type;

    return inclusive_scan_by_segment(std::forward<Policy>(policy), first1, last1, first2, result, binary_pred,
                                     std::plus<T>());
}

template <typename Policy, typename InputIter1, typename InputIter2, typename OutputIter>
oneapi::dpl::__internal::__enable_if_execution_policy<Policy, OutputIter>
inclusive_scan_by_segment(Policy&& policy, InputIter1 first1, InputIter1 last1, InputIter2 first2, OutputIter result)
{
    using T = typename std::iterator_traits<InputIter1>::value_type;

    return inclusive_scan_by_segment(std::forward<Policy>(policy), first1, last1, first2, result, std::equal_to<T>());
}

template <typename Policy, typename InputIterator1, typename InputIterator2, typename OutputIterator,
          typename BinaryPredicate, typename BinaryOperator>
oneapi::dpl::__internal::__enable_if_execution_policy<Policy, OutputIterator>
inclusive_scan_by_key(Policy&& policy, InputIterator1 first1, InputIterator1 last1, InputIterator2 first2,
                      OutputIterator result, BinaryPredicate binary_pred, BinaryOperator binary_op)
{
    return inclusive_scan_by_segment(std::forward<Policy>(policy), first1, last1, first2, result, binary_pred,
                                     binary_op);
}

template <typename Policy, typename InputIter1, typename InputIter2, typename OutputIter, typename BinaryPredicate>
oneapi::dpl::__internal::__enable_if_execution_policy<Policy, OutputIter>
inclusive_scan_by_key(Policy&& policy, InputIter1 first1, InputIter1 last1, InputIter2 first2, OutputIter result,
                      BinaryPredicate binary_pred)
{
    return inclusive_scan_by_segment(std::forward<Policy>(policy), first1, last1, first2, result, binary_pred);
}

template <typename Policy, typename InputIter1, typename InputIter2, typename OutputIter>
oneapi::dpl::__internal::__enable_if_execution_policy<Policy, OutputIter>
inclusive_scan_by_key(Policy&& policy, InputIter1 first1, InputIter1 last1, InputIter2 first2, OutputIter result)
{
    return inclusive_scan_by_segment(std::forward<Policy>(policy), first1, last1, first2, result);
}
} // end namespace dpl
} // end namespace oneapi

#endif // _ONEDPL_INCLUSIVE_SCAN_BY_SEGMENT_IMPL_H
