#pragma once
#include "../copy/copy_cfg.cuh"
#include "reduce_utils.cuh"
#include "../../cuda_utils/allocator.cuh"

namespace vidrial {

template<typename _T, int _thread_num,
         typename XTileShape,
         typename GXSlab, typename GxSlab,
         int _smem_buffer_size>
struct ReduceKernelCfg {
    using T = _T;
    static constexpr int thread_num = _thread_num;
    static constexpr int smem_buffer_size = _smem_buffer_size;
    using xTileShape = decltype(wrap(get<1>(XTileShape{})));
    using XCfg = decltype(make_tiling_cfg<T, thread_num>(GXSlab{}.shape(), XTileShape{}, GXSlab{}));
    using xCfg = decltype(make_tiling_cfg<T, thread_num>(GxSlab{}.shape(), xTileShape{}, GxSlab{}));
    // TODO: pick a smarter XFrgThr layout that allows as much thread level reduction as possible
    using XFrgThr = typename XCfg::TileCopy::FrgThr_t;
    // this is a reasonable choice for xFrgThr since it allows us to directly store the result back to gmem
    using xFrgThr = typename xCfg::TileCopy::FrgThr_t;
    using ReduceTile = CTAReduceCfg<T, XTileShape, XFrgThr, xFrgThr, smem_buffer_size, 0>;
    XCfg X;
    xCfg x;
    ReduceTile reduce_tile;
    static constexpr int block_num = size<1,1>(XCfg{}.TileBlock);
    static constexpr int reduce_tile_num = size<1,0>(XCfg{}.TileBlock);
};

template<typename Cfg, typename XT, typename xT, typename Fn>
__global__ void reduce_kernel(Cfg cfg, XT* gX_ptr, xT* gx_ptr, Fn fn) {
    using T = typename Cfg::T;
    static_assert(is_same_v<decltype(cfg.reduce_tile.x_FrgThr), decltype(cfg.x.tile_copy.FrgThr)>, "reduce and copy FrgThr are assumed to be the same");
    static_assert(is_same_v<T, XT> && is_same_v<T, xT>, "Invalid types");
    __shared__ T smem[cfg.smem_buffer_size];
    auto gX_slab = make_tensor(gX_ptr, cfg.X.gSlab);
    auto gx_slab = make_tensor(gx_ptr, cfg.x.gSlab);
    __shared__ alignas(16) T X_tile_smem[size(cfg.X.tile_copy.sTile)];
    __shared__ alignas(16) char reduce_smem_buff[cfg.reduce_tile.smem_buffer_size];
    auto sX_tile = make_tensor(make_smem_ptr(X_tile_smem), cfg.X.sTile);
    auto rx_frg = make_tensor<T>(get<0>(cfg.reduce_tile.x_FrgThr.shape()));
    auto main_loop_iter = [&](int reduce_tile_idx, auto carry_initial) {
        auto tile_coords = make_coord(reduce_tile_idx, blockIdx.x);
        auto gX_tile = slice_rest(gX_slab, cfg.X.TileBlock, tile_coords);
        CTA_copy_tile(cfg.X.tile_copy, gX_tile, sX_tile);
        cp_async_fence(); cp_async_wait<0>(); __syncthreads();
        auto rX_frg = make_tensor<T>(get<0>(cfg.reduce_tile.X_FrgThr.shape()));
        if (threadIdx.x <= size<1>(cfg.reduce_tile.X_FrgThr))
            copy(slice_rest(sX_tile, cfg.reduce_tile.X_FrgThr, threadIdx.x), rX_frg);
        CTA_reduce<carry_initial>(cfg.reduce_tile, rX_frg, rx_frg, fn, reduce_smem_buff);
    };
    main_loop_iter(0, _1{});
    static_assert(cfg.reduce_tile_num == 1, "Multiple tiles along reduce dimension is not supported yet");
    // for (int reduce_tile_idx = 1; reduce_tile_idx < cfg.reduce_tile_num; reduce_tile_idx++)
    //     main_loop_iter(reduce_tile_idx, _0{});
    if (threadIdx.x < size<1>(cfg.reduce_tile.x_FrgThr)) {
        auto gx_tile = slice_rest(gx_slab, cfg.x.TileBlock, blockIdx.x);
        auto gx_frg = slice_rest(gx_tile, cfg.reduce_tile.x_FrgThr, threadIdx.x);
        copy(rx_frg, gx_frg);
    }
}

template<typename Cfg, typename XT, typename xT, typename Fn>
void launch_reduce_kernel(Cfg cfg, XT* gX_ptr, xT* gx_ptr, Fn fn) {
    printf("launching reduce kernel with block_num=%d, thread_num=%d\n", cfg.block_num, cfg.thread_num);
    reduce_kernel<<<cfg.block_num, cfg.thread_num>>>(cfg, gX_ptr, gx_ptr, fn);
}

} // namespace vidrial