#include <cmath>
#include <cstddef>
#include "cnnl.h"
#include "cnrt.h"
#include "create_cos_sin_table.mluh"

namespace {
// constexpr int LINEAR_SCALING = 0;
// constexpr int FIX_NTK_SCALING = 1;
constexpr int DYNAMIC_NTK_SCALING = 2;
}  // namespace

namespace tmo {
namespace kernels {
__nram__ int8_t nram_buffer[__MLU_NRAM_SIZE__ * 1024 - 32 * 1024];
__nram__ const float range[64] = {
    0.0F,   2.0F,   4.0F,   6.0F,   8.0F,   10.0F,  12.0F,  14.0F,  16.0F,  18.0F,  20.0F,
    22.0F,  24.0F,  26.0F,  28.0F,  30.0F,  32.0F,  34.0F,  36.0F,  38.0F,  40.0F,  42.0F,
    44.0F,  46.0F,  48.0F,  50.0F,  52.0F,  54.0F,  56.0F,  58.0F,  60.0F,  62.0F,  64.0F,
    66.0F,  68.0F,  70.0F,  72.0F,  74.0F,  76.0F,  78.0F,  80.0F,  82.0F,  84.0F,  86.0F,
    88.0F,  90.0F,  92.0F,  94.0F,  96.0F,  98.0F,  100.0F, 102.0F, 104.0F, 106.0F, 108.0F,
    110.0F, 112.0F, 114.0F, 116.0F, 118.0F, 120.0F, 122.0F, 124.0F, 126.0F};

__mlu_func__ void genRangeDims(float *range_nram, int elem_count) {
  int count = 64;
  __bang_move(range_nram, range, std::min(count, elem_count) * sizeof(float));
  while (count < elem_count) {
    __bang_add_scalar(range_nram + count, range_nram, (float)count * 2.0F,
                      std::min(count, elem_count - count));
    count *= 2;
  }
}

__mlu_func__ int getBatchMaxSeqLen(int *seq_lens_nram, int *seq_lens, int batch) {
  __memcpy(seq_lens_nram, seq_lens, batch * sizeof(int), GDRAM2NRAM);
  __bang_argmax((float *)seq_lens_nram, (float *)seq_lens_nram, batch);
  return __load_nram(seq_lens_nram);
}

__mlu_func__ float getNTKAlpha(int curr_seq_len, int max_position_embeddings, int kv_seq_len) {
  int seq_len = kv_seq_len > max_position_embeddings ? curr_seq_len : kv_seq_len;
  float context_value = std::log2((float)seq_len / (float)max_position_embeddings) + 1.0F;
  float ntk_alpha = std::pow(2.0F, std::ceil(context_value)) - 1.0F;
  return std::max(ntk_alpha, 1.0F);
}

__mlu_func__ void getRotaryInvFreq(float *inv_freq_nram,
                                   float *base_nram,
                                   float *range_nram,
                                   float base,
                                   int rotary_dim,
                                   int elem_count) {
  // inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2) / dim))
  __bang_write_value(base_nram, elem_count, base);
  __bang_mul_scalar(inv_freq_nram, range_nram, 1.0F / (float)rotary_dim, elem_count);
  __bang_log(base_nram, base_nram, elem_count);
  __bang_mul(inv_freq_nram, inv_freq_nram, base_nram, elem_count);
  __bang_pow2(inv_freq_nram, inv_freq_nram, elem_count);
  __bang_recip(inv_freq_nram, inv_freq_nram, elem_count);
}

template <typename T>
__mlu_func__ void convertCosSinTable(float *cos_table, float *sin_table, int elem_count) {}

template <>
__mlu_func__ void convertCosSinTable<half>(float *cos_table, float *sin_table, int elem_count) {
  __bang_float2half((half *)cos_table, cos_table, elem_count);
  __bang_float2half((half *)sin_table, sin_table, elem_count);
}

template <>
__mlu_func__ void convertCosSinTable<bfloat16_t>(float *cos_table,
                                                 float *sin_table,
                                                 int elem_count) {
  __bang_float2bfloat16((bfloat16_t *)cos_table, cos_table, elem_count);
  __bang_float2bfloat16((bfloat16_t *)sin_table, sin_table, elem_count);
}

__mlu_global__ void MLUUpdateCachedAlpha(float *rotary_emb_alpha_cached,
                                         int *seq_lens,
                                         int max_position_embeddings,
                                         int batch) {
  int *seq_lens_nram = (int *)nram_buffer;  // [batch]
  int kv_seq_len = getBatchMaxSeqLen(seq_lens_nram, seq_lens, batch);
  rotary_emb_alpha_cached[taskIdY] =
      getNTKAlpha(seq_lens[taskIdY], max_position_embeddings, kv_seq_len);
}

template <typename T>
__mlu_global__ void MLUCreateCosSinTableKernel(void *cos_sin_table,
                                               float *rotary_emb_alpha_cached,
                                               int *seq_lens,
                                               int max_position_embeddings,
                                               int batch,
                                               int batch_stride,
                                               int rotary_seq_len,
                                               int rotary_dim,
                                               int rotary_stride,
                                               float rotary_base,
                                               float rotary_scaling,
                                               int rotary_scaling_type,
                                               int seq_seg,
                                               bool interleaved,
                                               cnnlDataType_t dtype) {
  int half_rotary_dim = rotary_dim / 2;
  float *base_nram = (float *)nram_buffer;               // [rotary_dim / 2]
  float *range_nram = base_nram + half_rotary_dim;       // [rotary_dim / 2]
  float *inv_freq_nram = range_nram + half_rotary_dim;   // [rotary_dim / 2]
  float *freqs_nram = inv_freq_nram + half_rotary_dim;   // [rotary_dim / 2]
  float *cos_nram = freqs_nram + half_rotary_dim;        // [rotary_dim]
  float *sin_nram = cos_nram + rotary_dim;               // [rotary_dim]
  float *swap_nram = sin_nram + rotary_dim;              // [rotary_dim]
  int *seq_lens_nram = (int *)(swap_nram + rotary_dim);  // [batch]

  genRangeDims(range_nram, half_rotary_dim);

  float adjust_base = rotary_base;
  if (rotary_scaling_type == DYNAMIC_NTK_SCALING) {
    int kv_seq_len = getBatchMaxSeqLen(seq_lens_nram, seq_lens, batch);
    float ntk_alpha = getNTKAlpha(seq_lens[taskIdY], max_position_embeddings, kv_seq_len);
    if (rotary_emb_alpha_cached[taskIdY] == ntk_alpha) {
      return;
    }
    adjust_base = rotary_base * std::pow(ntk_alpha, (float)rotary_dim / (float)(rotary_dim - 2));
  }

  getRotaryInvFreq(inv_freq_nram, base_nram, range_nram, adjust_base, rotary_dim, half_rotary_dim);

  int seq_start = taskIdX * seq_seg;
  int seq_end = (taskIdX + 1) * seq_seg > rotary_seq_len ? rotary_seq_len : (taskIdX + 1) * seq_seg;

  T *cos_table = (T *)cos_sin_table + (size_t)taskIdY * batch_stride;
  T *sin_table = cos_table + rotary_dim;

  for (int idx = seq_start; idx < seq_end; ++idx) {
    __bang_mul_scalar(freqs_nram, inv_freq_nram, idx, half_rotary_dim);
    __bang_cos(cos_nram, freqs_nram, half_rotary_dim);
    __bang_sin(sin_nram, freqs_nram, half_rotary_dim);
    convertCosSinTable<T>(cos_nram, sin_nram, half_rotary_dim);
    if (!interleaved) {
      __memcpy(cos_table + idx * rotary_stride, cos_nram, half_rotary_dim * sizeof(T), NRAM2GDRAM,
               half_rotary_dim * sizeof(T), 0, 1);
      __memcpy(sin_table + idx * rotary_stride, sin_nram, half_rotary_dim * sizeof(T), NRAM2GDRAM,
               half_rotary_dim * sizeof(T), 0, 1);
    } else {
      __bang_move((T *)cos_nram + half_rotary_dim, (T *)cos_nram, half_rotary_dim * sizeof(T));
      __bang_transpose((T *)swap_nram, (T *)cos_nram, 2, half_rotary_dim);
      __memcpy(cos_table + idx * rotary_stride, (T *)swap_nram, half_rotary_dim * 2 * sizeof(T),
               NRAM2GDRAM);
      __bang_move((T *)sin_nram + half_rotary_dim, (T *)sin_nram, half_rotary_dim * sizeof(T));
      __bang_transpose((T *)cos_nram, (T *)sin_nram, 2, half_rotary_dim);
      __memcpy((T *)sin_table + idx * rotary_stride, (T *)cos_nram, half_rotary_dim * 2 * sizeof(T),
               NRAM2GDRAM);
    }
  }
}

#if __BANG_ARCH__ < 592
template <>
__mlu_global__ void MLUCreateCosSinTableKernel<bfloat16_t>(void *cos_sin_table,
                                                           float *rotary_emb_alpha_cached,
                                                           int *seq_lens,
                                                           int max_position_embeddings,
                                                           int batch,
                                                           int batch_stride,
                                                           int rotary_seq_len,
                                                           int rotary_dim,
                                                           int rotary_stride,
                                                           float rotary_base,
                                                           float rotary_scaling,
                                                           int rotary_scaling_type,
                                                           int seq_seg,
                                                           bool interleaved,
                                                           cnnlDataType_t dtype) {}
#endif
}  // namespace kernels

KernelStatus invokeCreateCosSinTable(cnrtQueue_t queue,
                                     void *cos_sin_table,
                                     float *rotary_emb_alpha_cached,
                                     int *seq_lens,
                                     int max_position_embeddings,
                                     int batch,
                                     int batch_stride,
                                     int rotary_seq_len,
                                     int rotary_dim,
                                     int rotary_stride,
                                     float rotary_base,
                                     float rotary_scaling,
                                     int rotary_scaling_type,
                                     bool interleaved,
                                     cnnlDataType_t data_type) {
  bool is_supported_dtype = data_type == CNNL_DTYPE_HALF || data_type == CNNL_DTYPE_FLOAT ||
                            data_type == CNNL_DTYPE_BFLOAT16;
  if (!is_supported_dtype) {
    std::cerr << "[invokeCreateCosSinTable]: unsupport data type for create cos sin table kernel."
              << std::endl;
    return KernelStatus::KERNEL_STATUS_FAILED;
  }

  // clang-format off
  void (*create_sin_cos_kernels[])(
    void*,                                               /* cos_sin_table           */
    float*,                                              /* rotary_emb_alpha_cached */
    int*,                                                /* seq_lens                */
    int,                                                 /* max_position_embeddings */
    int,                                                 /* batch                   */
    int,                                                 /* batch_stride            */
    int,                                                 /* rotary_seq_len          */
    int,                                                 /* rotary_dim              */
    int,                                                 /* rotary_stride           */
    float,                                               /* rotary_base             */
    float,                                               /* rotary_scaling          */
    int,                                                 /* rotary_scaling_type     */
    int,                                                 /* seq_seg                 */
    bool,                                                /* interleaved             */
    cnnlDataType_t                                       /* data_type               */
  ) = {
      kernels::MLUCreateCosSinTableKernel<half>,
      kernels::MLUCreateCosSinTableKernel<float>,
      kernels::MLUCreateCosSinTableKernel<bfloat16_t>
  };
  // clang-format on

  int kernel_index = 0;
  if (data_type == CNNL_DTYPE_HALF) {
    kernel_index = 0;
  } else if (data_type == CNNL_DTYPE_FLOAT) {
    kernel_index = 1;
  } else {
    kernel_index = 2;
  }

  CNdev dev;
  cnCtxGetDevice(&dev);
  int cluster_num = 1;
  int core_num = 1;
  CNRT_CHECK(cnrtDeviceGetAttribute(&cluster_num, cnrtAttrClusterCount, dev));
  CNRT_CHECK(cnrtDeviceGetAttribute(&core_num, cnrtAttrMcorePerCluster, dev));

  int used_core_num = std::min(rotary_seq_len, cluster_num * core_num);
  int seq_seg = (rotary_seq_len + used_core_num - 1) / used_core_num;

  cnrtDim3_t dim1;
  dim1.x = used_core_num;
  dim1.y = rotary_scaling_type == DYNAMIC_NTK_SCALING ? batch : 1;
  dim1.z = 1;

  if (data_type == CNNL_DTYPE_BFLOAT16 && !isBf16Supported()) {
    std::cerr << "[invokeCreateCosSinTable]: MLU300 devices do not support bfloat16." << std::endl;
    return KernelStatus::KERNEL_STATUS_FAILED;
  }

  create_sin_cos_kernels[kernel_index]<<<dim1, cnrtFuncTypeBlock, queue>>>(
      cos_sin_table, rotary_emb_alpha_cached, seq_lens, max_position_embeddings, batch,
      batch_stride, rotary_seq_len, rotary_dim, rotary_stride, rotary_base, rotary_scaling,
      rotary_scaling_type, seq_seg, interleaved, data_type);

  if (rotary_scaling_type == DYNAMIC_NTK_SCALING) {
    cnrtDim3_t dim2;
    dim2.x = 1;
    dim2.y = batch;
    dim2.z = 1;
    kernels::MLUUpdateCachedAlpha<<<dim2, cnrtFuncTypeBlock, queue>>>(
        rotary_emb_alpha_cached, seq_lens, max_position_embeddings, batch);
  }

  return KernelStatus::KERNEL_STATUS_SUCCESS;
}

}  // namespace tmo
