#pragma once
#include "../../cuda_utils/utilities.cuh"
#include "../../cuda_utils/frgthr_tiling.cuh"

using namespace cute;

namespace vidrial {

/*----------------- Thread Reduce -----------------
 * X_Shape: [reduce, batch]
 * x_Shape: [outerReduce, batch]
 * 
 * At the CTA level, calling thread_reduce transforms:
 *                      X -> x
 * At the thread level, rX_frg -> rx_frg
 * 
 * Depending on the X_FrgThr, which maps [X_frg_idx, thr_idx] -> X_idx
 * The fragment X_frg is split into [innerReduce, innerBatch]
 *   innerReduceBatch_2_Frg: [innerReduce, innerBatch] -> X_Frg
 * Accordingly, reduce is split into [innerReduce, outerReduce]
 *      ReduceSplit: [innerReduce_idx, warp_reduce_idx] -> reduce_idx
 * CanonicalFrg: [innerReduce_idx, innerBatch_idx] -> X_frg_idx
 */
template<typename _XShape, typename _XFrgThr, typename _xShape, typename _xFrgThr, typename _CanonicalFrg>
struct ThreadReduceCfg {
    using XShape = _XShape;
    using XFrgThr = _XFrgThr;
    using xShape = _xShape;
    using xFrgThr = _xFrgThr;
    using CanonicalFrg = _CanonicalFrg;
    using xFrgShape = decltype(get<0>(xFrgThr{}).shape());
    static constexpr int local_reduce = size<0>(CanonicalFrg{});
    static constexpr int local_batch = size<1>(CanonicalFrg{});
    XShape X_Shape;
    XFrgThr X_FrgThr;
    xShape x_Shape;
    xFrgThr x_FrgThr;
    CanonicalFrg canonical_frg;
    xFrgShape x_Frg_Shape;
};
template<typename XShape, typename XFrgThr>
auto make_ThreadReduceCfg() {
    auto X_Frg = get<0>(XFrgThr{});
    auto X_Thr = get<1>(XFrgThr{});
    auto X2Red = projection_layout<0>(XShape{});
    auto X2Bat = projection_layout<1>(XShape{});
    auto innerBat2Frg = nullspace(X2Red.compose(X_Frg));
    auto innerRed2Frg = nullspace(X2Bat.compose(X_Frg));
    auto innerBat2X = X_Frg.compose(innerBat2Frg);
    auto innerRed2X = X_Frg.compose(innerRed2Frg);
    auto outerRed2X = complement(innerRed2X, size<0>(XShape{}));
    auto x2X = make_layout(XShape{}).compose(outerRed2X, _);
    auto x_FrgThr = left_inverse(x2X).compose(make_layout(innerBat2X, X_Thr));
    auto canonicalFrg = make_layout(innerRed2Frg, innerBat2Frg);
    return ThreadReduceCfg<XShape, XFrgThr, decltype(x2X.shape()), decltype(x_FrgThr), decltype(canonicalFrg)>{};
}
CUTE_DEVICE auto thread_reduce(auto& cfg, auto& fn, auto& rX_frg) {
    using T = TensorType(rX_frg);
    auto rx_frg = make_tensor<T>(cfg.x_Frg_Shape);
    auto rX_cfrg = rX_frg.compose(cfg.canonical_frg);
    CUTE_UNROLL
    for (int innerBatch_idx = 0; innerBatch_idx < cfg.local_batch; innerBatch_idx++) {
        T carry = rX_cfrg(_0{}, innerBatch_idx);
        CUTE_UNROLL
        for (int innerReduce_idx = 1; innerReduce_idx < cfg.local_reduce; innerReduce_idx++)
            carry = fn(carry, rX_cfrg(innerReduce_idx, innerBatch_idx));
        rx_frg(innerBatch_idx) = carry;
    }
    return rx_frg;
}

/*----------------- Warp Reduce -----------------
 * Accepts X_frg that are irreducible at the thread level and reduces X->x
 * as much as possible using warp_level shufl instructions.
 * Note that X_frg = x_frg.
*/
template<typename _XShape, typename _XFrgThr, typename _xShape, typename _xFrgThr, typename _XReduceSteps>
struct WarpReduceCfg {
    using XShape = _XShape;
    using XFrgThr = _XFrgThr;
    using xShape = _xShape;
    using xFrgThr = _xFrgThr;
    using XReduceSteps = _XReduceSteps; // e.g. ReduceSteps = 4:2, then the threads 0,2,4,6 of each warp will perform a reduce/broadcast
    XShape X_Shape;
    XFrgThr X_FrgThr;
    xShape x_Shape;
    xFrgThr x_FrgThr;
    XReduceSteps X_ReduceSteps;
};
template<typename XShape, typename XFrgThr>
auto make_WarpReduceCfg() {
    constexpr auto X_Frg = get<0>(XFrgThr{});
    constexpr auto X_Thr = get<1>(XFrgThr{});
    constexpr auto warp_size = static_min(32, size(X_Thr));
    auto ThreadPartition = zipped_divide(get<1>(XFrgThr{}), Int<warp_size>{});
    auto X_Shard = make_layout(get<0>(XFrgThr{}), get<0>(ThreadPartition));
    auto X_Warp = get<1>(ThreadPartition);
    auto ShardWarp = make_layout(X_Shard, X_Warp);
    auto X2Red = projection_layout<0>(XShape{});
    auto X2Bat = projection_layout<1>(XShape{});
    auto frgBat2Frg = nullspace(X2Red.compose(X_Frg));
    auto innerBat2Shard = nullspace(X2Red.compose(X_Shard));
    auto innerRed2Shard = nullspace(X2Bat.compose(X_Shard));
    auto frgBat2X = X_Frg.compose(frgBat2Frg);
    auto innerRed2X = X_Shard.compose(innerRed2Shard);
    auto outerRed2X = complement(innerRed2X, size<0>(XShape{}));
    auto x2X = make_layout(XShape{}).compose(outerRed2X, _);
    auto y2X = complement(x2X, size(XShape{}));
    auto xy2X = make_layout(x2X, y2X);
    auto X2xy = right_inverse(xy2X);
    auto xy2x = projection_layout<0>(xy2X.shape());
    auto X2x = xy2x.compose(X2xy);
    auto x_Thr = X2x.compose(X_Thr);
    auto x_Frg = left_inverse(x2X).compose(frgBat2X);
    auto x_FrgThr = make_layout(x_Frg, x_Thr);
    auto innerRed2Thr_ = left_inverse(X_Thr).compose(innerRed2X);
    auto innerRed2Thr = coalesce(innerRed2Thr_);
    return WarpReduceCfg<XShape, XFrgThr, decltype(x2X.shape()), decltype(x_FrgThr), decltype(innerRed2Thr)>{};
}
template<typename ReduceSteps>
__device__ void warp_reduce_value(auto& fn, auto& value) {
    static_assert(depth(ReduceSteps{})<=1);
    constexpr int reduce_size = get<0>(ReduceSteps{}.shape());
    constexpr int reduce_stride = get<0>(ReduceSteps{}.stride());
    static_assert(1<<static_log2<reduce_stride>()==reduce_stride, "reduce_stride must be a power of 2");
    static_assert(1<<static_log2<reduce_size>()==reduce_size, "reduce_size must be a power of 2");
    CUTE_UNROLL
    for (int lane_mask = reduce_stride; lane_mask < reduce_size*reduce_stride; lane_mask *= 2) {
        auto tmp = __shfl_xor_sync(0xffffffff, value, lane_mask, warpSize);
        value = fn(value, tmp);
    }
    if constexpr (rank(ReduceSteps{}) > 1)
        warp_reduce_value<decltype(drop<0>(ReduceSteps{}))>(fn, value);
}
__device__ void warp_reduce(auto& cfg, auto&& fn, auto& rX_frg) {
    // it is assumed that the fragment is irreducible. All innerReduce elements
    // belong to different threads in the same warp
    if constexpr (size(decltype(cfg.X_ReduceSteps){})>1) {
        CUTE_UNROLL
        for (int i=0; i<size(rX_frg); i++) {
            warp_reduce_value<decltype(cfg.X_ReduceSteps)>(fn, rX_frg(i));
        }
    }
}


/*----------------- SMEM Reduce ----------------- */
/* Reduces a shared memory tensor sX: [reduce, batch]
each thread takes care of an entire reduce column during each iteration
if thread_num < batch multiple iterations will be performed cover the entire batch dimension
*/
template<typename T, int thread_num, typename Fn>
CUTE_DEVICE void smem_reduce_batch_parallel(Fn& fn, auto sX) {
    CUTE_UNROLL
    for (int iter = 0; iter * thread_num < size<1>(sX); iter++) {
        int b = iter * thread_num + threadIdx.x; // batch index
        T carry;
        if (b < size<1>(sX)) { // b index within bounds
            carry = sX(_0{},b);
            CUTE_UNROLL
            for (int r = 1; r < size<0>(sX); r++) {
                carry = fn(carry,sX(r,b));
            }
        }
        __syncthreads(); // wait for all reductions to finish before writing to smem
        if (b < size<1>(sX))
            sX(_0{},b) = carry;
    }
}

/*----------------- CTA SMEM Reduce -----------------
The algorithm proceeds as follows:
- X: [reduce, batch] is broken into X_tiles: [reduce, batch_tile] 
- for each tile:
    - threads collaborate to store X_tile in smem
    - the sX_tile is reduced
    - results of the reduction are loaded into rx_frg
*/
template<typename _T, typename _XShape, typename _XFrgThr, typename _xFrgThr, int _smem_buffer_size>
struct CTASmemReduceCfg {
    using T = _T;
    using XShape = _XShape;
    using XFrgThr = _XFrgThr;
    using xFrgThr = _xFrgThr;
    static constexpr int smem_buffer_size = _smem_buffer_size;
    static constexpr int thread_num = size<1>(XFrgThr{});
    static constexpr int reduce = size<0>(XShape{});
    static constexpr int batch = size<1>(XShape{});
    static constexpr int bytes_per_reduce = reduce * sizeof(T);
    static constexpr int batch_tile = static_min(smem_buffer_size / bytes_per_reduce, batch);
    static constexpr int tile_num = batch / batch_tile;
    static_assert(thread_num >= size<1>(XFrgThr{}) && thread_num >= size<1>(xFrgThr{}));
    static_assert(batch_tile > 0, "Not enough smem to reduce a single dimension of the tile. Try a smaller tile size along the reduce dim.");
    static_assert(is_static<XShape>::value && is_static<XFrgThr>::value && is_static<xFrgThr>::value);
    static_assert(batch_tile * bytes_per_reduce <= smem_buffer_size);
    using xShape = Shape<Int<batch>>;
    using XTileShape = Shape<Int<reduce>, Int<batch_tile>>;
    using xTileShape = Shape<Int<batch_tile>>;
    using XTileRest = decltype(zipped_divide(Layout<XShape>{}, XTileShape{}));
    using xTileRest = decltype(zipped_divide(Layout<xShape>{}, xTileShape{}));
    using XFrgThrTilingCfg = decltype(make_FrgThrTilingCfg(XTileRest{}, XFrgThr{}));
    using xFrgThrTilingCfg = decltype(make_FrgThrTilingCfg(xTileRest{}, xFrgThr{}));
    XFrgThr X_FrgThr;
    xFrgThr x_FrgThr;
    decltype(make_layout(get<0>(XFrgThr{}.shape()))) X_Frg;
    decltype(make_layout(get<0>(xFrgThr{}.shape()))) x_Frg;
    XTileShape X_TileShape;
    XFrgThrTilingCfg X_frgthr_tiling;
    xFrgThrTilingCfg x_frgthr_tiling;
};
template<bool carry_initial, typename Cfg, typename Fn>
CUTE_DEVICE void CTA_smem_reduce(Cfg& cfg, auto& rX_frg, auto& rx_frg, Fn& fn, char* smem_buff) {
    using T = typename Cfg::T;
    auto sX_tile = make_tensor(make_smem_ptr(reinterpret_cast<T*>(smem_buff)), make_layout(cfg.X_TileShape));
    for (int tile_idx = 0; tile_idx < cfg.tile_num; tile_idx++) {
        FrgThr_store_tile(cfg.X_frgthr_tiling, rX_frg, sX_tile, tile_idx);
        __syncthreads();
        smem_reduce_batch_parallel<T, cfg.thread_num>(fn, sX_tile);
        __syncthreads();
        auto sx_tile = sX_tile(_0{},_); // the reduced values can be read from the 0th index
        FrgThr_load_tile(cfg.x_frgthr_tiling, rx_frg, sx_tile, tile_idx);
        __syncthreads();
    }
}

/* ----------------- CTA Reduce -----------------
The X data is held by thre CTA_threads according to X_FrgThr
The CTA_reduce algorithm consists of 3 steps:
    1 - each thread reduces it's own fragment as much as possible
    2 - each warp reduces using shufl instructions
    3 - smem reduction across warps finishes the job
*/ 
template<typename _T, typename _XShape, typename _XFrgThr, typename _xFrgThr, int _smem_buffer_size, int... Is>
struct CTAReduceCfg {
    using T = _T;
    using XShape = _XShape;
    using XFrgThr = _XFrgThr;
    using xFrgThr = _xFrgThr;
    XFrgThr X_FrgThr;
    xFrgThr x_FrgThr;
    static constexpr int smem_buffer_size = _smem_buffer_size;
    // Conversion to canonical form: A = [reduce, batch]
    using AShape = decltype(make_shape(get<Is...>(XShape{}), drop<Is...>(XShape{})));
    using XLayout = decltype(make_layout(XShape{}));
    using A2X = decltype(make_layout(get<Is...>(XLayout{}), drop<Is...>(XLayout{})));
    using AFrgThr = decltype(left_inverse(A2X{}).compose(XFrgThr{}));
    // Thread level reduction]
    using ThreadReduceCfg = decltype(make_ThreadReduceCfg<AShape, AFrgThr>());
    using BShape = typename ThreadReduceCfg::xShape;
    using BFrgThr = typename ThreadReduceCfg::xFrgThr;
    // Warp reduction
    using WarpReduceCfg = decltype(make_WarpReduceCfg<BShape, BFrgThr>());
    using CShape = typename WarpReduceCfg::xShape;
    using CFrgThr = typename WarpReduceCfg::xFrgThr;
    // Smem reduction
    using SmemReduce = CTASmemReduceCfg<T, CShape, CFrgThr, xFrgThr, smem_buffer_size>;
    ThreadReduceCfg thread_reduce;
    WarpReduceCfg warp_reduce;
    SmemReduce smem_reduce;
};

template<bool carry_initial, typename Cfg, typename Fn>
CUTE_DEVICE void CTA_reduce(Cfg& cfg, auto& rX_frg, auto& rx_frg, Fn& fn, char* smem_buff)  {
    auto rB_frg = thread_reduce(cfg.thread_reduce, fn, rX_frg);
    warp_reduce(cfg.warp_reduce, fn, rB_frg); // rC_frg = rB_frg
    CTA_smem_reduce<carry_initial>(cfg.smem_reduce, rB_frg, rx_frg, fn, smem_buff);
}


// -------------- Callables --------------
struct SumCallable {
    template<typename Ta, typename Tb> CUTE_HOST_DEVICE
    auto operator()(Ta a, Tb b) {
        return a + b;
    }
};

struct MaxCallable {
    template<typename Ta, typename Tb> CUTE_HOST_DEVICE
    auto operator()(Ta a, Tb b) {
        return a > b ? a : b;
    }
};



} // namespace vidrial
