/*************************************************************************
 * 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 <algorithm>
#include <cstdint>
#include <vector>
#include "torch_ops_api.h"
#include "utils.h"

namespace tmo {
namespace torch_api {

using GroupGemmDesc = op_desc::GroupGemmDesc;
using QuantMode = GroupGemmDesc::QuantMode;

at::Tensor group_gemm(const at::Tensor &a_tensor,
                      const at::Tensor &b_tensor,
                      const at::Tensor &m_list,
                      const c10::optional<at::Tensor> &gather_idx,
                      const c10::optional<at::Tensor> &c_tensor,
                      const c10::optional<at::Tensor> &alpha,
                      const c10::optional<at::Tensor> &beta,
                      const c10::optional<at::Tensor> &a_scale,
                      const c10::optional<at::Tensor> &b_scale,
                      const c10::optional<at::Tensor> &bias,
                      const c10::optional<std::string> &data_type,
                      const c10::optional<at::List<int64_t>> &quant_flag,
                      const c10::optional<at::Tensor> &b_offset,
                      const int64_t max_m) {
  bool has_dtype = data_type.has_value();
  bool has_c_tensor = c_tensor.has_value();
  bool has_gather_idx = gather_idx.has_value();
  bool has_a_scale = a_scale.has_value();
  bool has_b_scale = b_scale.has_value();
  bool has_alpha = alpha.has_value();
  bool has_beta = beta.has_value();
  bool has_bias = bias.has_value();
  bool has_quant_flag = quant_flag.has_value();
  QuantMode quant_mode = QuantMode::noQuant;
  bool quant_grouped = false;

  // check stride
  TORCH_CHECK(a_tensor.stride(-1) == 1, "a_tensor last dim must be contiguous");
  CHECK_TENSOR_CONTIGUOUS(m_list);
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(c_tensor);
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(gather_idx);
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(a_scale);
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(b_scale);
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(alpha);
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(beta);

  // check dim
  TORCH_CHECK(m_list.dim() == 1, "the shape of m_list should be 2D");
  TORCH_CHECK(a_tensor.dim() == 2, "the shape of a_tensor should be 2D");
  TORCH_CHECK(b_tensor.dim() == 3 || b_tensor.dim() == 4 || (b_tensor.dim() == 1 && has_quant_flag),
              "the shape of b_tensor should be 3D or 4D");
  if (has_c_tensor) {
    TORCH_CHECK(c_tensor.value().dim() == 2, "the shape of c_tensor should be 2D");
  }
  if (has_gather_idx) {
    TORCH_CHECK(gather_idx.value().dim() == 1, "the shape of gather_idx should be 1D");
  }
  if (has_a_scale) {
    TORCH_CHECK(a_scale.value().dim() == 1, "the shape of a_scale should be 1D");
  }
  if (has_b_scale) {
    TORCH_CHECK(b_scale.value().dim() == 2 || b_scale.value().dim() == 3,
                "the shape of b_scale should be 2D or 3D");
  }
  if (has_alpha) {
    TORCH_CHECK(alpha.value().dim() == 1, "the shape of alpha should be 1D");
  }
  if (has_beta) {
    TORCH_CHECK(beta.value().dim() == 1, "the shape of beta should be 1D");
  }

  // check shape
  int64_t experts_num = m_list.size(0);
  int64_t num_token = a_tensor.size(0);
  int64_t k = a_tensor.size(-1);
  int64_t n = has_quant_flag ? b_scale.value().size(2) : b_tensor.size(1);
  int64_t lda = a_tensor.stride(0);
  int64_t total_m = has_gather_idx ? gather_idx.value().size(0) : num_token;

  TORCH_CHECK(experts_num > 0, "number of groups must be large than zero");
  CHECK_SHAPE(a_tensor, num_token, k);
  if (has_a_scale) {
    CHECK_SHAPE(a_scale.value(), total_m);
  }
  if (has_quant_flag) {
    quant_grouped = true;
    quant_mode = QuantMode::W4W8;
    TORCH_CHECK(has_b_scale, "if quant_flag given, b_scale must exist");
    CHECK_SHAPE(b_scale.value(), b_scale.value().size(0), experts_num, n);
  } else if (has_b_scale) {
    int64_t b_k = b_tensor.size(-1);
    quant_grouped = b_scale.value().dim() == 3 ? true : false;
    if (quant_grouped) {
      CHECK_SHAPE(b_scale.value(), b_scale.value().size(0), experts_num, n);
    } else {
      CHECK_SHAPE(b_scale.value(), experts_num, n);
    }
    TORCH_CHECK(k == b_k || k == b_k * 2, "k == b_k || k == b_k * 2.");
    quant_mode = (b_k == k) ? QuantMode::W8 : QuantMode::W4;
  }
  if (has_alpha) {
    CHECK_SHAPE(alpha.value(), experts_num);
  }
  if (has_beta) {
    CHECK_SHAPE(beta.value(), experts_num);
  }
  auto b_shape = b_tensor.sizes();
  cnnlDataType_t a_dtype = getCnnlDataType(a_tensor.scalar_type());
  cnnlDataType_t b_dtype = getCnnlDataType(b_tensor.scalar_type());
  torch::Dtype dtype = a_tensor.scalar_type();
  if (quant_mode == QuantMode::W8 || quant_mode == QuantMode::noQuant) {
    if (b_tensor.dim() == 3) {
      CHECK_SHAPE(b_tensor, experts_num, n, k);
    } else {
      TORCH_CHECK(b_shape[0] * b_shape[2] == experts_num, "b_shape[0] * b_shape[2] == experts_num");
      CHECK_SHAPE(b_tensor, b_shape[0], n, b_shape[2], k);
    }
  } else if (quant_mode == QuantMode::W4) {
    CHECK_SHAPE(b_tensor, experts_num, n, k / 2);
    b_dtype = CNNL_DTYPE_INT4X2;
  }

  TORCH_CHECK(m_list.dtype() == torch::kInt32, "data type of m_list should be int32");
  if (has_gather_idx) {
    TORCH_CHECK(gather_idx.value().dtype() == torch::kInt32,
                "data type of gather_idx should be int32");
  }

  if (has_a_scale) {
    TORCH_CHECK(has_dtype, "data_type must given when a_scale and b_scale are given");
    TORCH_CHECK(data_type.value() == "float" || data_type.value() == "half" ||
                    data_type.value() == "bfloat16",
                "data_type must be 'float', 'half' or 'bfloat16'.");
    dtype = data_type.value() == "float"  ? torch::kFloat32
            : data_type.value() == "half" ? torch::kFloat16
                                          : torch::kBFloat16;
  }

  const torch_mlu::mlu::MLUGuard device_guard(a_tensor.device());
  at::Tensor d_tensor = at::empty({total_m, n}, a_tensor.options().dtype(dtype));
  cnnlDataType_t d_dtype = getCnnlDataType(d_tensor.scalar_type());

  // check device
  TORCH_CHECK(isMlu(m_list), "m_list must on mlu");
  TORCH_CHECK(isMlu(a_tensor), "a_tensor must on mlu");
  TORCH_CHECK(isMlu(b_tensor), "b_tensor must on mlu");
  TORCH_CHECK(isMlu(d_tensor), "d_tensor must on mlu");
  if (has_c_tensor) {
    TORCH_CHECK(isMlu(c_tensor.value()), "c_tensor must on mlu");
  }
  if (has_gather_idx) {
    TORCH_CHECK(isMlu(gather_idx.value()), "gather_idx must on mlu");
  }
  if (has_a_scale) {
    TORCH_CHECK(isMlu(a_scale.value()), "a_scale must on mlu");
  }
  if (has_b_scale) {
    TORCH_CHECK(isMlu(b_scale.value()), "b_scale must on mlu");
  }
  if (has_alpha) {
    TORCH_CHECK(isMlu(alpha.value()), "alpha must on mlu");
  }
  if (has_beta) {
    TORCH_CHECK(isMlu(beta.value()), "beta must on mlu");
  }

  cnnlHandle_t handle = torch_mlu::getCurrentHandle();
  GroupGemmDesc group_gemm_desc(experts_num, max_m, n, k, d_dtype, has_gather_idx, quant_mode);
  group_gemm_desc.setInputOutputTensor(a_dtype, b_dtype, d_dtype, CNNL_DTYPE_INT32,
                                       getAtTensorPtr(a_tensor), getAtTensorPtr(b_tensor),
                                       getAtTensorPtr(c_tensor), getAtTensorPtr(d_tensor),
                                       getAtTensorPtr(gather_idx), k, lda, total_m, has_c_tensor,
                                       (int64_t *)getAtTensorPtr(b_offset));
  std::vector<int> flag_vec;
  if (quant_mode != QuantMode::noQuant || has_bias) {
    if (has_quant_flag) {
      auto vec = quant_flag.value().vec();
      flag_vec.resize(vec.size());
      std::copy(vec.begin(), vec.end(), flag_vec.begin());
    }
    group_gemm_desc.setPerRowColScaleBiasAct(
        has_a_scale ? getAtTensorPtr(a_scale) : nullptr,
        has_b_scale ? getAtTensorPtr(b_scale) : nullptr, flag_vec.data(),
        has_bias ? getAtTensorPtr(bias) : nullptr, d_dtype, quant_grouped ? k : 0,
        quant_grouped ? k / b_scale.value().size(0) : 0, 0);
  }

  size_t group_gemm_wsize =
      tmo::ops::getGroupGemmWorkspaceSize(handle, group_gemm_desc, experts_num);
  auto group_gemm_workspace =
      at::empty({static_cast<int64_t>(group_gemm_wsize)}, a_tensor.options().dtype(at::kByte));

  std::vector<int> ldb_array;
  // if (quant_mode == QuantMode::W4W8) {
  //   ldb_array.resize(experts_num);
  //   auto q_group = b_scale.value().size(0);
  //   for (auto i = 0; i < experts_num; ++i) {
  //     int sum = 0;
  //     for (auto j = 0; j < q_group; j++) {
  //       sum += flag_vec[i * q_group + j];
  //     }
  //     ldb_array[i] = (sum / 4) * (k / q_group / 2);
  //   }
  // }
  if (quant_mode != QuantMode::W4W8) {
    int ldb = b_tensor.dim() == 3 ? k : b_shape[2] * k;
    ldb_array.assign(experts_num, ldb);
  }
  GroupGemmTheory obj(total_m, experts_num, k, n, has_c_tensor, a_dtype, getCnnlDataType(dtype));
  cnpxPush(obj);
  ops::GroupGemm(handle, group_gemm_desc, getAtTensorPtr(m_list), getAtTensorPtr(alpha),
                 getAtTensorPtr(beta), getAtTensorPtr(group_gemm_workspace), group_gemm_wsize,
                 experts_num, k, n, lda, ldb_array);
  cnpxPop();
  return d_tensor;
}
}  //  namespace torch_api
}  //  namespace tmo
