#pragma once
#include "../sympow_mma/sympow_mma_cfg.cuh"
#include "../sympow_bwd/kernel.cuh"
#include "../sympow_mma/kernel.cuh"
#include "../../cutlass/include/cute/tensor.hpp"
#include "../../cutlass/include/cute/tensor_impl.hpp"
#include "../../cuda_utils/perf_cfg.cuh"
#include "../../cuda_utils/allocator.cuh"
#include "../../cuda_utils/pipeline.cuh"

namespace vidrial {

/* Cfg where the expanded object is C along the M dimension
X = c
*/
template<typename _T, int _pow, typename MmaAtom, typename MNKAtomPlacement, typename _MNKPSlabShape, typename MNKTileShape,
         int _d, int _d_tile, typename _GASlab, typename _GBSlab, typename _GcSlab, int _reduce_buffer_size, typename _PerfCfg = DefaultPerfCfg>
struct SympowCMmaCfg {
    using T = _T;
    using acc_T = float;
    static constexpr int pow = _pow;
    static constexpr int d = _d;
    static constexpr int d_tile = _d_tile;
    using PerfCfg = _PerfCfg;
    static constexpr PerfCfg perf{};
    static constexpr int reduce_buffer_size = _reduce_buffer_size;
    using GASlab = decltype(static_tree_cast<int64_t>(_GASlab{}));
    using GBSlab = decltype(static_tree_cast<int64_t>(_GBSlab{}));
    using GcSlab = decltype(static_tree_cast<int64_t>(_GcSlab{}));
    using MNKAtomLayout = Layout<Shape<_1,_1,_1>>; // generic layouts not implemented
    using MNKPSlabShape = decltype(static_tree_cast<int64_t>(_MNKPSlabShape{}));
    static constexpr long M=size<0>(MNKPSlabShape{}),N=size<1>(MNKPSlabShape{}),K=size<2>(MNKPSlabShape{}),P=size<3>(MNKPSlabShape{}); 
    static constexpr long M_tile=get<0>(MNKTileShape{}),N_tile=get<1>(MNKTileShape{}),K_tile=get<2>(MNKTileShape{}); 
    static constexpr int D = get<0>(MNKPSlabShape{});
    static constexpr int D_tile = static_pow<pow>(d_tile);
    static_assert(D == sympow_dim<pow, d, d_tile>());
    static_assert(D_tile == get<0>(MNKTileShape{}), "D_tile mismatch with M_tile");
    static constexpr int D_tile_num = sympow_dim<pow, d/d_tile>();
    using CSlabShape = decltype(ABC_get_MNKP(C_t{}, MNKPSlabShape{}));
    using CTileShape = decltype(ABC_get_MNK(C_t{}, MNKTileShape{}));
    using cSlabShape = Shape<Int<d>,Int<N>,Int<P>>; // similar to XSlabShape = [d,[N,P]]
    using XSlabShape = Shape<Int<d>,Shape<Int<N>,Int<P>>>;
    using cTileShape = Shape<Int<d_tile>,Int<N_tile>>; // same as XTileShape
    using XTileShape = cTileShape; // same as XTileShape
    using c2XSlab = Layout<XSlabShape>; // [d,N,P] -> [d,[N,P]]
    using X2cSlab = decltype(group<1,3>(Layout<cSlabShape>{})); // [d,[N,P]] -> [d,N,P]
    using c2XTile = Layout<XTileShape>; // [d_tile,N_tile] -> [d_tile,N_tile]
    using C2ZTile = Layout<CTileShape>; // [D_tile,N_tile] -> [d_tile^p,N_tile]
    using ZSlabShape = decltype(static_tree_cast<int64_t>(Shape<Int<D>,Shape<Int<N>,Int<P>>>{})); // [D,[N,P]]
    using C2ZSlab = decltype(flatten(Layout<ZSlabShape>{})); // [D,N,P] -> [D,[N,P]]
    using ZTileShape = decltype(tpow_shape<pow>(XTileShape{}));
    using GXSlab = decltype(GcSlab{}.compose(X2cSlab{})); // gX_slab is a view of gc_slab
    using GZSlab = decltype(make_layout(static_tree_cast<int64_t>(sympow_shape<pow,d_tile>(XSlabShape{})))); // virtual
    using GCSlab = decltype(GZSlab{}.compose(C2ZSlab{})); // virtual
    GXSlab gXSlab;
    static_assert(size(GZSlab{}) == size(GCSlab{}));
    using MmaCfg = decltype(make_mma_cfg<T, PerfCfg, MmaAtom, MNKAtomPlacement>(MNKTileShape{}, MmaAtom{}, MNKAtomPlacement{}, GASlab{}, GBSlab{}, GCSlab{}));
    static constexpr MmaCfg mma{};
    MmaAtom mma_Atom{};
    decltype(mma.A) A;
    decltype(mma.B) B;
    decltype(mma.C) C;
    using FrgTypeA = typename MmaCfg::FrgTypeA;
    using FrgTypeB = typename MmaCfg::FrgTypeB;
    using FrgTypeC = typename MmaCfg::FrgTypeC;
    static constexpr int thread_num = mma.thread_num;
    static_assert(size(ZTileShape{}) == size(decltype(C){}.tileShape));
    using _ZMmaFrgThr = decltype(C2ZTile{}.compose(C.mma_FrgThr));
    using _ZTprodFrgThr = decltype(colayout(ZTileShape{}, get<0>(_ZMmaFrgThr{}))); // colayout transforms the mma_frg layout into a format compatible with tprod
    using ZTprodFrgThr = decltype(make_layout(_ZTprodFrgThr{}, get<1>(_ZMmaFrgThr{})));
    static_assert(size(ZTprodFrgThr{}) == size(ZTileShape{}));
    using CTprodFrgThr_Frg = decltype(left_inverse(C2ZTile{}).compose(get<0>(ZTprodFrgThr{}))); // maps tprod_frg_coords -> C_coords
    using ZTprod__2__CMma__frg = decltype(left_inverse(get<0>(C.mma_FrgThr)).compose(CTprodFrgThr_Frg{})); // C_tprod_frg_coords -> Y_tprod_frg_coords
    ZTprod__2__CMma__frg Z_tprod_frg__2__C_mma_frg;
    using SZTile = Layout<ZTileShape>; // virtual (we don't even store the expansions in smem)
    using Sympow = SympowCfg<T, pow, XSlabShape, XTileShape, ZTprodFrgThr, GZSlab, GXSlab, SZTile>;
    Sympow sympow;
    using SympowCoords = typename Sympow::SympowCoords;
    using Xi_t = decltype(typename Sympow::Xi_t{});
    Xi_t Xi;
    Sympow::XBatching X;
    Sympow::Z_t Z;
    using Xi_size_t = decltype(transform(Xi_t{}, [](auto const& x) { return size(x.sTile); }));
    using Xi_offset_t = decltype(get<1>(fold(Xi_size_t{}, make_tuple(_0{},make_tuple()),
            [](auto const& carry, auto const& Xi_size) {
                    auto [current, offsets] = carry;
                    return make_tuple(current+Xi_size, append(offsets, current)); })));
    using X_size_t = decltype(fold(Xi_size_t{}, _0{}, [](auto const& a, auto const& b) { return a + b; }));
    static constexpr Xi_size_t Xi_smem_size{};
    static constexpr Xi_offset_t Xi_smem_offset{};
    static constexpr X_size_t X_smem_size{};
    static int smem_size() {
        using X_SBatch = decltype(X)::SXBatch;
        using A_STile = decltype(A)::STile;
        using B_STile = decltype(B)::STile;
        constexpr static int smempipe = static_min(perf.smempipe, mma.K_tile_num * mma.M_tile_num);
        constexpr static int smem_size = Allocator<16>::total<T, T, T, T, acc_T>(
            int(cosize(X_SBatch{})), // X_smem
            int(cosize(X_SBatch{})), // Xgrd_smem
            int(size(A_STile{}) * smempipe), // A_smem
            int(size(B_STile{}) * smempipe), // B_smem
            reduce_buffer_size // X_reduce_buffer
        );
        return smem_size;
    }
};

template<auto... Is, typename Maker>
CUTE_DEVICE auto make_variadic_tuple(int_sequence<Is...>, Maker maker) {
    return make_tuple(maker(Is)...);
}

/* ----- fused sympow mma kernel -----
 * This kernel is a function A,B,c -> cdot = sympow_bwd(c, A@B)
 *     (used in the implementation of sympowA_mma_bwd)
 * Kernel uses 2 different coordinate systems for the 2 operations
 *   dims M,N,K,P
 *   A_slab = [M,K,P]  A_tile = [M_tile,K_tile]
 *   B_slab = [N,K,P]  B_tile = [N_tile,K_tile]
 *   C_slab = [M,N,P]  C_tile = [M_tile,N_tile]
 *   dims d,b: using d=M and b=[K,P]
 *   X_slab = [d,b]    X_tile = [d_tile,b_tile]
 *   Z_slab = [D,b]  Z_tile = [D_tile,b_tile]
 *   where D_tile = [d_tile,...]  (repeated power times)
 *   and D = [D_tile,d_tile_num]
 * The two systems are connected via
 *   D = M   and  b = [N,P] and b_tile = [N_tile,1]
 *   Z_slab = [D,[N,P]], X_slab = [d,[N,P]]
 *   C_tile = Z_tile = [D_tile,N_tile], c_tile = X_tile = [d_tile,N_tile]
 * Each CTA in the grid computes an Xgrad_batch of shape [d, b_tile]
 *   CTA_idx = [N_tile_idx, P_tile_idx] (constant through execution)
 *   so that each CTA computes a different c_tile
 * The mainloop has 2 nested loops:
 *   D_tile_idx  ranges  0 -> cfg.D_tile_num
 *       K_tile_idx  ranges  0 -> cfg.K_tile_num
*/ 
template <bool duplicate_correction, bool scale_A, int smempipe_, auto... Is, typename Cfg, typename AT, typename BT, typename CT>
__device__ void mma_sympow_bwd_kernel_impl(int_sequence<Is...>, Cfg cfg, AT* A_ptr, BT* B_ptr, CT* c_ptr, CT* cgrad_ptr, float sA) {
    int tid = threadIdx.x;
    int N_tile_idx = blockIdx.y;
    int P_tile_idx = blockIdx.z;
    // ----- Iterators of the kernel -----
    auto mma_tile_coords = MmaMNKCoords(cfg.mma.MNK_tile_shape);
    mma_tile_coords.step_N(N_tile_idx);
    mma_tile_coords.step_P(P_tile_idx);
    auto sympow_tile_coords = SympowCoords<decltype(cfg.sympow)>{};
    sympow_tile_coords.step_b(N_tile_idx + cfg.mma.N_tile_num * P_tile_idx); // b_tile_idx = [N_tile_idx, P_tile_idx]
    // ------ Global memory slabs ------
    auto gA_slab = make_tensor(make_gmem_ptr(A_ptr), cfg.A.gSlab);
    auto gB_slab = make_tensor(make_gmem_ptr(B_ptr), cfg.B.gSlab);
    auto gX_slab = make_tensor(make_gmem_ptr(c_ptr), cfg.X.gSlab); // gc and gX are equivalent
    auto gXgrad_slab = make_tensor(make_gmem_ptr(cgrad_ptr), cfg.gXSlab);
    auto gX_batch = sympow_tile_coords.slice_X_batch(gX_slab);
    // ------ Shared memory tensors ------
    constexpr static int smempipe = static_min(cfg.perf.smempipe, cfg.mma.K_tile_num * cfg.mma.M_tile_num);
    auto pipe = SmemPipe<smempipe>();
    extern __shared__ char smem[];
    Allocator<16> alloc(smem);
    CT* X_smem = alloc.allocate<CT>(int(cosize(cfg.X.sBatch))); // [d, b_tile]
    CT* Xgrd_smem = alloc.allocate<CT>(int(cosize(cfg.X.sBatch))); // [d, b_tile]
    AT* A_smem = alloc.allocate<AT>(size(cfg.A.sTile) * smempipe);
    BT* B_smem = alloc.allocate<BT>(size(cfg.B.sTile) * smempipe);
    auto sA_tile_pipe = pipe.create(A_smem, cfg.A.sTile, cfg.A.tile_copy);
    auto sB_tile_pipe = pipe.create(B_smem, cfg.B.sTile, cfg.B.tile_copy);
    auto sX_batch = make_tensor(make_smem_ptr(X_smem), cfg.X.sBatch); // [d, b_tile]
    auto sXgrad_batch = make_tensor(make_smem_ptr(Xgrd_smem), cfg.X.sBatch); // [d, b_tile]
    using C_Frg_layout = decltype(cfg.C.make_mma_frg().layout());
    using YFrgType = TensorType(declval<decltype(cfg.C.make_mma_frg())>());
    using Ygrad_Frg_layout = decltype(C_Frg_layout{}.compose(cfg.Z_tprod_frg__2__C_mma_frg));
    auto create_rYgrad_tprod_frg_cache = [&](size_t i) {
        auto rYgrad_tprod_frg = make_tensor<YFrgType>(cfg.Z.Frg);
        clear(rYgrad_tprod_frg);
        return rYgrad_tprod_frg;
    };
    auto rYgrad_tprod_buffers = make_variadic_tuple(make_int_sequence<sizeof...(Is) - 1>{}, create_rYgrad_tprod_frg_cache);
    CTA_copy_tile(cfg.X.batch_copy, gX_batch, sX_batch);
    clear(slice_rest(sXgrad_batch, cfg.X.batch_copy.FrgThr, tid)); // Every thread clears a fragment of the Xgrad_batch
    cp_async_fence(); cp_async_wait<0>(); __syncthreads();
    // TODO: properly compute the size of buffer needed
    char* X_reduce_buffer = alloc.allocate<char>(cfg.reduce_buffer_size);
    // ------ Pipeline fetch A_tile and B_tile ------
    auto pipe_fetch = [&]() {
        if (mma_tile_coords.valid_K_tile(cfg.K) && mma_tile_coords.valid_M_tile(cfg.M)) {
            auto gA_tile = mma_tile_coords.slice_A_tile(gA_slab);
            auto gB_tile = mma_tile_coords.slice_B_tile(gB_slab);
            pipe.fetch(gA_tile, sA_tile_pipe, cfg.A.tile_copy);
            pipe.fetch(gB_tile, sB_tile_pipe, cfg.B.tile_copy);
            mma_tile_coords.step_K();
            if (mma_tile_coords.K_coord() == cfg.mma.K_tile_num) {
                mma_tile_coords.reset_K();
                mma_tile_coords.step_M();
            }
        }
        pipe.commit();
    };
    // ------ Prefill pipeline ------
    for (; pipe.stage < smempipe - 1; pipe.step())
        pipe_fetch();
    // ------ Main loop ------
    while (sympow_tile_coords.valid_D_tile()) {
        // ------ Register fragments for mma and tprod operations ------
        auto rA_frg_mma = cfg.A.make_mma_frg();
        auto rB_frg_mma = cfg.B.make_mma_frg();
        auto rC_frg_mma = cfg.C.make_mma_frg();
        clear(rC_frg_mma);
        for (int k_tile = 0; k_tile < cfg.mma.K_tile_num; k_tile++) {
            pipe_fetch();
            pipe.ready();
            auto sA_tile = pipe.read(sA_tile_pipe);
            auto sB_tile = pipe.read(sB_tile_pipe);
            vidrial::pipeAB_mma<scale_A>(cfg, sA_tile, sB_tile, rC_frg_mma, sA);
            pipe.step();
        }
        auto rYgrad_tprod_frg = rC_frg_mma.compose(cfg.Z_tprod_frg__2__C_mma_frg);
        if constexpr (duplicate_correction)
            tensor_scalar_prod(rYgrad_tprod_frg, static_cast<YFrgType>(sympow_tile_coords.scale_correction()));
        // We are ready to perform an accumulation step onto Xigrad using the Ygrad_frg
        auto rXi_tprod_frg = make_tuple(make_tensor<YFrgType>(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_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)...);
        (..., [&]() {
            if constexpr (Is != (sizeof...(Is) - 1)) {
                tprod_bwd_accumulate<Is, YFrgType>(cfg, rYgrad_tprod_frg, rXi_tprod_frg, get<Is>(rYgrad_tprod_buffers));
                if (sympow_tile_coords.template will_change<Is>() || sympow_tile_coords.is_last()) {
                    tprod_bwd<Is>(cfg, X_reduce_buffer, get<Is>(rYgrad_tprod_buffers), sXigrad_copy_frg, tid);
                    clear(get<Is>(rYgrad_tprod_buffers));
                }
            } else {
                CTA_tprod_bwd<Is>(cfg, X_reduce_buffer, rYgrad_tprod_frg, rXi_tprod_frg, sXigrad_copy_frg, tid);
            }
        }());
        sympow_tile_coords.step_D();
        __syncthreads(); // sXigrad_batch is being accumulated. Wait untill all threads are done.
    }
    // ------ Write X_tile to global memory ------
    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, bool scale_A, int smempipe, typename Cfg, typename AT, typename BT, typename CT>
__global__ void mma_sympow_bwd_kernel(Cfg cfg, AT* a_ptr, BT* B_ptr, CT* c_ptr, CT* Cgrad_ptr, float sA) {
    mma_sympow_bwd_kernel_impl<duplicate_correction, scale_A, smempipe>(make_int_sequence<cfg.pow>{}, cfg, a_ptr, B_ptr, c_ptr, Cgrad_ptr, sA);
}

template<bool duplicate_correction=true, bool scale_A=false, int smempipe=1, typename Cfg, typename AT, typename BT, typename CT>
int launch_mma_sympow_bwd_kernel(Cfg cfg, AT* a_ptr, BT* B_ptr, CT* c_ptr, CT* cgrad_ptr, float sA) {
    dim3 blocks(1, cfg.mma.N_tile_num, cfg.mma.P);
    int threads = cfg.thread_num;
    int smem_size = cfg.smem_size(); // Use calculated instead of hardcoded
    
    auto kernel = mma_sympow_bwd_kernel<duplicate_correction, scale_A, smempipe, Cfg, AT, BT, CT>;
    CUDA_CHECK_RETURN(adjust_dynamic_smem_size(kernel, smem_size), "Error setting max dynamic shared memory size");
    kernel<<<blocks,threads,smem_size>>>(cfg, a_ptr, B_ptr, c_ptr, cgrad_ptr, sA);
    CUDA_CHECK_LAST_ERROR("CUDA kernel launch error");
    return 0;
}

} // namespace vidrial