/*
    Copyright (c) 2013, Taiga Nomi and the respective contributors
    All rights reserved.

    Use of this source code is governed by a BSD-style license that can be found
    in the LICENSE file.
*/
#pragma once

#include <algorithm>
#include <limits>
#include <vector>

namespace tiny_dnn {
namespace core {
namespace kernels {

template <class T>
T highest() {
  return (std::numeric_limits<T>::max)();
}

template <class T>
T lowest() {
  return std::numeric_limits<T>::is_integer
           ? (std::numeric_limits<T>::min)()
           : (-(std::numeric_limits<T>::max)());
}

// We have to be able to detect and handle overflows in int32, so this function
// uses doubles and int64's to make sure we have enough room.
template <class T>
int64_t float_to_quantized_unclamped(float_t input,
                                     float_t range_min,
                                     float_t range_max) {
  if (range_min == range_max) {
    return 0;
  }
  const int number_of_bits      = sizeof(T) * 8;
  const int64_t number_of_steps = static_cast<int64_t>(1) << number_of_bits;
  const double range_adjust     = (number_of_steps / (number_of_steps - 1.0));
  const double range            = ((range_max - range_min) * range_adjust);
  const double range_scale      = (number_of_steps / range);
  int64_t quantized = static_cast<int64_t>(round(input * range_scale) -
                                           round(range_min * range_scale));
  const int64_t lowest_quantized = static_cast<int64_t>(lowest<T>());
  quantized += lowest_quantized;
  return quantized;
}

inline int32_t int64_to_int32(int64_t src) {
  assert(src <= std::numeric_limits<int32_t>::max() &&
         src >= std::numeric_limits<int32_t>::min());
  return static_cast<int32_t>(src);
}

// This converts the float into the final quantized type, clamping/saturating
// any over or underflows.
template <class T>
T float_to_quantized(float_t input, float_t range_min, float_t range_max) {
  int64_t quantized =
    float_to_quantized_unclamped<T>(input, range_min, range_max);
  const int64_t lowest_quantized  = static_cast<int64_t>(lowest<T>());
  const int64_t highest_quantized = static_cast<int64_t>(highest<T>());
  quantized = std::max<int64_t>(quantized, lowest_quantized);
  quantized = std::min<int64_t>(quantized, highest_quantized);
  return static_cast<T>(static_cast<int32_t>(quantized));
}

template <class T>
float quantized_to_float(T input, float_t range_min, float_t range_max) {
  if (range_min == range_max) {
    return range_min;
  }
  const int number_of_bits       = sizeof(T) * 8;
  const int64_t number_of_steps  = static_cast<int64_t>(1) << number_of_bits;
  const double range_adjust      = (number_of_steps / (number_of_steps - 1.0));
  const double range             = ((range_max - range_min) * range_adjust);
  const double range_scale       = (range / number_of_steps);
  const int64_t lowest_quantized = static_cast<int64_t>(lowest<T>());
  const double offset_input = static_cast<double>(input) - lowest_quantized;
  const double result       = range_min + (offset_input * range_scale);
  return static_cast<float_t>(result);
}

template <class T>
float float_for_one_quantized_level(float_t range_min, float_t range_max) {
  const int64_t highest_ = static_cast<int64_t>(highest<T>());
  const int64_t lowest_  = static_cast<int64_t>(lowest<T>());
  const float float_for_one_quantized_level =
    (range_max - range_min) / (highest_ - lowest_);
  return float_for_one_quantized_level;
}

template <class T1, class T2, class T3>
void quantization_range_for_multiplication(float_t min_a,
                                           float_t max_a,
                                           float_t min_b,
                                           float_t max_b,
                                           float_t *min_c,
                                           float_t *max_c) {
  const float_t a_float_for_one_quant_level =
    float_for_one_quantized_level<T1>(min_a, max_a);
  const float_t b_float_for_one_quant_level =
    float_for_one_quantized_level<T2>(min_b, max_b);

  const int64_t c_highest = static_cast<int64_t>(highest<T3>());
  const int64_t c_lowest  = static_cast<int64_t>(lowest<T3>());
  const float c_float_for_one_quant_level =
    a_float_for_one_quant_level * b_float_for_one_quant_level;

  *min_c = c_float_for_one_quant_level * c_lowest;
  *max_c = c_float_for_one_quant_level * c_highest;
}

template <class T1, class T2>
inline T2 requantize_in_new_range(T1 input,
                                  float_t min_input,
                                  float_t max_input,
                                  float_t min_new,
                                  float_t max_new) {
  const float_t input_float =
    quantized_to_float<T1>(input, min_input, max_input);
  return float_to_quantized<T2>(input_float, min_new, max_new);
}

template <class T1, class T2>
inline void requantize_many_in_new_range(T1 *input,
                                         size_t count,
                                         float_t min_input,
                                         float_t max_input,
                                         float_t min_output,
                                         float_t max_output,
                                         T2 *output) {
  for (size_t index = 0; index < count; ++index) {
    const float_t input_float =
      quantized_to_float<T1>(input[index], min_input, max_input);
    output[index] = float_to_quantized<T2>(input_float, min_output, max_output);
  }
}

// Because converting 32-bit accumulated results down to eight bit is a common
// case, we have a specialized code path to handle it as efficiently as
// possible using only fixed-point math for the inner loop.
template <>
inline void requantize_many_in_new_range<int32_t, uint8_t>(int32_t *input,
                                                           size_t count,
                                                           float_t min_input,
                                                           float_t max_input,
                                                           float_t min_output,
                                                           float_t max_output,
                                                           uint8_t *output) {
  // Initially we calculate all the constants we need once, before we go into
  // the inner loop.
  const int fp_shift             = 16;
  const float input_range        = max_input - min_input;
  const float output_range       = max_output - min_output;
  const float recip_output_range = (255.0f / output_range);
  const int64_t recip_output_range_fp =
    static_cast<int64_t>(recip_output_range * (1 << fp_shift));
  const int64_t range_scale_fp =
    static_cast<int64_t>(255.0f * (1 << fp_shift) * input_range / output_range);
  const int64_t input_offset_fp = static_cast<int64_t>(
    (min_input * recip_output_range_fp) + (range_scale_fp >> 1));
  const int64_t output_offset_fp =
    static_cast<int64_t>(round((min_output * 255.0f) / output_range));
  const int64_t rounding_delta = 1 << (fp_shift - 1);
  // Inside this loop we just do minimal adds, multiplies, and shifts, in a
  // way
  // that could be easily adapted for a SIMD implementation. It should also be
  // possible to perform all the calculations in 32-bit rather than 64, but
  // that's not been implemented yet.
  for (size_t index = 0; index < count; ++index) {
    const int64_t input_value = static_cast<int64_t>(input[index]);
    const int64_t fp_value =
      ((input_value * range_scale_fp) >> 32) + input_offset_fp;
    const int64_t round_intermediate =
      ((fp_value >= 0) ? (fp_value + rounding_delta)
                       : (fp_value - rounding_delta)) >>
      fp_shift;
    int64_t quantized_int64 = (round_intermediate - output_offset_fp);
    quantized_int64         = std::max<int64_t>(quantized_int64, 0LL);
    quantized_int64         = std::min<int64_t>(quantized_int64, 255LL);
    output[index] = static_cast<uint8_t>(static_cast<int32_t>(quantized_int64));
  }
}

// REQUIRES: 'result->NumElements() == input.NumElements()'
template <class T>
void float_tensor_to_quantized_in_place(const vec_t &input,
                                        float_t min,
                                        float_t max,
                                        std::vector<T> *result) {
  const size_t data_size = input.size();
  for (size_t i = 0; i < data_size; ++i) {
    (*result)[i] = float_to_quantized<T>(input[i], min, max);
  }
}

template <class T>
std::vector<T> float_tensor_to_quantized(const vec_t &input,
                                         float_t min,
                                         float_t max) {
  std::vector<T> result(input.size(), static_cast<T>(0));
  float_tensor_to_quantized_in_place<T>(input, min, max, &result);
  return result;
}

// REQUIRES: 'result->NumElements() == input.NumElements()'
template <class T>
void quantized_tensor_to_float_in_place(const std::vector<T> &input,
                                        float_t min,
                                        float_t max,
                                        vec_t *result) {
  const size_t data_size = input.size();
  for (size_t i = 0; i < data_size; ++i) {
    (*result)[i] = quantized_to_float<T>(input[i], min, max);
  }
}

template <class T>
vec_t quantized_tensor_to_float(const std::vector<T> &input,
                                float_t min,
                                float_t max) {
  vec_t result(input.size(), static_cast<float_t>(0));
  quantized_tensor_to_float_in_place<T>(input, min, max, &result);
  return result;
}

template <class T1, class T2>
void quantize_down_and_shrink_range(std::vector<T1> &input,
                                    float_t min_input,
                                    float_t max_input,
                                    float_t *min_new,
                                    float_t *max_new,
                                    std::vector<T2> *output) {
  const int32_t input_lowest_quantized  = static_cast<int32_t>(lowest<T1>());
  const int32_t input_highest_quantized = static_cast<int32_t>(highest<T1>());
  T1 actual_min_quantized               = input_highest_quantized;
  T1 actual_max_quantized               = input_lowest_quantized;
  for (size_t i = 0; i < input.size(); ++i) {
    const T1 value       = input[i];
    actual_min_quantized = std::min(actual_min_quantized, value);
    actual_max_quantized = std::max(actual_max_quantized, value);
  }
  // We want to make sure that the minimum is no larger than zero, so that the
  // convolution operation can run efficiently.
  *min_new = std::min(
    0.0f, quantized_to_float(actual_min_quantized, min_input, max_input));
  *max_new = quantized_to_float(actual_max_quantized, min_input, max_input);
  requantize_many_in_new_range<int32_t, uint8_t>(&input[0], input.size(),
                                                 min_input, max_input, *min_new,
                                                 *max_new, &(*output)[0]);
}

}  // namespace kernels
}  // namespace core
}  // namespace tiny_dnn
