#include "xpu/kernel/cluster.h"
#include "xpu/kernel/cluster_partition.h"
#include "xpu/kernel/cluster_primitive.h"
#include "xpu/kernel/xtdk_io.h"

namespace xpu3 {
namespace plugin {

template <typename T>
static __device__ inline void text_image_gather(
        __global_ptr__ T* input,
        __global_ptr__ T* text_input,
        __global_ptr__ T* image_input,
        __global_ptr__ int* token_type_ids,
        __global_ptr__ int* text_index,
        __global_ptr__ int* image_index,
        int64_t token_num,
        int64_t text_token_num,
        int64_t image_token_num,
        int64_t hidden_size,
        T* input_lm) {
    int cid = core_id();
    int clusterid = cluster_id();
    int token_start_cluster;
    int token_end_cluster;
    int token_start_core;
    int token_end_core;

    const int BUFSIZE = 2 * 1024 / sizeof(T);       // 1024 for bf16, 512 for fp32
    // cluster partition
    partition(cluster_id(), cluster_num(), (int)token_num, 1, &token_start_cluster, &token_end_cluster);
    if (token_start_cluster >= token_end_cluster) {
        return;
    }
    int rows_cluster = token_end_cluster - token_start_cluster;    // total rows for a cluster
    // core partition
    partition(core_id(), core_num(), rows_cluster, 1, &token_start_core, &token_end_core);
    int rows_core = token_end_core - token_start_core;    // total rows for a core
    token_start_core += token_start_cluster;
    token_end_core += token_start_cluster;

    int read_len;
    for (int i = token_start_core; i < token_end_core; i += 1) {
        int token_type, text_image_token_idx;
        __global_ptr__ T* text_image_input = nullptr;
        __global_ptr__ int* text_image_index = nullptr;

        GM2LM(token_type_ids + i, &token_type, sizeof(int));
        if (token_type == 0) {
            text_image_input = text_input;
            text_image_index = text_index;
        } else {
            text_image_input = image_input;
            text_image_index = image_index;
        }
        GM2LM(text_image_index + i, &text_image_token_idx, sizeof(int));
        int input_offset = i * hidden_size;
        int text_image_offset = text_image_token_idx * hidden_size;

        for (int j = 0; j < hidden_size; j += BUFSIZE) {
            read_len = min(hidden_size - j, BUFSIZE);
            GM2LM(text_image_input + text_image_offset + j, input_lm, sizeof(T) * read_len);
            LM2GM(input_lm, input + input_offset + j, sizeof(T) * read_len);
        }
    }
}

template <typename T>
static __device__ inline void text_image_scatter(
        __global_ptr__ T* input,
        __global_ptr__ T* text_input,
        __global_ptr__ T* image_input,
        __global_ptr__ int* token_type_ids,
        __global_ptr__ int* text_index,
        __global_ptr__ int* image_index,
        int64_t token_num,
        int64_t text_token_num,
        int64_t image_token_num,
        int64_t hidden_size,
        T* input_lm) {
    int cid = core_id();
    int clusterid = cluster_id();
    int token_start_cluster;
    int token_end_cluster;
    int token_start_core;
    int token_end_core;

    const int BUFSIZE = 2 * 1024 / sizeof(T);       // 1024 for bf16, 512 for fp32
    // cluster partition
    partition(cluster_id(), cluster_num(), (int)token_num, 1, &token_start_cluster, &token_end_cluster);
    if (token_start_cluster >= token_end_cluster) {
        return;
    }
    int rows_cluster = token_end_cluster - token_start_cluster;    // total rows for a cluster
    // core partition
    partition(core_id(), core_num(), rows_cluster, 1, &token_start_core, &token_end_core);
    int rows_core = token_end_core - token_start_core;    // total rows for a core
    token_start_core += token_start_cluster;
    token_end_core += token_start_cluster;

    int read_len;
    for (int i = token_start_core; i < token_end_core; i += 1) {
        int token_type, text_image_token_idx;
        __global_ptr__ T* text_image_input = nullptr;
        __global_ptr__ int* text_image_index = nullptr;

        GM2LM(token_type_ids + i, &token_type, sizeof(int));
        if (token_type == 0) {
            text_image_input = text_input;
            text_image_index = text_index;
        } else {
            text_image_input = image_input;
            text_image_index = image_index;
        }
        GM2LM(text_image_index + i, &text_image_token_idx, sizeof(int));
        int input_offset = i * hidden_size;
        int text_image_offset = text_image_token_idx * hidden_size;

        for (int j = 0; j < hidden_size; j += BUFSIZE) {
            read_len = min(hidden_size - j, BUFSIZE);
            GM2LM(input + input_offset + j, input_lm, sizeof(T) * read_len);
            LM2GM(input_lm, text_image_input + text_image_offset + j, sizeof(T) * read_len);
        }
    }
}

template <typename T>
__global__ void text_image_gather_scatter(
        T* input,
        T* text_input,
        T* image_input,
        int* token_type_ids,
        int* text_index,
        int* image_index,
        int64_t token_num,
        int64_t text_token_num,
        int64_t image_token_num,
        int64_t hidden_size,
        bool is_scatter) {
    int cid = core_id();
    int ncores = core_num();
    int clusterid = cluster_id();
    int nclusters = cluster_num();
    const int BUFSIZE = 2 * 1024 / sizeof(T);   // 1024 for bf16, 512 for fp32
    __simd__ T input_lm[BUFSIZE];               // 2KB for bf16 and fp32
    if (is_scatter) {
        text_image_scatter(
            input, text_input, image_input, token_type_ids, text_index, image_index,
            token_num, text_token_num, image_token_num, hidden_size, input_lm);
    } else {
        text_image_gather(
            input, text_input, image_input, token_type_ids, text_index, image_index,
            token_num, text_token_num, image_token_num, hidden_size, input_lm);
    }
}


#define _XPU_DEF_TEXT_IMAGE_GATHER_SCATTER(T)              \
  template __global__ void text_image_gather_scatter<T>(    \
    T* input,                                               \
    T* text_input,                                          \
    T* image_input,                                         \
    int* token_type_ids,                                    \
    int* text_index,                                        \
    int* image_index,                                       \
    int64_t token_num,                                      \
    int64_t text_token_num,                                 \
    int64_t image_token_num,                                \
    int64_t hidden_size,                                    \
    bool is_scatter);

_XPU_DEF_TEXT_IMAGE_GATHER_SCATTER(bfloat16);

}  // namespace plugin
}  // namespace xpu3
