/*
 * GridTools
 *
 * Copyright (c) 2014-2019, ETH Zurich
 * All rights reserved.
 *
 * Please, refer to the LICENSE file in the root directory.
 * SPDX-License-Identifier: BSD-3-Clause
 */
#pragma once

#include <utility>

#include "../../common/defs.hpp"
#include "../../common/generic_metafunctions/utility.hpp"
#include "../../common/gt_math.hpp"
#include "../../common/host_device.hpp"
#include "../../common/tuple_util.hpp"
#include "../../meta.hpp"
#include "../dim.hpp"
#include "../execution_types.hpp"
#include "../sid/blocked_dim.hpp"
#include "../sid/concept.hpp"
#include "k_cache.hpp"

namespace gridtools {
    namespace cuda {
        namespace fused_mss_loop_cuda_impl_ {
            GT_FUNCTION_DEVICE void syncthreads(std::true_type) { __syncthreads(); }
            GT_FUNCTION_DEVICE void syncthreads(std::false_type) {}

            template <class Deref, class Info, class Ptr, class Strides, class Validator>
            GT_FUNCTION_DEVICE void exec_cells(
                Info, Ptr const &ptr, Strides const &strides, Validator const &validator) {
                device::for_each<typename Info::cells_t>([&](auto cell) GT_FORCE_INLINE_LAMBDA {
                    syncthreads(cell.need_sync());
                    if (validator(cell.extent()))
                        cell.template operator()<Deref>(ptr, strides);
                });
            }

            template <class Deref,
                class Mss,
                class Sizes,
                class = typename has_k_caches<Mss>::type,
                class = typename Mss::execution_t>
            struct k_loop_f;

            template <class Deref, class Mss, class Sizes, class Execution>
            struct k_loop_f<Deref, Mss, Sizes, std::true_type, Execution> {
                Sizes m_sizes;

                template <class Ptr, class Strides, class Validator>
                GT_FUNCTION_DEVICE void operator()(Ptr ptr, Strides const &strides, Validator validator) const {
                    k_caches_type<Mss> k_caches;
                    auto mixed_ptr = hymap::device::merge(k_caches.ptr(), wstd::move(ptr));
                    tuple_util::device::for_each(
                        [&](const int_t size, auto info) GT_FORCE_INLINE_LAMBDA {
#ifdef __HIPCC__
// unroll factor estimate based on GT perftests on AMD Mi50
#pragma unroll 3
#else
// unroll factor estimate based on COSMO dycore performance on NVIDIA V100: no unrolling
#endif
                            for (int_t i = 0; i < size; ++i) {
                                exec_cells<Deref>(info, mixed_ptr, strides, validator);
                                k_caches.slide(info.k_step());
                                info.inc_k(mixed_ptr.secondary(), strides);
                            }
                        },
                        m_sizes,
                        Mss::interval_infos());
                }
            };

            template <class Deref, class Mss, class Sizes, class Execution>
            struct k_loop_f<Deref, Mss, Sizes, std::false_type, Execution> {
                Sizes m_sizes;

                template <class Ptr, class Strides, class Validator>
                GT_FUNCTION_DEVICE void operator()(Ptr ptr, Strides const &strides, Validator validator) const {
                    tuple_util::device::for_each(
                        [&](const int_t size, auto info) GT_FORCE_INLINE_LAMBDA {
#ifdef __HIPCC__
// unroll factor estimate based on GT perftests on AMD Mi50
#pragma unroll 3
#else
// unroll factor estimate based on COSMO dycore performance on NVIDIA V100: no unrolling
#endif
                            for (int_t i = 0; i < size; ++i) {
                                exec_cells<Deref>(info, ptr, strides, validator);
                                info.inc_k(ptr, strides);
                            }
                        },
                        m_sizes,
                        Mss::interval_infos());
                }
            };

            template <class Deref, class Mss, class Sizes, int_t BlockSize>
            struct k_loop_f<Deref, Mss, Sizes, std::false_type, execute::parallel_block<BlockSize>> {
                Sizes m_sizes;

                template <class Ptr, class Strides, class Validator>
                GT_FUNCTION_DEVICE void operator()(Ptr ptr, Strides const &strides, Validator validator) const {
                    int_t cur = -(int_t)blockIdx.z * BlockSize;
                    sid::shift(ptr, sid::get_stride<dim::k>(strides), -cur);
                    tuple_util::device::for_each(
                        [&](int_t size, auto info) GT_FORCE_INLINE_LAMBDA {
                            if (cur >= BlockSize)
                                return;
                            int_t lim = math::min(cur + size, BlockSize) - math::max(cur, 0);
                            cur += size;
#pragma unroll
                            for (int_t i = 0; i < BlockSize; ++i) {
                                if (i >= lim)
                                    break;
                                exec_cells<Deref>(info, ptr, strides, validator);
                                info.inc_k(ptr, strides);
                            }
                        },
                        m_sizes,
                        Mss::interval_infos());
                }
            };

            template <class Sid, class KLoop>
            struct kernel_f {
                sid::ptr_holder_type<Sid> m_ptr_holder;
                sid::strides_type<Sid> m_strides;
                KLoop k_loop;

                template <class Validator>
                GT_FUNCTION_DEVICE void operator()(int_t i_block, int_t j_block, Validator validator) const {
                    auto ptr = m_ptr_holder();
                    sid::shift(ptr, sid::get_stride<sid::blocked_dim<dim::i>>(m_strides), blockIdx.x);
                    sid::shift(ptr, sid::get_stride<sid::blocked_dim<dim::j>>(m_strides), blockIdx.y);
                    sid::shift(ptr, sid::get_stride<dim::i>(m_strides), i_block);
                    sid::shift(ptr, sid::get_stride<dim::j>(m_strides), j_block);
                    k_loop(wstd::move(ptr), m_strides, wstd::move(validator));
                }
            };

            template <class Deref, class Mss, class Grid, class Composite>
            auto make_kernel_fun(Grid const &grid, Composite &composite) {
                sid::ptr_diff_type<Composite> offset{};
                auto strides = sid::get_strides(composite);
                sid::shift(offset, sid::get_stride<dim::k>(strides), grid.k_start(Mss::interval(), Mss::execution()));
                auto origin = sid::get_origin(composite) + offset;
                auto k_sizes = stage_matrix::make_k_sizes(Mss::interval_infos(), grid);
                using k_sizes_t = decltype(k_sizes);
                using k_loop_t = k_loop_f<Deref, Mss, k_sizes_t>;

                return kernel_f<Composite, k_loop_t>{
                    sid::get_origin(composite) + offset, std::move(strides), {std::move(k_sizes)}};
            }
        } // namespace fused_mss_loop_cuda_impl_
        using fused_mss_loop_cuda_impl_::make_kernel_fun;
    } // namespace cuda
} // namespace gridtools
