// Copyright (c) Microsoft Corporation. All rights reserved.
// Copyright (c) Huawei. All rights reserved.
// Licensed under the MIT License.

#include <unistd.h>
#include <algorithm>

#include "core/providers/cann/cann_utils.h"

namespace onnxruntime {
namespace cann {

template <typename T>
aclDataType getACLType() {
  return ACL_DT_UNDEFINED;
}

#define GET_ACL_TYPE(onnx_type, cann_type) \
  template <>                              \
  aclDataType getACLType<onnx_type>() {    \
    return cann_type;                      \
  }

GET_ACL_TYPE(int8_t, ACL_INT8);
GET_ACL_TYPE(int16_t, ACL_INT16);
GET_ACL_TYPE(int32_t, ACL_INT32);
GET_ACL_TYPE(int64_t, ACL_INT64);
GET_ACL_TYPE(uint8_t, ACL_UINT8);
GET_ACL_TYPE(uint16_t, ACL_UINT16);
GET_ACL_TYPE(uint32_t, ACL_UINT32);
GET_ACL_TYPE(uint64_t, ACL_UINT64);
GET_ACL_TYPE(float, ACL_FLOAT);
GET_ACL_TYPE(MLFloat16, ACL_FLOAT16);
GET_ACL_TYPE(BFloat16, ACL_BF16);
GET_ACL_TYPE(double, ACL_DOUBLE);
GET_ACL_TYPE(bool, ACL_BOOL);

template <typename T>
Status Fill(Tensor* y, void* addr, aclrtStream stream) {
  int64_t one = 1;
  int64_t dims = static_cast<int64_t>(y->Shape().NumDimensions());

  const aclDataType aclType = getACLType<T>();
  aclFormat format = ACL_FORMAT_ND;

  CannPreparation prepare;

  ORT_TRY {
    CANN_PREPARE_INPUTDESC(prepare, ACL_INT64, 1, &dims, format);
    CANN_CONST_INPUTDESC(prepare, 0, const_cast<int64_t*>(y->Shape().GetDims().data()), dims * sizeof(int64_t));
    CANN_PREPARE_INPUTDESC(prepare, aclType, 1, &one, format);
    CANN_PREPARE_OUTPUTDESC(prepare, aclType, dims, y->Shape().GetDims().data(), format);

    CANN_PREPARE_INPUTBUFFER(prepare, const_cast<int64_t*>(y->Shape().GetDims().data()), dims * sizeof(int64_t));
    CANN_PREPARE_INPUTBUFFER(prepare, addr, sizeof(T));
    CANN_PREPARE_OUTPUTBUFFER(prepare, y->MutableDataRaw(), y->SizeInBytes());
  }
  ORT_CATCH(const std::exception& e) {
    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, e.what());
  }

  CANN_RETURN_IF_ERROR(aclopCompileAndExecute("Fill",
                                              prepare.inputDesc_.size(),
                                              prepare.inputDesc_.data(),
                                              prepare.inputBuffers_.data(),
                                              prepare.outputDesc_.size(),
                                              prepare.outputDesc_.data(),
                                              prepare.outputBuffers_.data(),
                                              prepare.opAttr_,
                                              ACL_ENGINE_SYS,
                                              ACL_COMPILE_SYS,
                                              NULL,
                                              stream));

  return Status::OK();
}

template Status Fill<uint8_t>(Tensor* y, void* addr, aclrtStream stream);
template Status Fill<uint16_t>(Tensor* y, void* addr, aclrtStream stream);
template Status Fill<uint32_t>(Tensor* y, void* addr, aclrtStream stream);
template Status Fill<uint64_t>(Tensor* y, void* addr, aclrtStream stream);
template Status Fill<int8_t>(Tensor* y, void* addr, aclrtStream stream);
template Status Fill<int16_t>(Tensor* y, void* addr, aclrtStream stream);
template Status Fill<int32_t>(Tensor* y, void* addr, aclrtStream stream);
template Status Fill<int64_t>(Tensor* y, void* addr, aclrtStream stream);
template Status Fill<MLFloat16>(Tensor* y, void* addr, aclrtStream stream);
template Status Fill<float>(Tensor* y, void* addr, aclrtStream stream);
template Status Fill<double>(Tensor* y, void* addr, aclrtStream stream);
template Status Fill<bool>(Tensor* y, void* addr, aclrtStream stream);

template <typename T>
Status Broadcast(const Tensor* x, Tensor* y, void* addr, aclrtStream stream) {
  const aclDataType aclType = getACLType<T>();
  aclFormat format = ACL_FORMAT_ND;

  CannPreparation prepare;

  CANN_RETURN_IF_ERROR(aclopSetAttrListInt(prepare.opAttr_, "shape", static_cast<int64_t>(y->Shape().NumDimensions()),
                                           y->Shape().GetDims().data()));

  ORT_TRY {
    CANN_PREPARE_INPUTDESC(prepare, aclType, x->Shape().NumDimensions(), x->Shape().GetDims().data(), format);
    CANN_PREPARE_OUTPUTDESC(prepare, aclType, y->Shape().NumDimensions(), y->Shape().GetDims().data(), format);

    CANN_PREPARE_INPUTBUFFER(prepare, const_cast<void*>(x->DataRaw()), x->SizeInBytes());
    CANN_PREPARE_OUTPUTBUFFER(prepare, addr, y->SizeInBytes());
  }
  ORT_CATCH(const std::exception& e) {
    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, e.what());
  }

  CANN_RETURN_IF_ERROR(aclopCompileAndExecute("BroadcastToD",
                                              prepare.inputDesc_.size(),
                                              prepare.inputDesc_.data(),
                                              prepare.inputBuffers_.data(),
                                              prepare.outputDesc_.size(),
                                              prepare.outputDesc_.data(),
                                              prepare.outputBuffers_.data(),
                                              prepare.opAttr_,
                                              ACL_ENGINE_SYS,
                                              ACL_COMPILE_SYS,
                                              NULL,
                                              stream));

  return Status::OK();
}

template Status Broadcast<uint8_t>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);
template Status Broadcast<uint16_t>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);
template Status Broadcast<uint32_t>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);
template Status Broadcast<uint64_t>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);
template Status Broadcast<int8_t>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);
template Status Broadcast<int16_t>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);
template Status Broadcast<int32_t>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);
template Status Broadcast<int64_t>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);
template Status Broadcast<MLFloat16>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);
template Status Broadcast<BFloat16>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);
template Status Broadcast<float>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);
template Status Broadcast<double>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);
template Status Broadcast<bool>(const Tensor* x, Tensor* y, void* addr, aclrtStream stream);

Status aclrtblasGemmEx(aclTransType transA,
                       aclTransType transB,
                       aclTransType transC,
                       int m,
                       int n,
                       int k,
                       const void* alpha,
                       const void* matrixA,
                       int lda,
                       aclDataType dataTypeA,
                       const void* matrixB,
                       int ldb,
                       aclDataType dataTypeB,
                       const void* beta,
                       void* matrixC,
                       int ldc,
                       aclDataType dataTypeC,
                       aclComputeType type,
                       aclrtStream stream) {
  ORT_UNUSED_PARAMETER(transC);
  ORT_UNUSED_PARAMETER(lda);
  ORT_UNUSED_PARAMETER(ldb);
  ORT_UNUSED_PARAMETER(ldc);
  ORT_UNUSED_PARAMETER(type);

  TensorShape A;
  TensorShape B;
  TensorShape C{m, n};
  TensorShape shape{1};

  A = transA ? TensorShape{k, m} : TensorShape{m, k};
  B = transB ? TensorShape{n, k} : TensorShape{k, n};

  aclFormat format = ACL_FORMAT_ND;

  CannPreparation prepare;

  CANN_RETURN_IF_ERROR(aclopSetAttrBool(prepare.opAttr_, "transpose_a", transA));
  CANN_RETURN_IF_ERROR(aclopSetAttrBool(prepare.opAttr_, "transpose_b", transB));

  ORT_TRY {
    CANN_PREPARE_INPUTDESC(prepare, dataTypeA, A.NumDimensions(), A.GetDims().data(), format);
    CANN_PREPARE_INPUTDESC(prepare, dataTypeB, B.NumDimensions(), B.GetDims().data(), format);
    CANN_PREPARE_INPUTDESC(prepare, dataTypeC, C.NumDimensions(), C.GetDims().data(), format);
    CANN_PREPARE_INPUTDESC(prepare, dataTypeC, shape.NumDimensions(), shape.GetDims().data(), format);
    CANN_PREPARE_INPUTDESC(prepare, dataTypeC, shape.NumDimensions(), shape.GetDims().data(), format);
    CANN_PREPARE_OUTPUTDESC(prepare, dataTypeC, C.NumDimensions(), C.GetDims().data(), format);

    CANN_PREPARE_INPUTBUFFER(prepare, const_cast<void*>(matrixA), A.Size() * aclDataTypeSize(dataTypeA));
    CANN_PREPARE_INPUTBUFFER(prepare, const_cast<void*>(matrixB), B.Size() * aclDataTypeSize(dataTypeB));
    CANN_PREPARE_INPUTBUFFER(prepare, const_cast<void*>(matrixC), C.Size() * aclDataTypeSize(dataTypeC));
    CANN_PREPARE_INPUTBUFFER(prepare, const_cast<void*>(alpha), aclDataTypeSize(dataTypeC));
    CANN_PREPARE_INPUTBUFFER(prepare, const_cast<void*>(beta), aclDataTypeSize(dataTypeC));
    CANN_PREPARE_OUTPUTBUFFER(prepare, const_cast<void*>(matrixC), C.Size() * aclDataTypeSize(dataTypeC));
  }
  ORT_CATCH(const std::exception& e) {
    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, e.what());
  }

  CANN_RETURN_IF_ERROR(aclopCompileAndExecute("GEMM",
                                              prepare.inputDesc_.size(),
                                              prepare.inputDesc_.data(),
                                              prepare.inputBuffers_.data(),
                                              prepare.outputDesc_.size(),
                                              prepare.outputDesc_.data(),
                                              prepare.outputBuffers_.data(),
                                              prepare.opAttr_,
                                              ACL_ENGINE_SYS,
                                              ACL_COMPILE_SYS,
                                              NULL,
                                              stream));

  return Status::OK();
}

bool FileExist(const std::string& file_name) {
  return (access(file_name.c_str(), F_OK) != -1);
}

void GenerateHashValue(const std::string string, HashValue& hash_value) {
  uint32_t hash[4] = {0, 0, 0, 0};
  MurmurHash3::x86_128(string.data(), gsl::narrow_cast<int32_t>(string.size()), hash[0], &hash);
  hash_value = hash[0] | (uint64_t(hash[1]) << 32);
}

Status ComputeOutputShape(const std::string& node_name, const TensorShape& lhs_shape,
                          const TensorShape& rhs_shape, TensorShape& out_shape) {
  size_t lhs_rank = lhs_shape.NumDimensions();
  size_t rhs_rank = rhs_shape.NumDimensions();
  size_t out_rank = std::max(lhs_rank, rhs_rank);

  std::vector<int64_t> output_dims(out_rank, 0);
  for (size_t i = 0; i < out_rank; ++i) {
    int64_t lhs_dim = 1;
    if (i < lhs_rank)
      lhs_dim = lhs_shape[lhs_rank - 1 - i];
    int64_t rhs_dim = 1;
    if (i < rhs_rank)
      rhs_dim = rhs_shape[rhs_rank - 1 - i];
    int64_t max = std::max(lhs_dim, rhs_dim);
    int64_t min = std::min(lhs_dim, rhs_dim);
    int64_t out_dim = (min == 0 ? min : max);  // special case a dim value of 0.
    if (lhs_dim != out_dim && lhs_dim != 1)
      return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, node_name, ": left operand cannot broadcast on dim ", lhs_rank - 1 - i,
                             " LeftShape: ", lhs_shape.ToString(), ", RightShape: ", rhs_shape.ToString());
    if (rhs_dim != out_dim && rhs_dim != 1)
      return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, node_name, ": right operand cannot broadcast on dim ", rhs_rank - 1 - i,
                             " LeftShape: ", lhs_shape.ToString(), ", RightShape: ", rhs_shape.ToString());
    output_dims[out_rank - 1 - i] = out_dim;
  }
  out_shape = TensorShape(output_dims);
  return Status::OK();
}

}  // namespace cann
}  // namespace onnxruntime
