// Copyright 2018 The MACE Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "mace/core/runtime/apu/apu_wrapper.h"

#include <algorithm>

#include "mace/core/quantize.h"

namespace mace {

ApuWrapper::ApuWrapper(Device *device)
    : quantize_util_uint8_(&device->cpu_runtime()->thread_pool()),
      quantize_util_int16_(&device->cpu_runtime()->thread_pool()) {
}

apu_data_type ApuWrapper::MapToApuDataType(DataType mace_type) {
  switch (mace_type) {
    case DT_FLOAT:
      return APU_DATA_TYPE_FLOAT;
    case DT_INT32:
      return APU_DATA_TYPE_INT32;
    case DT_HALF:
      return APU_DATA_TYPE_HALF;
    case DT_FLOAT16:
      return APU_DATA_TYPE_HALF;
    case DT_UINT8:
      return APU_DATA_TYPE_UINT8;
    case DT_INT16:
      return APU_DATA_TYPE_INT16;
    default:
      MACE_CHECK(false, "unsupport mace data type");
      break;
  }
  return APU_DATA_TYPE_UNDEFINED;
}

apu_pooling_mode ApuWrapper::MapToApuPoolingMode(int mace_mode) {
  switch (mace_mode) {
    case 1:
      return APU_POOLING_AVG;
    case 2:
      return APU_POOLING_MAX;
    default:
      MACE_CHECK(false, "unsupport mace pooling mode");
      break;
  }
  return APU_POOLING_UNDEFINED;
}

apu_eltwise_mode ApuWrapper::MapToApuEltwiseMode(int mace_mode) {
  switch (mace_mode) {
    case 0:
      return APU_ELTWISE_ADD;
    case 1:
      return APU_ELTWISE_SUB;
    case 2:
      return APU_ELTWISE_MUL;
    case 4:
      return APU_ELTWISE_MIN;
    case 5:
      return APU_ELTWISE_MAX;
    default:
      MACE_CHECK(false, "unsupport mace eltwise mode");
      break;
  }
  return APU_ELTWISE_UNDEFINED;
}

bool ApuWrapper::Init(const NetDef &net_def, unsigned const char *model_data,
                      const char *file_name, bool load, bool store) {
  frontend = new ApuFrontend();

  MACE_CHECK(!(load & store),
            "Should not load and store the model simultaneously.");

  // parse model argument
  int const_data_num = 0;
  int apu_data_type = -1;
  for (auto arg : net_def.arg()) {
    if (arg.name().compare("const_data_num") == 0) {
      const_data_num = arg.i();
    } else if (arg.name().compare("apu_data_type") == 0) {
      apu_data_type = arg.i();
    }
  }
  // input tensors
  std::vector<apu_tensor> input_tensors;
  for (auto input_info : net_def.input_info()) {
    apu_tensor tensor;
    tensor.tensor_id = input_info.node_id();
    tensor.tensor_type = APU_TENSOR_MODEL_INPUT;
    tensor.data_type = MapToApuDataType(static_cast<DataType>(apu_data_type));
    tensor.scale = input_info.has_scale() ? input_info.scale() : -1.0f;
    tensor.zero_point = input_info.has_zero_point() ?
                            input_info.zero_point() : 0;
    tensor.dim_size = input_info.dims_size();
    MACE_CHECK(tensor.dim_size <= APU_TENSOR_MAX_DIMS,
               "tensor dimension size not supported");
    ApuTensorInfo info;
    info.name = input_info.name();
    info.size = 1;
    info.data_type = tensor.data_type;
    int byte_per_element = GetByteNum(tensor.data_type);
    for (auto i = 0 ; i < tensor.dim_size ; i++) {
      tensor.dims[i] = input_info.dims(i);
      info.size *= input_info.dims(i);
      info.shape.push_back(input_info.dims(i));
    }
    info.buf
    = std::shared_ptr<uint8_t>(new uint8_t[info.size * byte_per_element],
                               std::default_delete<uint8_t[]>());
    info.scale = tensor.scale;
    info.zero_point = tensor.zero_point;
    input_infos.push_back(info);
    tensor.data_buf = info.buf.get();
    input_tensors.push_back(tensor);
  }
  // output tensors
  std::vector<apu_tensor> output_tensors;
  for (auto output_info : net_def.output_info()) {
    apu_tensor tensor;
    tensor.tensor_id = output_info.node_id();
    tensor.tensor_type = APU_TENSOR_MODEL_OUTPUT;
    tensor.data_type = MapToApuDataType(static_cast<DataType>(apu_data_type));
    tensor.dim_size = output_info.dims_size();
    ApuTensorInfo info;
    info.name = output_info.name();
    info.size = 1;
    info.data_type = tensor.data_type;
    int byte_per_element = GetByteNum(tensor.data_type);
    for (auto i = 0 ; i < tensor.dim_size ; i++) {
      tensor.dims[i] = output_info.dims(i);
      info.size *= output_info.dims(i);
      info.shape.push_back(output_info.dims(i));
    }
    info.buf =
    std::shared_ptr<uint8_t>(new uint8_t[info.size * byte_per_element],
                             std::default_delete<uint8_t[]>());
    for (auto op_def : net_def.op()) {
      if (output_info.name() == op_def.output(0)) {
        if (info.data_type == static_cast<int>(APU_DATA_TYPE_UINT8) ||
            info.data_type == static_cast<int>(APU_DATA_TYPE_INT16)) {
          info.scale = op_def.quantize_info(0).scale();
          info.zero_point = op_def.quantize_info(0).zero_point();
        } else {
          info.scale = 0.0;
          info.zero_point = 0;
        }
      }
    }
    output_infos.push_back(info);
    tensor.data_buf = info.buf.get();
    output_tensors.push_back(tensor);
  }
  // const tensors
  std::vector<apu_tensor> const_tensors;
  // operators
  std::vector<apu_operator> ops;
  std::vector<std::vector<int>> cached_op_inputs;
  if (!load) {
    // const tensors
    for (auto const_tensor : net_def.tensors()) {
      apu_tensor tensor;
      tensor.tensor_id = const_tensor.node_id();
      tensor.tensor_type = (tensor.tensor_id < const_data_num) ?
                               APU_TENSOR_CONST_DATA :
                               APU_TENSOR_CONST_ARGUMENT;
      tensor.data_type = MapToApuDataType(const_tensor.data_type());
      tensor.scale = const_tensor.has_scale() ? const_tensor.scale() : 0.0f;
      tensor.zero_point = const_tensor.has_zero_point() ?
                              const_tensor.zero_point() : 0;
      tensor.dim_size = const_tensor.dims_size();
      MACE_CHECK(tensor.dim_size <= APU_TENSOR_MAX_DIMS,
                 "tensor dimension size not supported");
      for (auto i = 0 ; i < tensor.dim_size ; i++) {
        tensor.dims[i] = const_tensor.dims(i);
      }
      tensor.data_buf =
          const_cast<unsigned char*>(model_data + const_tensor.offset());
      const_tensors.push_back(tensor);
    }
    // operators
    for (auto op_def : net_def.op()) {
      apu_operator op;
      strncpy(op.type, op_def.type().c_str(), APU_OP_TYPE_MAX_SIZE);
      op.input_size = op_def.node_input_size();
      std::vector<int> input_ids;
      for (auto i = 0 ; i < op.input_size ; i++) {
        input_ids.push_back(op_def.node_input(i).node_id());
      }
      cached_op_inputs.push_back(input_ids);
      op.input_ids = cached_op_inputs.back().data();
      op.output.tensor_id = op_def.node_id();
      op.output.tensor_type = APU_TENSOR_OP_OUTPUT;
      op.output.data_type = MapToApuDataType(op_def.output_type(0));
      if (op.output.data_type == APU_DATA_TYPE_UINT8 ||
          op.output.data_type == APU_DATA_TYPE_INT16) {
        op.output.scale = op_def.quantize_info(0).scale();
        op.output.zero_point = op_def.quantize_info(0).zero_point();
      } else {
        op.output.scale = 0.0f;
        op.output.zero_point = 0;
      }
      op.output.dim_size = op_def.output_shape(0).dims_size();
      MACE_CHECK(op.output.dim_size <= APU_TENSOR_MAX_DIMS,
                 "tensor dimension size not supported");
      for (auto i = 0 ; i < op.output.dim_size ; i++) {
        op.output.dims[i] = op_def.output_shape(0).dims(i);
      }
      op.output.data_buf = nullptr;
      // get op mode and activation mode
      bool is_pooling = (strcmp(op.type, "Pooling") == 0);
      bool is_eltwise = (strcmp(op.type, "Eltwise") == 0);
      std::string activation;
      float max_limit = 0.0f;
      for (auto arg : op_def.arg()) {
        if (arg.name().compare("activation") == 0) {
          activation = arg.s();
        }
        if (arg.name().compare("max_limit") == 0) {
          max_limit = arg.f();
        }
        if (is_pooling && arg.name().compare("pooling_type") == 0) {
          op.op_mode = static_cast<int>(MapToApuPoolingMode(arg.i()));
        }
        if (is_eltwise && arg.name().compare("type") == 0) {
          op.op_mode = static_cast<int>(MapToApuEltwiseMode(arg.i()));
        }
      }
      if (activation.compare("RELU") == 0) {
        op.act_mode = APU_ACT_RELU;
      } else if (activation.compare("RELUX") == 0 && max_limit == 6.0) {
        op.act_mode = APU_ACT_RELU6;
      } else if (activation.compare("SIGMOID") == 0) {
        op.act_mode = APU_ACT_SIGMOID;
      } else if (activation.compare("TANH") == 0) {
        op.act_mode = APU_ACT_TANH;
      } else {
        op.act_mode = APU_ACT_NONE;
      }
      ops.push_back(op);
    }
  }
  bool print_model = false;
  bool ret = frontend->InitGraph(
                 const_tensors.size(), const_tensors.data(),
                 input_tensors.size(), input_tensors.data(),
                 output_tensors.size(), output_tensors.data(),
                 ops.size(), ops.data(),
                 print_model, file_name, load, store);
  cached_op_inputs.clear();
  return ret;
}

bool ApuWrapper::Run(const std::map<std::string, Tensor *> &input_tensors,
                     std::map<std::string, Tensor *> *output_tensors) {
  MACE_ASSERT(input_tensors.size() == input_infos.size(), "Wrong inputs num");
  MACE_ASSERT(output_tensors.size() == output_infos.size(),
              "Wrong outputs num");
  // prepare input
  for (int i = 0 ; i < static_cast<int>(input_tensors.size()) ; i++) {
    Tensor* tensor = input_tensors.at(input_infos[i].name);

    // check size
    int element_size = input_infos[i].size;
    int byte_per_element = GetByteNum(input_infos[i].data_type);
    MACE_ASSERT(element_size == static_cast<int>(tensor->size()),
                "Wrong input size");
    // quantize
    if (input_infos[i].data_type == APU_DATA_TYPE_INT16) {
      quantize_util_int16_.QuantizeWithScaleAndZeropoint(
          (const float*)tensor->raw_data(),
          element_size,
          input_infos[i].scale,
          input_infos[i].zero_point,
          reinterpret_cast<int16_t*>(input_infos[i].buf.get()));
    } else if (input_infos[i].data_type == APU_DATA_TYPE_FLOAT) {
      std::memcpy(input_infos[i].buf.get(),
                    (const float*)tensor->raw_data(),
                    element_size * byte_per_element);
    } else {
      quantize_util_uint8_.QuantizeWithScaleAndZeropoint(
          (const float*)tensor->raw_data(),
          element_size,
          input_infos[i].scale,
          input_infos[i].zero_point,
          input_infos[i].buf.get());
    }
  }

  // run model
  bool ret = frontend->RunGraph();
  MACE_CHECK(ret == true, "neuron run model failed");

  // process output
  for (int i = 0 ; i < static_cast<int>(output_tensors->size()) ; i++) {
    Tensor* tensor = output_tensors->at(output_infos[i].name);

    // prepare out buffer
    tensor->SetDtype(DT_FLOAT);
    tensor->Resize(output_infos[i].shape);
    int element_size = output_infos[i].size;
    int byte_per_element = GetByteNum(output_infos[i].data_type);
    MACE_ASSERT(element_size == static_cast<int>(tensor->size()),
                "Wrong output size");
    // dequantize
    if (output_infos[i].data_type == APU_DATA_TYPE_INT16) {
      quantize_util_int16_.Dequantize(
          reinterpret_cast<int16_t*>(output_infos[i].buf.get()),
          element_size,
          output_infos[i].scale,
          output_infos[i].zero_point,
          reinterpret_cast<float*>(tensor->raw_mutable_data()));
    } else if (output_infos[i].data_type == APU_DATA_TYPE_FLOAT) {
        std::memcpy(reinterpret_cast<float*>(tensor->raw_mutable_data()),
                    output_infos[i].buf.get(),
                    element_size * byte_per_element);
    } else {
      quantize_util_uint8_.Dequantize(
          output_infos[i].buf.get(),
          element_size,
          output_infos[i].scale,
          output_infos[i].zero_point,
          reinterpret_cast<float*>(tensor->raw_mutable_data()));
    }
  }

  return true;
}

bool ApuWrapper::Uninit() {
  bool ret = frontend->UninitGraph();
  frontend = nullptr;
  input_infos.clear();
  output_infos.clear();
  return ret;
}

int ApuWrapper::GetByteNum(apu_data_type data_type) {
  int byte_per_element;
  if (data_type == APU_DATA_TYPE_FLOAT || data_type == APU_DATA_TYPE_INT32) {
    byte_per_element = 4;
  } else if (data_type == APU_DATA_TYPE_HALF ||
             data_type == APU_DATA_TYPE_INT16) {
    byte_per_element = 2;
  } else if (data_type == APU_DATA_TYPE_UINT8) {
    byte_per_element = 1;
  } else {
    byte_per_element = 1;
    MACE_CHECK(false, "unsupport data type");
  }
  return byte_per_element;
}

}  // namespace mace
