//  Copyright (c) 2021 Intel Corporation
//
//  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 "common.hpp"

namespace executor {

unordered_map<string, int> type2bytes = {{"fp32", sizeof(float)},
                                          {"int8", sizeof(char)},
                                          {"int32", sizeof(int)},
                                          {"u8", sizeof(unsigned char)},
                                          {"s8", sizeof(char)},
                                          {"s32", sizeof(int)}};

void GlobalInit(int* pargc, char*** pargv) {
  // Google flags.
  ::gflags::ParseCommandLineFlags(pargc, pargv, true);
  // Google logging.
  ::google::InitGoogleLogging(*(pargv)[0]);
  FLAGS_logtostderr = 1;
  // Provide a backtrace on segfault.
  ::google::InstallFailureSignalHandler();
}
// read weight file to data
/*
Args:
    root is the model const tensors like weights, bias .bin file path.
    type is for setting the tensor datatype, like 'float32'
    shape is the data shape
    location gives the info of tensor data location in .bin file
      location[0] is the start idx when sotre the data bytes
      location[1] is the data bytes length
Return:
    void* ptr, points a consecutive memory that sotres the data
*/
void* read_file_to_type(const string& root, const string& type,
                        const vector<int64_t>& shape, const vector<int64_t>& location) {
  int b = type2bytes[type];
  if (b == 0) {
    LOG(INFO) << type << " not implemented yet...";
  }

  int64_t size = Product(shape);
  // int used_count_ = 1;
  // from file tensor will directly malloc memory
  void* p = reinterpret_cast<void*>(malloc(size * b));

  std::ifstream inFile(root, std::ios::in | std::ios::binary);
  if (inFile) {
    inFile.seekg(location[0], std::ios::beg);
    inFile.read(reinterpret_cast<char*>(p), location[1]);
    inFile.close();
  } else {
    std::memcpy(p, &root[location[0]], location[1]);
  }
  // for (int i=0; i<5; i++){
  //     cout << *((float*)p+i) << endl;
  //
  return p;
}

void InitVector(float* v, int buffer_size) {
  std::mt19937 gen;
  std::uniform_real_distribution<float> u(-10, 10);
  for (int i = 0; i < buffer_size; ++i) {
    v[i] = u(gen);
  }
}

float Time(string state) {
  static std::chrono::milliseconds millis_start = std::chrono::milliseconds();
  static std::chrono::milliseconds millisec_end = std::chrono::milliseconds();
  std::chrono::system_clock::time_point now_time = std::chrono::system_clock::now();
  std::chrono::nanoseconds nanos = now_time.time_since_epoch();
  if (state == "start") {
    millis_start = std::chrono::duration_cast<std::chrono::milliseconds>(nanos);
    return 0.;
  } else if (state == "end") {
    millisec_end = std::chrono::duration_cast<std::chrono::milliseconds>(nanos);
    return millisec_end.count() - millis_start.count();
  } else {
    LOG(FATAL) << "not supported state for time, only start and end...";
  }
}

int64_t Product(const vector<int64_t>& shape) {
  return std::accumulate(shape.begin(), shape.end(), (int64_t)1,
                        std::multiplies<int64_t>());
}

// Get the shapes vector with the absolute perm. Default or empty perm is (0, 1, 2, 3, ...).
// e.g.: shape_before = (64, 384, 16, 64), perm = (0, 2, 1, 3), return (64, 16, 384, 64)
vector<int64_t> GetShapes(const vector<int64_t>& origin_shape,
                          const vector<int64_t>& absolute_perm) {
  if (absolute_perm.empty()) {
    return origin_shape;
  }
  int shape_size = origin_shape.size();
  vector<int64_t> transed_shape(shape_size, 0);
  #pragma omp parallel for
  for (int i = 0; i < shape_size; ++i) {
    int trans_axis_id = absolute_perm[i];
    transed_shape[i] = origin_shape[trans_axis_id];
  }
  return transed_shape;
}

// Get the strides vector with the absolute perm. Default or empty perm is (0, 1, 2, 3, ...).
// Tensor's each stride is the product of all higher dimensions
// Stride[0] = Shape(1)*Shape(2)*...*Shape(n).
// e.g.: axis = (0, 1, 2, 3), shape = (64, 16, 384, 64), return stride = (16*384*64, 384*64, 64, 1)
vector<int64_t> GetStrides(const vector<int64_t>& origin_shape,
                           const vector<int64_t>& absolute_perm) {
  int shape_size = origin_shape.size();
  vector<int64_t> origin_strides(shape_size, 1);
  if (shape_size >= 2) {
    for (int i = shape_size - 2; i >= 0; --i) {
      origin_strides[i] = origin_shape[i + 1] * origin_strides[i + 1];
    }
  }
  return GetShapes(origin_strides, absolute_perm);
}

template <typename T>
bool CompareData(const void* buf1, int64_t elem_num1,
                const void* buf2, int64_t elem_num2, float eps) {
  if (buf1 == buf2) {
    return false;
  }
  if (elem_num1 != elem_num2) {
    return false;
  }
  const auto buf1_data = static_cast<const T*>(buf1);
  const auto buf2_data = static_cast<const T*>(buf2);
  for (int64_t i = 0; i < elem_num1; ++i) {
    auto err = fabs(buf1_data[i] - buf2_data[i]);
    if (err > eps) {
      return false;
    }
  }
  return true;
}
template bool CompareData<float>(const void* buf1, int64_t elem_num1,
                                const void* buf2, int64_t elem_num2, float eps);

vector<float> GetScales(const void* mins, const void* maxs,
                        const int64_t size, const string& dtype) {
  const float* mins_p = static_cast<const float*>(mins);
  const float* maxs_p = static_cast<const float*>(maxs);

  vector<float> scales;
  if (dtype == "u8") {
    for (int i = 0; i < size; i++) {
      float max_sub_min = maxs_p[i]-mins_p[i];
      max_sub_min = max_sub_min == 0.f ? FLT_MIN : max_sub_min;
      scales.emplace_back(255.f/max_sub_min);
    }
  } else if (dtype == "s8") {
    for (int i = 0; i < size; i++) {
      float abs_max = max(abs(maxs_p[i]), abs(mins_p[i]));
      abs_max = abs_max == 0.f ? FLT_MIN : abs_max;
      scales.emplace_back(127.f / abs_max);
    }
  } else if (dtype == "fp32") {
    for (int i = 0; i < size; i++) {
      scales.emplace_back(1.f);
    }
  } else {
    LOG(ERROR) << "Can't suppport dst_dtype: " << dtype << " now!";
  }
  return scales;
}

vector<float> GetRescales(const vector<float>& src0_scales,
  const vector<float>& src1_scales, const vector<float>& dst_scales,
  const string& dst_dtype, const bool append_eltwise) {
  vector<float> rescales;
  if (dst_dtype == "fp32") {
    for (int i = 0; i < src1_scales.size(); i++) {
      rescales.emplace_back(1. / (src0_scales[0] * src1_scales[i]));
    }
  } else if (dst_dtype == "s8" && !dst_scales.empty()) {
    for (int i = 0; i < src1_scales.size(); i++) {
      auto rescale = append_eltwise ? 1. / (src0_scales[0] * src1_scales[i]) :
                                      dst_scales[0] / (src0_scales[0] * src1_scales[i]);
      rescales.emplace_back(rescale);
    }
  } else if (dst_dtype == "u8" && !dst_scales.empty()) {
    for (int i = 0; i < src1_scales.size(); i++) {
      auto rescale = append_eltwise ? 1. / (src0_scales[0] * src1_scales[i]) :
                                      dst_scales[0] / (src0_scales[0] * src1_scales[i]);
      rescales.emplace_back(rescale);
    }
  } else {
    LOG(ERROR) << "Can't suppport dst_dtype: " << dst_dtype << " now!";
  }
  return rescales;
}

vector<int> GetZeroPoints(const void* mins, const vector<float>& scales, const string& dtype) {
  const float* mins_p = static_cast<const float*>(mins);
  vector<int> zerops;
  if (dtype == "u8") {
    for (int i = 0; i < scales.size(); i++)
        zerops.emplace_back(nearbyint(-mins_p[i] * scales[i]));
  } else if (dtype == "s8") {
    for (int i = 0; i < scales.size(); i++)
      zerops.emplace_back(nearbyint(-128 - mins_p[i] * scales[i]));
  } else {
    LOG(ERROR) << "Can't suppport dtype: " << dtype << " now!";
  }
  return zerops;
}

void AddZeroPoints(const int size, const string& dtype, const float* src_data,
                   const float* range_mins, const vector<float>& scales, float* dst_data) {
  if (dtype == "u8") {
    #pragma omp parallel for
    for (int i = 0; i < size; i++) {
      dst_data[i] = src_data[i] - range_mins[0];
    }
  } else if (dtype == "s8") {
    #pragma omp parallel for
    for (int i = 0; i < size; i++) {
      dst_data[i] = src_data[i] -128/scales[0] - range_mins[0];
    }
  } else {
    LOG(ERROR) << "Can't suppport dst_dtype: " << dtype << " now!";
  }
  return;
}

#if __AVX512F__
void Quantize_avx512(const int size, const string& dtype, const void* src_data,
  const float* range_mins, const vector<float>& scales, void* dst_data) {
    const float* src_data_ = static_cast<const float*>(src_data);
    __m512 _min_with_scale_u8 = _mm512_set1_ps(range_mins[0] * scales[0]);
    __m512 _min_with_scale_s8 = _mm512_set1_ps(range_mins[0] * scales[0] + 128);
    __m512 _scale = _mm512_set1_ps(scales[0]);
    __m512i zero = _mm512_setzero_epi32();

    int avx512_loop_len = size >> 4;

    if (dtype == "u8") {
      unsigned char* dst_data_ = static_cast<unsigned char*>(dst_data);
      #pragma omp parallel for
      for (int i = 0; i < avx512_loop_len; ++i) {
        __m512 _src_data = _mm512_loadu_ps(src_data_ + (i << 4));
        __m512 data = _mm512_fmsub_ps(_src_data, _scale, _min_with_scale_u8);
        __m512i data_x32 = _mm512_cvt_roundps_epi32(
          data, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
        data_x32 = _mm512_max_epi32(data_x32, zero);
        _mm_storeu_si128(reinterpret_cast<__m128i*>(
          dst_data_ + (i << 4)), _mm512_cvtusepi32_epi8(data_x32));
      }

      #pragma omp parallel for
      for (int i = (avx512_loop_len << 4); i < size; i++) {
        int32_t data = nearbyint((src_data_[i] - range_mins[0]) * scales[0]);
        data = data < 0 ? 0 : data;
        data = data > 255 ? 255 : data;
        dst_data_[i] = static_cast<unsigned char>(data);
      }
    } else if (dtype == "s8") {
      char* dst_data_ = static_cast<char*>(dst_data);
      #pragma omp parallel for
      for (int i = 0; i < avx512_loop_len; ++i) {
        __m512 _src_data = _mm512_loadu_ps(src_data_ + (i << 4));
        __m512 data = _mm512_fmsub_ps(_src_data, _scale, _min_with_scale_s8);
        __m512i data_x32 = _mm512_cvt_roundps_epi32(
          data, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
        _mm_storeu_si128(reinterpret_cast<__m128i*>(
          dst_data_ + (i << 4)), _mm512_cvtsepi32_epi8(data_x32));
      }
      #pragma omp parallel for
      for (int i = (avx512_loop_len << 4); i < size; i++) {
        int32_t data = nearbyint((src_data_[i] - range_mins[0]) * scales[0] - 128);
        data = data < -128 ? -128 : data;
        data = data > 127 ? 127 : data;
        dst_data_[i] = static_cast<char>(data);
      }
    } else {
      LOG(ERROR) << "Can't suppport dst_dtype: " << dtype << " now!";
    }
    return;
  }

#else
void Quantize(const int size, const string& dtype, const void* src_data,
  const float* range_mins, const vector<float>& scales, void* dst_data) {
  const float* src_data_ = static_cast<const float*>(src_data);
  if (dtype == "u8") {
    unsigned char* dst_data_ = static_cast<unsigned char*>(dst_data);
    #pragma omp parallel for
    for (int i = 0; i < size; i++) {
      int32_t data = nearbyint((src_data_[i] - range_mins[0]) * scales[0]);
      data = data < 0 ? 0 : data;
      data = data > 255 ? 255 : data;
      dst_data_[i] = static_cast<unsigned char>(data);
    }
  } else if (dtype == "s8") {
    char* dst_data_ = static_cast<char*>(dst_data);
    #pragma omp parallel for
    for (int i = 0; i < size; i++) {
      int32_t data = nearbyint((src_data_[i] - range_mins[0]) * scales[0] - 128);
      data = data < -128 ? -128 : data;
      data = data > 127 ? 127 : data;
      dst_data_[i] = static_cast<char>(data);
    }
  } else {
    LOG(ERROR) << "Can't suppport dst_dtype: " << dtype << " now!";
  }
  return;
}
#endif

// Transpose from A to default B.
// e.g.: transpose from {2, 0, 1} to default {0, 1, 2} is {1, 2, 0}
vector<int64_t> ReversePerm(const vector<int64_t>& perm_to) {
  if (perm_to.empty()) return {};
  int dsize = perm_to.size();
  vector<int64_t> perm_from(dsize, 0);
  std::iota(perm_from.begin(), perm_from.end(), 0);
  if (perm_to.empty()) {
    return perm_from;
  }
  #pragma omp parallel for
  for (int i = 0; i < dsize; ++i) {
    int index = perm_to[i];
    perm_from[index] = i;
  }
  return perm_from;
}

template <typename T>
T StringToNum(const string& str) {
  std::istringstream iss(str);
  T num;
  iss >> num;
  return num;
}

template float StringToNum<float>(const string& str);
template int64_t StringToNum<int64_t>(const string& str);

template<typename T>
void PrintToFile(const T* data, const std::string& name, size_t size) {
  // print output file
  auto pos = name.rfind("/");
  string output_file = (pos != string::npos ? name.substr(pos + 1) : name) + ".txt";
  std::ofstream output_data(output_file);
  for (size_t i = 0; i < size; ++i) {
    output_data << static_cast<float>(data[i]) << "\n";
  }
  output_data.close();
}
template void PrintToFile<float>(const float* data, const std::string& name, size_t size);
template void PrintToFile<unsigned char>(
  const unsigned char* data, const std::string& name, size_t size);
template void PrintToFile<char>(const char* data, const std::string& name, size_t size);
template void PrintToFile<int32_t>(const int32_t* data, const std::string& name, size_t size);
template void PrintToFile<int64_t>(const int64_t* data, const std::string& name, size_t size);

template <typename T>
void StringSplit(vector<T>* split_list, const string& str_list, const string& split_op) {
  std::string::size_type pos1 = 0;
  std::string::size_type pos2 = str_list.find(split_op);
  while (std::string::npos != pos2) {
    T element = StringToNum<T>(str_list.substr(pos1, pos2));
    split_list->push_back(element);
    pos1 = pos2 + split_op.size();
    pos2 = str_list.find(split_op, pos1);
  }
  if (pos1 != str_list.size()) {
      T element = StringToNum<T>(str_list.substr(pos1));
      split_list->push_back(element);
  }
}
template void StringSplit<float>(vector<float>* split_list, const string& string_list,
                                 const string& split_op);
template void StringSplit<int64_t>(vector<int64_t>* split_list, const string& string_list,
                                   const string& split_op);
template void StringSplit<char>(vector<char>* split_list, const string& string_list,
                                const string& split_op);
template void StringSplit<unsigned char>(vector<unsigned char>* split_list,
                                         const string& string_list,
                                         const string& split_op);

}  // namespace executor
