#pragma once
#include "../sympow/sympow_cfg.cuh"
#include "tprod_bwd.cuh"

namespace vidrial {

// -------------- 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));
        }());
        (..., [&]() {
            auto sXigrad_tile = sympow_tile_coords.slice_X_tile_from_batch<Is>(sXgrad_batch);
            CTA_tprod_bwd(get<Is>(cfg.tprod_bwd), rXi_tprod_frg, rYgrad_tprod_frg, sXigrad_tile, X_reduce_buffer);
        }());
        __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