#pragma once
#include "../sympow/sympow_cfg.cuh"
#include "../reduce/reduce_utils.cuh"

namespace vidrial {

// The bwd process of sympow involves objects of the following shapes:
// Y_grad: [[d_tile, ...], b_tile]
// X_i: [[d_tile], b_tile]
// X_i_grad: [[d_tile], b_tile]
// Y_grad_for_i: [[d_tile, ...], b_tile]
// Y_grad_for_i_warp_reduced: [warp_reduce([d_tile, ...]), b_tile]
//
// The following steps are performed for each factor i:
// i. Y_grad_for_i = Y_grad * \prod_{j \neq i} X_j
// ii. Y_grad_for_i_warp_reduced = warp_reduce(Y_grad_for_i)
// iii. X_i_grad = CTA_reduce(Y_grad_for_i_warp_reduced)

template <int dim, auto... Is, typename YTensor, typename XTensors>
CUTE_HOST_DEVICE void tprod_bcast_multiply_all_except(index_sequence<Is...>, YTensor& Y, const XTensors& Xs) {
    (..., [&]() {
        if constexpr (Is != dim)
            tprod_bcast_multiply<Is>(Y, get<Is>(Xs));
    }());
}

template <int dim, typename YTensor, typename XTensors>
CUTE_HOST_DEVICE void tprod_bcast_multiply_all_except(YTensor& Y, const XTensors& Xs) {
    constexpr int rnk = rank(decltype(Xs){});
    tprod_bcast_multiply_all_except<dim>(make_index_sequence<rnk>{}, Y, Xs);
}

template<int keep_dim, typename YShape, typename YFrgThr, typename XFrgThr, typename T, int smem_buffer_size>
CUTE_DEVICE auto make_reduce_cfg() {
    constexpr int thread_num = size<1>(YFrgThr{});
    // _Y: [reduce_dims, keep_feature, batch]
    auto _YShape = make_shape(drop<keep_dim>(get<0>(YShape{})), get<0,keep_dim>(YShape{}), get<1>(YShape{}));
    auto _YLayout = make_layout(_YShape);
    // Convert Y -> _Y
    auto Y_convert0 = layout_insert<keep_dim>(get<0>(_YLayout), get<1>(_YLayout));
    auto Y_convert = make_layout(Y_convert0, get<2>(_YLayout)); // [[d_tile, ...], b_tile] with keep_dim has the highest stride in feature dimensions
    auto _Y_FrgThr = Y_convert.compose(YFrgThr{});
    auto reduce_cfg = CTAReduceCfg<T, decltype(_YShape), decltype(_Y_FrgThr), XFrgThr, smem_buffer_size, 0>{};
    return reduce_cfg;
}

// ---- Thread and warp reduce ----
template<int keep_dim>
CUTE_DEVICE auto tprod_wrap_reduce(auto& cfg, const auto& Y_frg) {
    auto Y_frg_shape = Y_frg.layout().shape();
    auto _Y_frg_shape = make_shape(drop<keep_dim>(get<0>(Y_frg_shape)), get<0,keep_dim>(Y_frg_shape), get<1>(Y_frg_shape));
    auto _Y_frg_layout = make_layout(_Y_frg_shape);
    auto _Y_frg_convert0 = layout_insert<keep_dim>(get<0>(_Y_frg_layout), get<1>(_Y_frg_layout));
    auto _Y_frg_convert = make_layout(_Y_frg_convert0, get<2>(_Y_frg_layout));
    auto _Y_frg = make_tensor(Y_frg.data(), _Y_frg_layout);
    auto fn = SumCallable{};
    auto rB_frg = thread_reduce(cfg.thread_reduce, fn, _Y_frg);
    warp_reduce(cfg.warp_reduce, fn, rB_frg);
    return rB_frg;
}

// ---- CTA reduce sum, including thread/warp/smem reduce ----
template<int keep_dim>
CUTE_DEVICE void tprod_CTA_reduce_sum(auto& cfg, const auto& Y_frg, auto& X_frg, char* sBuff) {
    auto rB_frg = tprod_wrap_reduce<keep_dim>(cfg, Y_frg);
    auto fn = SumCallable{};
    CTA_smem_reduce<false>(cfg.smem_reduce, rB_frg, X_frg, fn, sBuff);
}

// ---- Smem reduce sum ----
template<int keep_dim>
CUTE_DEVICE void tprod_smem_reduce_sum(auto& cfg, auto& rB_frg, auto& X_frg, int tid, char* sBuff) {
    auto fn = SumCallable{};
    CTA_smem_reduce<false>(cfg.smem_reduce, rB_frg, X_frg, fn, sBuff);
}

// ---- Compute and accumulate Y_grad for a factor ----
template<auto Is, typename T, typename Cfg>
CUTE_DEVICE void tprod_bwd_accumulate(Cfg& cfg, auto& rYgrad_tprod_frg, auto& rXi_tprod_frg, auto& rY_grad_frg_buff_i) {
    auto rY_tprod_frg_buff = make_tensor<T>(cfg.Z.Frg); // buffer so we can operate on rY_tprod_frg without losing the original
    copy(rYgrad_tprod_frg, rY_tprod_frg_buff);
    tprod_bcast_multiply_all_except<Is>(rY_tprod_frg_buff, rXi_tprod_frg);
    add_tensor(rY_grad_frg_buff_i, rY_tprod_frg_buff);
}

// ---- Compute and accumulate Y_grad_for_i_warp_reduced for a factor ----
template<auto Is, typename T, typename Cfg>
CUTE_DEVICE void tprod_bwd_accumulate_warp_reduce(Cfg& cfg, auto& rYgrad_tprod_frg, auto& rXi_tprod_frg, auto& rB_buffer_for_i) {
    // rB_buffer_for_i contains the accumulated gradients for each factor, but is not fully reduced.
    using YType = TensorType(rYgrad_tprod_frg);
    using YTileShape = decltype(cfg.Z.TileShape);
    using YTprodFrgThr = decltype(cfg.Z.tprod_FrgThr);
    using XFrgThr = decltype(get<Is>(cfg.Xi).tile_copy.FrgThr);
    auto rY_tprod_frg_buff = make_tensor<T>(cfg.Z.Frg); // buffer so we can operate on rY_tprod_frg without losing the original
    copy(rYgrad_tprod_frg, rY_tprod_frg_buff);
    tprod_bcast_multiply_all_except<Is>(rY_tprod_frg_buff, rXi_tprod_frg);
    auto reduce_cfg = make_reduce_cfg<Is, YTileShape, YTprodFrgThr, XFrgThr, YType, cfg.reduce_buffer_size>();
    auto rB_frg = tprod_wrap_reduce<Is>(reduce_cfg, rY_tprod_frg_buff);
    add_tensor(rB_buffer_for_i, rB_frg);
}

// ---- Use Y_grad_for_i and compute X_i_grad for a factor ----
template<auto Is, typename Cfg, typename T>
CUTE_DEVICE void tprod_bwd(Cfg& cfg, T* smem, auto& rYgrad_tprod_frg_for_i, auto& Xigrad_copy_frg, int tid) {
    using YType = TensorType(rYgrad_tprod_frg_for_i);
    using YTileShape = decltype(cfg.Z.TileShape);
    using YTprodFrgThr = decltype(cfg.Z.tprod_FrgThr);
    auto rXigrad_copy_frg_buff = make_tensor<YType>(get<Is>(cfg.Xi).tile_copy.Frg);
    using XCopyFrgThr = decltype(get<Is>(cfg.Xi).tile_copy.FrgThr);
    auto reduce_cfg = make_reduce_cfg<Is, YTileShape, YTprodFrgThr, XCopyFrgThr, YType, cfg.reduce_buffer_size>();
    tprod_CTA_reduce_sum<Is>(reduce_cfg, rYgrad_tprod_frg_for_i, rXigrad_copy_frg_buff, smem);
    if (tid < size<1>(get<Is>(cfg.Xi).tile_copy.FrgThr))
        add_tensor(get<Is>(Xigrad_copy_frg), rXigrad_copy_frg_buff);
}

// ---- Use Y_grad_for_i_warp_reduced and compute X_i_grad for a factor ----
template<auto Is, typename Cfg, typename T>
CUTE_DEVICE void tprod_bwd_warp_reduced(Cfg& cfg, T* smem, auto& rBfrg_i, auto& Xigrad_copy_frg, int tid) {
    using YType = TensorType(rBfrg_i);
    using XCopyFrgThr = decltype(get<Is>(cfg.Xi).tile_copy.FrgThr);
    using YTileShape = decltype(cfg.Y.TileShape);
    using YTprodFrgThr = decltype(cfg.Y.tprod_FrgThr);
    auto rXigrad_copy_frg_buff = make_tensor<YType>(get<Is>(cfg.Xi).tile_copy.Frg);
    auto reduce_cfg = make_reduce_cfg<Is, YTileShape, YTprodFrgThr, XCopyFrgThr, YType, cfg.reduce_buffer_size>();
    tprod_smem_reduce_sum<Is>(reduce_cfg, rBfrg_i, rXigrad_copy_frg_buff, tid, smem);
    if (tid < size<1>(get<Is>(cfg.Xi).tile_copy.FrgThr))
        add_tensor(get<Is>(Xigrad_copy_frg), rXigrad_copy_frg_buff);
}

// ---- Use Y_grad and compute X_i_grad for all factors ----
template<auto... Is, typename Cfg>
__device__ void CTA_tprod_bwd(Cfg& cfg, char* smem, auto& rYgrad_tprod_frg, auto& rXi_tprod_frg, auto& Xigrad_copy_frg, int tid) {
    using T = TensorType(rYgrad_tprod_frg);
    using YTileShape = decltype(cfg.Z.TileShape); // [[d_tile, ...], b_tile]
    using YTprodFrgThr = decltype(cfg.Z.tprod_FrgThr);
    (..., [&]() { // accumulate sXigrad for each factor
        auto rY_tprod_frg_buff = make_tensor<T>(cfg.Z.Frg); // buffer so we can operate on rY_tprod_frg without losing the original
        copy(rYgrad_tprod_frg, rY_tprod_frg_buff);
        tprod_bcast_multiply_all_except<Is>(rY_tprod_frg_buff, rXi_tprod_frg);
        auto rXigrad_copy_frg_buff = make_tensor<T>(get<Is>(cfg.Xi).tile_copy.Frg);
        using XCopyFrgThr = decltype(get<Is>(cfg.Xi).tile_copy.FrgThr);
        auto reduce_cfg = make_reduce_cfg<Is, YTileShape, YTprodFrgThr, XCopyFrgThr, T, cfg.reduce_buffer_size>();
        tprod_CTA_reduce_sum<Is>(reduce_cfg, rY_tprod_frg_buff, rXigrad_copy_frg_buff, smem);
        if (tid < size<1>(get<Is>(cfg.Xi).tile_copy.FrgThr))
            add_tensor(get<Is>(Xigrad_copy_frg), rXigrad_copy_frg_buff);
    }());
}

// -------------- Example Kernels using TprodCfg --------------
template <bool duplicate_correction, auto... Is, typename Cfg, typename XPtr, typename ZGradPtr, typename XGradPtr>
__device__ void tiled_sympow_bwd_kernel_impl(index_sequence<Is...>, Cfg cfg, XPtr X_ptr, ZGradPtr Zgrad_ptr, XGradPtr Xgrad_ptr) {
    using T = typename Cfg::T;
    int tid = threadIdx.x; int bid = blockIdx.x;
    // ----- Iterators of the kernel -----
    auto sympow_tile_coords = SympowCoords<Cfg>{};
    sympow_tile_coords.step_b(bid);
    // ------ Global memory slabs ------
    auto gX_slab = make_tensor(make_gmem_ptr(X_ptr), cfg.X.gSlab);
    auto gZgrad_slab = make_tensor(make_gmem_ptr(Zgrad_ptr), cfg.Z.gSlab);
    auto gXgrad_slab = make_tensor(make_gmem_ptr(Xgrad_ptr), cfg.X.gSlab);
    auto gX_batch = sympow_tile_coords.slice_X_batch(gX_slab);
    // ------ Shared memory tensors ------
    __shared__ alignas(16) T Z_smem[int(cosize(cfg.Z.sTile))];
    __shared__ alignas(16) T X_smem[int(cosize(cfg.X.sBatch))];
    __shared__ alignas(16) T Xgrad_smem[int(cosize(cfg.X.sBatch))];
    __shared__ alignas(16) char X_reduce_buffer[cfg.reduce_buffer_size];
    auto sZgrad_tile = make_tensor(make_smem_ptr(Z_smem), cfg.Z.sTile);
    auto sX_batch = make_tensor(make_smem_ptr(X_smem), cfg.X.sBatch);
    auto sXgrad_batch = make_tensor(make_smem_ptr(Xgrad_smem), cfg.X.sBatch);
    CTA_copy_tile(cfg.X.batch_copy, gX_batch, sX_batch);
    clear(sXgrad_batch); // clear the gradient accumulators
    cp_async_fence(); cp_async_wait<0>(); __syncthreads();
    auto sZgrad_copy_frg = slice_rest(sZgrad_tile, cfg.Z.tile_copy.FrgThr, tid);
    auto rYgrad_tprod_frg = make_tensor<T>(cfg.Z.Frg);
    while (sympow_tile_coords.valid_D_tile()) {
        auto gZgrad_tile = sympow_tile_coords.slice_Z_tile(gZgrad_slab);
        CTA_copy_tile(cfg.Z.tile_copy, gZgrad_tile, sZgrad_tile);
        cp_async_fence(); cp_async_wait<0>(); __syncthreads(); // ensure the sZgrad_tile is ready
        // load the Zgrad_frg into registers (applying the duplicate correction if necessary)
        copy(slice_rest(sZgrad_tile, cfg.Z.tprod_FrgThr, tid), rYgrad_tprod_frg);
        if constexpr (duplicate_correction)
            tensor_scalar_prod(rYgrad_tprod_frg, static_cast<T>(sympow_tile_coords.scale_correction()));
        // load the relevant Xi_frg into registers
        auto rXi_tprod_frg = make_tuple(make_tensor<T>(get<Is>(cfg.Xi).Frg)...);
        (..., [&]() { // s2r load of Xi_tprod_frg. Could be optimized by only loading the Xi fragment when the index has changed
            auto sXi_tile = sympow_tile_coords.slice_X_tile_from_batch<Is>(sX_batch);
            auto sXi_tprod_frg = slice_rest(sXi_tile, get<Is>(cfg.Xi).tprod_FrgThr, tid);
            copy(sXi_tprod_frg, get<Is>(rXi_tprod_frg));
        }());
        // slice the Xgrad fragments for accumulation
        auto sXigrad_copy_frg = make_tuple(slice_rest(sympow_tile_coords.slice_X_tile_from_batch<Is>(sXgrad_batch),
                                                      get<Is>(cfg.Xi).tile_copy.FrgThr,
                                                      tid)...);
        CTA_tprod_bwd<Is...>(cfg, X_reduce_buffer, rYgrad_tprod_frg, rXi_tprod_frg, sXigrad_copy_frg, tid);
        __syncthreads(); // sZgrad_tile cannot be overwritten until all threads have used it to compute their gradients 
        sympow_tile_coords.step_D();
    }
    auto gXgrad_batch = sympow_tile_coords.slice_X_batch(gXgrad_slab);
    CTA_copy_tile(cfg.X.batch_copy, sXgrad_batch, gXgrad_batch);
}
template <bool duplicate_correction, typename Cfg, typename XPtr, typename ZGradPtr, typename XGradPtr>
__global__ void tiled_sympow_bwd_kernel(Cfg cfg, XPtr X_ptr, ZGradPtr Zgrad_ptr, XGradPtr Xgrad_ptr) {
    tiled_sympow_bwd_kernel_impl<duplicate_correction>(make_index_sequence<cfg.p>{}, cfg, X_ptr, Zgrad_ptr, Xgrad_ptr);
}
template<bool duplicate_correction, typename Cfg, typename XPtr, typename ZGradPtr, typename XGradPtr>
int launch_tiled_sympow_bwd_kernel(Cfg cfg, XPtr X_ptr, ZGradPtr Zgrad_ptr, XGradPtr Xgrad_ptr) {
    int blocks = cfg.b_tile_num;
    int threads = cfg.thread_num;
    tiled_sympow_bwd_kernel<duplicate_correction><<<blocks, threads>>>(cfg, X_ptr, Zgrad_ptr, Xgrad_ptr);
    CUDA_CHECK_LAST_ERROR("CUDA kernel launch error");
    return 0;
}

} // namespace vidrial