/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/
#include <cnnl.h>
#include <torch/extension.h>
#include "kernels/rotary_embedding.mluh"
#include "register_api.h"
#include "utils.h"

namespace tmo {
namespace kernel_test_api {
at::Tensor rotary_embedding_test(const at::Tensor &input,
                                 const at::Tensor &sin_table,
                                 const at::Tensor &cos_table,
                                 const c10::optional<at::Tensor> &seq_offsets,
                                 const c10::optional<at::Tensor> &cu_seq_lens,
                                 const bool &interleaved,
                                 const bool &discrete,
                                 const bool &dynamic_ntk,
                                 const bool &rope_2d,
                                 const int &max_seq_len,
                                 const bool &no_offset) {
  MLU_TENSOR_CHECK_FATAL(input);
  MLU_TENSOR_CHECK_FATAL(sin_table);
  MLU_TENSOR_CHECK_FATAL(cos_table);

  bool has_seq_offsets = seq_offsets.has_value();
  bool has_cu_seq_lens = cu_seq_lens.has_value();
  int *seq_offsets_ptr =
      has_seq_offsets ? reinterpret_cast<int *>(seq_offsets.value().data_ptr()) : nullptr;
  int *cu_seq_lens_ptr =
      has_cu_seq_lens ? reinterpret_cast<int *>(cu_seq_lens.value().data_ptr()) : nullptr;

  int batch = 0;
  int total_seq_len = 0;
  int head_size = input.size(-1);

  if (input.dim() == 3) {  // pack mode
    TORCH_CHECK(has_cu_seq_lens,
                "input has 3 dims: (total_seq_len, head_num, head_size),"
                " which means pack mode, cu_seqlens should not be None");
    total_seq_len = input.size(0);
    batch = cu_seq_lens.value().size(0) - 1;
  } else if (input.dim() == 4) {
    TORCH_CHECK(!has_cu_seq_lens,
                "input has 4 dims: (batch_size, seq_len, head_num, head_size),"
                " which means pad mode, cu_seqlens should be None");
    TORCH_CHECK(max_seq_len == input.size(1),
                "input has 4 dims: (batch_size, seq_len, head_num, head_size),"
                " which means pad mode, max_seqlen must be equtals to input.size(1)");
    batch = input.size(0);
    total_seq_len = batch * input.size(1);
  } else {
    TORCH_CHECK(false, "input only support 3 or 4 dims");
  }

  const int rope_seqlen = dynamic_ntk ? sin_table.size(1) : sin_table.size(0);
  const int rope_dim = dynamic_ntk ? sin_table.size(2) : sin_table.size(1);

#if 0
  if (has_seq_offsets) {
    if (discrete) {
      CHECK_SHAPE(seq_offsets.value(), total_seq_len);
    } else {
      CHECK_SHAPE(seq_offsets.value(), batch);
    }
  }
  if (dynamic_ntk) {
    CHECK_SHAPE(sin_table, batch, rope_seqlen, rope_dim);
    CHECK_SHAPE(cos_table, batch, rope_seqlen, rope_dim);
  } else {
    CHECK_SHAPE(sin_table, rope_seqlen, rope_dim);
    CHECK_SHAPE(cos_table, rope_seqlen, rope_dim);
  }
#endif

  TORCH_CHECK(input.stride(-1) == 1, "input last dim must be contiguous");
  if (dynamic_ntk) {
    TORCH_CHECK(sin_table.stride(1) == cos_table.stride(1),
                "sin_table second stride must be equal to cos_table second stride");
  } else {
    TORCH_CHECK(sin_table.stride(0) == cos_table.stride(0),
                "sin_table first stride must be equal to cos_table second stride");
  }

  if (has_seq_offsets) {
    TORCH_CHECK(seq_offsets.value().is_contiguous(), "seq_offsets must be contiguous");
  }

  if (has_cu_seq_lens) {
    TORCH_CHECK(cu_seq_lens.value().is_contiguous(), "cu_seq_lens must be contiguous");
  }

  // auto output = input;
  // if (is_qkv) {
  // }

  int dims = input.dim();
  int head_num = input.size(dims - 2);  // head_num_qk
  TORCH_CHECK(head_size <= 256, "only support input head_size <= 256");
  int head_dim = input.size(dims - 1);
  int rotary_seq_len = dynamic_ntk ? sin_table.size(1) : sin_table.size(0);
  int rotary_dim = dynamic_ntk ? sin_table.size(2) : sin_table.size(1);

#if 0
  if (rope_2d) {
    total_seq_len = input.size(0);
    rotary_seq_len *= 2;
  }
#endif
  int rotary_stride = dynamic_ntk ? sin_table.stride(1) : sin_table.stride(0);

  size_t input_seq_stride = input.stride(dims - 3);
  size_t input_head_stride = input.stride(dims - 2);

  cnnlHandle_t handle = torch_mlu::getCurrentHandle();
  cnrtQueue_t queue;
  cnnlGetQueue(handle, &queue);

#if 0
  std::shared_ptr<torch::Device> dev = std::make_shared<torch::Device>(input.device());

  std::cout
    << " batch * max_seq_len : " << batch * max_seq_len
    << " max_seq_len : " << max_seq_len
    << " total_seq_len : " << total_seq_len
    << " head_num : " << head_num
    << " head_size : " << head_size << "\n";

  // at::IntArrayRef output_shape = {batch, max_seq_len, head_num, head_size};
  // if (input.dim() == 3) {
  //   output_shape = {batch * max_seq_len, head_num, head_size};
  // }
#endif

  cnnlDataType_t dtype = torch2cnnlDataType(torch::typeMetaToScalarType(input.dtype()));
  auto output = input;

  size_t output_seq_stride = output.stride(dims - 3);
  size_t output_head_stride = output.stride(dims - 2);

#if 0
  std::cout << " <<<< batch:  " << batch << "\n";
  std::cout << " <<<< max_seq_len:  " << max_seq_len << "\n";
  std::cout << " <<<< total_seq_len:  " << total_seq_len << "\n";
  std::cout << " <<<< head_num:  " << head_num << "\n";
  std::cout << " <<<< head_size:  " << head_size << "\n";
  std::cout << " <<<< rotary_seq_len_:  " << rotary_seq_len << "\n";
  std::cout << " <<<< rotary_dim:  " << rotary_dim << "\n";
  std::cout << " <<<< rotary_stride:  " << rotary_stride << "\n";
  std::cout << " <<<< input_seq_stride:  " << input_seq_stride << "\n";
  std::cout << " <<<< input_head_stride:  " << input_head_stride << "\n";
  std::cout << " <<<< output_seq_stride:  " << output_seq_stride << "\n";
  std::cout << " <<<< output_head_stride:  " << output_head_stride << "\n";
  std::cout << " <<<< interleaved:  " << interleaved << "\n";
  std::cout << " <<<< discrete:     " << discrete << "\n";
  std::cout << " <<<< dynamic_ntk:  " << dynamic_ntk << "\n";
#endif

  size_t io_bytes =
      (size_t)total_seq_len * head_num * (head_size + rotary_dim) * dtype_size_map[dtype] * 2;
  io_bytes += (discrete && !no_offset) ? total_seq_len * sizeof(int) : batch * sizeof(int);

  cnrtNotifier_t time_begin;
  cnrtNotifier_t time_end;

  CNRT_CHECK(cnrtNotifierCreate(&time_begin));
  CNRT_CHECK(cnrtNotifierCreate(&time_end));

  CNRT_CHECK(cnrtPlaceNotifier(time_begin, queue));
  if (rope_2d) {
    TMO_KERNEL_CHECK_FATAL(invokeGlm6BRotaryEmbedding(
        queue, output.data_ptr(), input.data_ptr(), sin_table.data_ptr(), cos_table.data_ptr(),
        seq_offsets_ptr, cu_seq_lens_ptr, batch, max_seq_len, total_seq_len, head_num, head_size,
        rotary_seq_len, rotary_stride, input_seq_stride, input_head_stride, output_seq_stride,
        output_head_stride, interleaved, dtype));
  } else {
    TMO_KERNEL_CHECK_FATAL(invokeRotaryEmbedding(
        queue, output.data_ptr(), input.data_ptr(), sin_table.data_ptr(), cos_table.data_ptr(),
        seq_offsets_ptr, cu_seq_lens_ptr, batch, max_seq_len, head_num, head_size, rotary_seq_len,
        rotary_dim, rotary_stride, input_seq_stride, input_head_stride, output_seq_stride,
        output_head_stride, interleaved, discrete, dynamic_ntk, dtype));
  }
  CNRT_CHECK(cnrtPlaceNotifier(time_end, queue));
  cnrtQueueSync(queue);
  float usec = 0.0f;
  CNRT_CHECK(cnrtNotifierDuration(time_begin, time_end, &usec));
  print_info(usec, io_bytes);

  return output;
}

}  // namespace kernel_test_api
}  // namespace tmo
