// Copyright (c) 2022-2025 Advanced Micro Devices, Inc. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#ifndef ROCPRIM_DEVICE_DETAIL_DEVICE_SCAN_BY_KEY_HPP_
#define ROCPRIM_DEVICE_DETAIL_DEVICE_SCAN_BY_KEY_HPP_

#include "device_config_helper.hpp"
#include "device_scan_common.hpp"
#include "lookback_scan_state.hpp"

#include "../../block/block_discontinuity.hpp"
#include "../../block/block_load.hpp"
#include "../../block/block_scan.hpp"
#include "../../block/block_store.hpp"
#include "../../config.hpp"
#include "../../detail/binary_op_wrappers.hpp"
#include "../../intrinsics/thread.hpp"
#include "../../types/tuple.hpp"
#include "../detail/ordered_block_id.hpp"

#include <type_traits>

BEGIN_ROCPRIM_NAMESPACE

namespace detail
{
template<bool         Exclusive,
         unsigned int block_size,
         unsigned int items_per_thread,
         typename key_type,
         typename result_type,
         ::rocprim::block_load_method load_keys_method,
         ::rocprim::block_load_method load_values_method>
struct load_values_flagged
{
    using block_load_keys
        = ::rocprim::block_load<key_type, block_size, items_per_thread, load_keys_method>;

    using block_discontinuity = ::rocprim::block_discontinuity<key_type, block_size>;

    using block_load_values
        = ::rocprim::block_load<result_type, block_size, items_per_thread, load_keys_method>;

    union storage_type
    {
        struct
        {
            typename block_load_keys::storage_type     load;
            typename block_discontinuity::storage_type flag;
        } keys;
        typename block_load_values::storage_type load_values;
    };

    // Load flagged values
    // - if the scan is exclusive, the last item of each segment (range where the keys compare equal)
    //   is flagged and reset to the initial value. Adding the last item of the range to the
    //   second to last using `headflag_scan_op_wrapper` will return the initial_value,
    //   which is exactly what should be saved at the start of the next range.
    // - if the scan is inclusive, then the first item of each segment is marked, and it will
    //   restart the scan from that value
    template<typename KeyIterator, typename ValueIterator, typename CompareFunction>
    ROCPRIM_DEVICE
    void load(
        KeyIterator        keys_input,
        ValueIterator      values_input,
        CompareFunction    compare,
        const result_type  initial_value,
        const unsigned int flat_block_id,
        const size_t       starting_block,
        const size_t       number_of_blocks,
        const unsigned int flat_thread_id,
        const size_t       size,
        bool               use_last_keys,
        const typename std::iterator_traits<KeyIterator>::value_type* const __restrict__ last_keys,
        rocprim::tuple<result_type, bool> (&wrapped_values)[items_per_thread],
        storage_type& storage)
    {
        constexpr static unsigned int items_per_block = items_per_thread * block_size;
        const unsigned int            block_offset    = flat_block_id * items_per_block;
        KeyIterator                   block_keys      = keys_input + block_offset;
        ValueIterator                 block_values    = values_input + block_offset;

        key_type    keys[items_per_thread];
        result_type values[items_per_thread];
        bool        flags[items_per_thread];

        auto not_equal = [compare](const auto& a, const auto& b) mutable { return !compare(a, b); };
        const auto global_block_id = starting_block + flat_block_id;

        const auto flag_segment_boundaries = [&]()
        {
            if constexpr(Exclusive)
            {
                const key_type tile_successor
                    = global_block_id < number_of_blocks - 1
                          ? static_cast<key_type>(block_keys[items_per_block])
                          : static_cast<key_type>(*block_keys);
                block_discontinuity{}.flag_tails(flags,
                                                 tile_successor,
                                                 keys,
                                                 not_equal,
                                                 storage.keys.flag);
            }
            else
            { // Inclusive
                const key_type tile_predecessor
                    = global_block_id > 0
                          ? (use_last_keys ? static_cast<key_type>(last_keys[global_block_id - 1])
                                           : static_cast<key_type>(block_keys[-1]))
                          : static_cast<key_type>(*block_keys);
                block_discontinuity{}.flag_heads(flags,
                                                 tile_predecessor,
                                                 keys,
                                                 not_equal,
                                                 storage.keys.flag);
            }
        };

        if(global_block_id < number_of_blocks - 1)
        {
            block_load_keys{}.load(block_keys, keys, storage.keys.load);

            flag_segment_boundaries();
            // Reusing shared memory for loading values
            ::rocprim::syncthreads();

            block_load_values{}.load(block_values, values, storage.load_values);

            ROCPRIM_UNROLL
            for(unsigned int i = 0; i < items_per_thread; ++i)
            {
                rocprim::get<0>(wrapped_values[i])
                    = (Exclusive && flags[i]) ? initial_value : values[i];
                rocprim::get<1>(wrapped_values[i]) = flags[i];
            }
        }
        else
        {
            const unsigned int valid_in_last_block
                = static_cast<unsigned int>(size - items_per_block * (number_of_blocks - 1));

            block_load_keys{}.load(
                block_keys,
                keys,
                valid_in_last_block,
                *block_keys, // Any value is okay, so discontinuity doesn't access undefined items
                storage.keys.load);

            flag_segment_boundaries();
            // Reusing shared memory for loading values
            ::rocprim::syncthreads();

            block_load_values{}.load(block_values,
                                     values,
                                     valid_in_last_block,
                                     storage.load_values);

            ROCPRIM_UNROLL
            for(unsigned int i = 0; i < items_per_thread; ++i)
            {
                if(flat_thread_id * items_per_thread + i >= valid_in_last_block)
                {
                    break;
                }

                rocprim::get<0>(wrapped_values[i])
                    = (Exclusive && flags[i]) ? initial_value : values[i];
                rocprim::get<1>(wrapped_values[i]) = flags[i];
            }
        }
    }
};

template<unsigned int block_size,
         unsigned int items_per_thread,
         typename result_type,
         ::rocprim::block_store_method store_method>
struct unwrap_store
{
    using block_store_values
        = ::rocprim::block_store<result_type, block_size, items_per_thread, store_method>;

    using storage_type = typename block_store_values::storage_type;

    template<typename OutputIterator>
    ROCPRIM_DEVICE
    void store(OutputIterator     output,
               const unsigned int flat_block_id,
               const size_t       starting_block,
               const size_t       number_of_blocks,
               const unsigned int flat_thread_id,
               const size_t       size,
               const rocprim::tuple<result_type, bool> (&wrapped_values)[items_per_thread],
               storage_type& storage)
    {
        constexpr static unsigned int items_per_block = items_per_thread * block_size;
        const unsigned int            block_offset    = flat_block_id * items_per_block;
        OutputIterator                block_output    = output + block_offset;

        result_type thread_values[items_per_thread];

        if(starting_block + flat_block_id < number_of_blocks - 1)
        {
            ROCPRIM_UNROLL
            for(unsigned int i = 0; i < items_per_thread; ++i)
            {
                thread_values[i] = rocprim::get<0>(wrapped_values[i]);
            }

            // Reusing shared memory from scan to perform store
            rocprim::syncthreads();

            block_store_values{}.store(block_output, thread_values, storage);
        }
        else
        {
            const unsigned int valid_in_last_block
                = static_cast<unsigned int>(size - items_per_block * (number_of_blocks - 1));

            ROCPRIM_UNROLL
            for(unsigned int i = 0; i < items_per_thread; ++i)
            {
                if(flat_thread_id * items_per_thread + i >= valid_in_last_block)
                {
                    break;
                }

                thread_values[i] = rocprim::get<0>(wrapped_values[i]);
            }

            // Reusing shared memory from scan to perform store
            rocprim::syncthreads();

            block_store_values{}.store(block_output, thread_values, valid_in_last_block, storage);
        }
    }
};

    template<typename ArchConfig,
             lookback_scan_determinism Determinism,
             bool                      Exclusive,
             typename KeyInputIterator,
             typename InputIterator,
             typename OutputIterator,
             typename ResultType,
             typename CompareFunction,
             typename BinaryFunction,
             typename LookbackScanState,
             typename WrappedBlockId>
ROCPRIM_DEVICE ROCPRIM_FORCE_INLINE auto
        device_scan_by_key_kernel_impl(KeyInputIterator,
                                       InputIterator,
                                       OutputIterator,
                                       ResultType,
                                       const CompareFunction,
                                       const BinaryFunction,
                                       LookbackScanState,
                                       const size_t,
                                       const size_t,
                                       const size_t,
                                       const rocprim::tuple<ResultType, bool>* const,
                                       bool,
                                       const typename std::iterator_traits<KeyInputIterator>::value_type* const __restrict__,
                                       WrappedBlockId)
            -> std::enable_if_t<!is_lookback_kernel_runnable<LookbackScanState>()>
    {
        // No need to build the kernel with sleep on a device that does not require it
    }

    template<typename ArchConfig,
             lookback_scan_determinism Determinism,
             bool                      Exclusive,
             typename KeyInputIterator,
             typename InputIterator,
             typename OutputIterator,
             typename ResultType,
             typename CompareFunction,
             typename BinaryFunction,
             typename LookbackScanState,
             typename WrappedBlockId>
    ROCPRIM_DEVICE ROCPRIM_FORCE_INLINE auto device_scan_by_key_kernel_impl(
        KeyInputIterator                              keys,
        InputIterator                                 values,
        OutputIterator                                output,
        ResultType                                    initial_value,
        const CompareFunction                         compare,
        const BinaryFunction                          scan_op,
        LookbackScanState                             scan_state,
        const size_t                                  size,
        const size_t                                  starting_block,
        const size_t                                  number_of_blocks,
        const rocprim::tuple<ResultType, bool>* const previous_last_value,
        bool use_last_keys,
        const typename std::iterator_traits<KeyInputIterator>::value_type* const __restrict__ last_keys,
        WrappedBlockId ordered_bid)
        -> std::enable_if_t<is_lookback_kernel_runnable<LookbackScanState>()>
    {
        using result_type = ResultType;
        static_assert(std::is_same<rocprim::tuple<ResultType, bool>,
                                   typename LookbackScanState::value_type>::value,
                      "value_type of LookbackScanState must be tuple of result type and flag");
        static constexpr scan_by_key_config_params params = ArchConfig::params;

        constexpr auto block_size         = params.kernel_config.block_size;
        constexpr auto items_per_thread   = params.kernel_config.items_per_thread;
        constexpr auto load_keys_method   = params.block_load_method;
        constexpr auto load_values_method = load_keys_method;

        using key_type     = typename std::iterator_traits<KeyInputIterator>::value_type;
        using load_flagged = load_values_flagged<Exclusive,
                                                 block_size,
                                                 items_per_thread,
                                                 key_type,
                                                 result_type,
                                                 load_keys_method,
                                                 load_values_method>;

        auto wrapped_op    = headflag_scan_op_wrapper<result_type, bool, BinaryFunction>{scan_op};
        using wrapped_type = rocprim::tuple<result_type, bool>;

        using block_scan_type
            = ::rocprim::block_scan<wrapped_type, block_size, params.block_scan_method>;

        constexpr auto store_method = params.block_store_method;
        using store_unwrap = unwrap_store<block_size, items_per_thread, result_type, store_method>;

        ROCPRIM_SHARED_MEMORY union
        {
            typename WrappedBlockId::storage_type  ordered_bid;
            typename load_flagged::storage_type    load;
            typename block_scan_type::storage_type scan;
            typename store_unwrap::storage_type    store;
        } storage;

        const auto flat_thread_id = ::rocprim::detail::block_thread_id<0>();
        const auto flat_block_id
            = ordered_bid.get(flat_thread_id,
                              storage.ordered_bid); // ::rocprim::detail::block_id<0>();

        ::rocprim::syncthreads();

        // Load input
        wrapped_type wrapped_values[items_per_thread];
        load_flagged{}.load(keys,
                            values,
                            compare,
                            initial_value,
                            flat_block_id,
                            starting_block,
                            number_of_blocks,
                            flat_thread_id,
                            size,
                            use_last_keys,
                            last_keys,
                            wrapped_values,
                            storage.load);

        // Reusing the storage from load to perform the scan
        ::rocprim::syncthreads();

        // Perform look back scan scan
        if(flat_block_id == 0)
        {
            auto wrapped_initial_value = rocprim::make_tuple(initial_value, false);

            // previous_last_value is used to pass the value from the previous grid, if this is a
            // multi grid launch
            if(previous_last_value != nullptr)
            {
                if constexpr(Exclusive)
                {
                    rocprim::get<0>(wrapped_initial_value) = rocprim::get<0>(*previous_last_value);
                }
                else if(flat_thread_id == 0)
                {
                    wrapped_values[0] = wrapped_op(*previous_last_value, wrapped_values[0]);
                }
            }

            wrapped_type reduction;
            if constexpr(Exclusive)
            {
                lookback_block_scan<Exclusive, block_scan_type>(wrapped_values,
                                                                wrapped_initial_value,
                                                                reduction,
                                                                storage.scan,
                                                                wrapped_op);
            }
            else
            {
                lookback_block_scan<Exclusive, block_scan_type>(wrapped_values,
                                                                reduction,
                                                                storage.scan,
                                                                wrapped_op);
            }

            if(flat_thread_id == 0)
            {
                scan_state.set_complete(flat_block_id, reduction);
            }
        }
        else
        {
            auto prefix_op
                = lookback_scan_prefix_op<wrapped_type,
                                          decltype(wrapped_op),
                                          decltype(scan_state),
                                          Determinism>{flat_block_id, wrapped_op, scan_state};

            // Scan of block values
            lookback_block_scan<Exclusive, block_scan_type>(wrapped_values,
                                                            storage.scan,
                                                            prefix_op,
                                                            wrapped_op);
        }

        // Store output
        // synchronization is inside the function after unwrapping
        store_unwrap{}.store(output,
                             flat_block_id,
                             starting_block,
                             number_of_blocks,
                             flat_thread_id,
                             size,
                             wrapped_values,
                             storage.store);
    }
    } // namespace detail

    END_ROCPRIM_NAMESPACE

#endif // ROCPRIM_DEVICE_DETAIL_DEVICE_SCAN_BY_KEY_HPP_
