#ifndef MIGRAPHX_GUARD_RTGLIB_MULTI_INDEX_HPP
#define MIGRAPHX_GUARD_RTGLIB_MULTI_INDEX_HPP

#include <migraphx/config.hpp>
#include <migraphx/gpu/device/launch.hpp>
#include <migraphx/gpu/device/shape.hpp>
#include <migraphx/functional.hpp>

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {
namespace device {

template <index_int N>
struct multi_index
{
    using hip_index = hip_array<index_int, N>;
    hip_index id{};
    hip_index stride{};

    MIGRAPHX_DEVICE_CONSTEXPR auto for_stride(hip_index n) const
    {
        // f should return void, but this helps with type deduction
        return [=](auto f) -> decltype(f(hip_index{})) {
            for(hip_index i = id; i < n; i = n.carry(i + stride))
            {
                f(i);
            }
        };
    }
};

template <class ForStride>
__device__ __host__ auto deduce_for_stride(ForStride fs) -> decltype(fs(id{}));

MIGRAPHX_DEVICE_CONSTEXPR multi_index<1> make_multi_index(index_int i, index_int n)
{
    return {{i}, {n}};
}

template <index_int N>
MIGRAPHX_DEVICE_CONSTEXPR multi_index<N>
make_multi_index(const hip_shape<N>& s, index_int i, index_int n)
{
    return {s.multi(i), s.multi(n)};
}

template <index_int N>
MIGRAPHX_DEVICE_CONSTEXPR multi_index<N>
make_multi_index(const hip_shape<N>& s, index_int i, const hip_array<index_int, N>& n)
{
    return {s.multi(i), n};
}

template <index_int N>
inline auto mi_nglobal(const hip_shape<N>& s, index_int nlocal)
{
    assert(s.standard);
    assert(s.elements() > 0);
    index_int n       = s.elements();
    index_int groups  = (n + nlocal - 1) / nlocal;
    index_int nglobal = std::min<index_int>(128, groups) * nlocal;

    assert(groups > 0);
    assert(nglobal > 0);
    auto nglobal_multi = s.multi(nglobal);

    // Skip checking this, since this will cause metadata to not be generated
    // for some unknown reason.
    //
    // assert(std::any_of(nglobal_multi.begin(), nglobal_multi.end(), [](auto x){return x>0;}));

    return nglobal_multi;
}

template <index_int N>
inline auto mi_nlocal(const hip_shape<N>& s, index_int local)
{
    assert(s.standard);
    assert(s.elements() > 0);
    auto nlocal_multi = s.multi(local);

    // Skip checking this, since this will cause metadata to not be generated
    // for some unknown reason.
    //
    // assert(std::any_of(nlocal_multi.begin(), nlocal_multi.end(), [](auto x){return x>0;}));

    return nlocal_multi;
}

template <index_int N>
inline auto mi_launch(hipStream_t stream, const hip_shape<N>& global, index_int nlocal = 1024)
{
    auto nglobal_multi = mi_nglobal(global, nlocal);
    auto nglobal       = global.index(nglobal_multi);

    return [=](auto f) {
        launch(stream, nglobal, nlocal)([=](auto idx) __device__ {
            auto midx = make_multi_index(global, idx.global, nglobal_multi);
            f(idx, midx.for_stride(global.lens));
        });
    };
}

template <index_int N>
inline auto mi_launch(hipStream_t stream,
                      const hip_shape<N>& global,
                      const hip_shape<N>& local,
                      index_int nlocal = 1024)
{
    auto nglobal_multi = mi_nglobal(global, 1);
    auto nglobal       = global.index(nglobal_multi);
    auto nlocal_multi  = mi_nlocal(local, nlocal);

    return [=](auto f) {
        launch(stream, nglobal * nlocal, nlocal)([=](auto idx) {
            // TODO: Use fast div for nlocal
            auto midx = make_multi_index(global, idx.global / nlocal, nglobal_multi);
            auto lidx = make_multi_index(local, idx.local, nlocal_multi);
            f(idx, midx.for_stride(global.lens), lidx.for_stride(local.lens));
        });
    };
}

template <index_int N>
inline auto mi_gs_launch(hipStream_t stream, const hip_shape<N>& global, index_int nlocal = 1024)
{
    return [=](auto f) {
        mi_launch(stream, global, nlocal)([=](auto, auto g) { g([&](auto i) { f(i); }); });
    };
}

} // namespace device
} // namespace gpu
} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx

#endif
