// Copyright 2017 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "base/containers/span.h"
#include "base/numerics/safe_conversions.h"

// This file is automatically generated using tfNative from a neural network,
// trained by TensorFlow. Please do not edit.

#include "darkmode_classifier.h"

#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <limits>
#include <tuple>
#if USE_EIGEN
#include "third_party/eigen3/Eigen/Core"
#endif
namespace darkmode_tfnative_model {
namespace {

// -----------------------------------------------------------------------------
// OP LIBRARY
// Copied here to make sure that the inferece code always stays in sync with the
// lib that it was generated for.
// -----------------------------------------------------------------------------

// Default to using std::copy and std::fill over memcpy and memset as they
// are usually faster, thanks to the compiler getting stricter alignment
// guarantees.
#ifndef USE_TYPED_MEMSETMEMCPY
#define USE_TYPED_MEMSETMEMCPY 1
#endif
#define USE_EIGEN 0
#ifndef USE_EIGEN
#error Please define USE_EIGEN to either 0 or 1
#endif

// Helper to reinterpret memory as Eigen matrices.
#if USE_EIGEN
template <typename Scalar>
using ConstMatrixMap = typename Eigen::Map<
    const Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>>;
template <typename Scalar>
using ConstRowVectorMap =
    typename Eigen::Map<const Eigen::Matrix<Scalar, Eigen::Dynamic, 1>>;
template <typename Scalar>
using RowVectorMap =
    typename Eigen::Map<Eigen::Matrix<Scalar, Eigen::Dynamic, 1>>;
template <typename Scalar>
using MatrixMap =
    typename Eigen::Map<Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>>;
#endif

#define BENCHMARK_TIMER(...)

// The size of a shape in terms of number of coefficients.
inline int ShapeSize(const int32_t rank, base::span<const int32_t> shape) {
  int size = 1;
  for (int i = 0; i < rank; ++i)
    size *= shape[i];
  return size;
}

template <typename T>
void FullyConnected(base::span<const int32_t> input_shape,
                    base::span<const T> input_values,
                    base::span<const int32_t> weight_shape,
                    base::span<const T> weight_values,
                    base::span<const int32_t> bias_shape,
                    base::span<const T> bias_values,
                    base::span<T> output_values) {
  BENCHMARK_TIMER("FullyConnected");
#if USE_EIGEN
  const auto in =
      ConstMatrixMap<T>(input_values, input_shape[1], input_shape[0]);
  const auto weight =
      ConstMatrixMap<T>(weight_values, weight_shape[1], weight_shape[0]);
  const auto bias = ConstRowVectorMap<T>(bias_values, bias_shape[0]);
  auto result = MatrixMap<T>(output_values, weight_shape[1], input_shape[0]);
  result.noalias() = (weight * in).colwise() + bias;
#else
  const int batch_size = input_shape[0];
  const int num_inputs = weight_shape[0];
  const int num_outputs = weight_shape[1];
  assert(input_shape[1] == num_inputs);
  assert(bias_shape[0] == num_outputs);
  for (int batch = 0; batch < batch_size; ++batch) {
    for (int out_i = 0; out_i < num_outputs; ++out_i) {
      T value = 0;
      for (int in_i = 0; in_i < num_inputs; ++in_i) {
        value += input_values[batch * num_inputs + in_i] *
                 weight_values[in_i * num_outputs + out_i];
      }
      value += bias_values[out_i];
      output_values[batch * num_outputs + out_i] = value;
    }
  }
#endif
}

// -----------------------------------------------------------------------------
// Simple unary ops
// -----------------------------------------------------------------------------

// We use macros instead of template functions with templated functors here
// because it's a lot less verbose and easier for the compiler to optimize.

#if USE_EIGEN

#define SIMPLE_UNARY_OP(OP_NAME, _, EXPR_EIGEN)                            \
  template <typename T>                                                    \
  void OP_NAME(const int32_t rank, base::span<const int32_t> input_shape,  \
               base::span<const T> input_values,                           \
               base::span<T> output_values) {                              \
    BENCHMARK_TIMER(#OP_NAME);                                             \
    const int size = ShapeSize(rank, input_shape);                         \
    auto values = ConstRowVectorMap<T>(input_values.data(), size).array(); \
    auto output = RowVectorMap<T>(output_values.data(), size).array();     \
    output = EXPR_EIGEN;                                                   \
  }

#else

#define SIMPLE_UNARY_OP(OP_NAME, EXPR, _)                                 \
  template <typename T>                                                   \
  void OP_NAME(const int32_t rank, base::span<const int32_t> input_shape, \
               base::span<const T> input_values,                          \
               base::span<T> output_values) {                             \
    BENCHMARK_TIMER(#OP_NAME);                                            \
    const int size = ShapeSize(rank, input_shape);                        \
    for (int i = 0; i < size; ++i) {                                      \
      const T value = input_values[i];                                    \
      output_values[i] = EXPR;                                            \
    }                                                                     \
  }

#endif

// Second macro param is value expression, third entry is Eigen vector
// expression.
SIMPLE_UNARY_OP(Relu, std::max(value, static_cast<T>(0)), values.max(0))

// -----------------------------------------------------------------------------
// CONSTANTS
// Note that for now, endianness of the target machine needs to match that of
// the one training was performed on.
// -----------------------------------------------------------------------------
const int32_t dnn_hiddenlayer_0_weights_part_0_shape[2] = {4, 10};
const union {
  uint8_t bytes[160];
  float values[40];
} dnn_hiddenlayer_0_weights_part_0 = {{
    0xbc, 0x22, 0x0a, 0xbf, 0xb4, 0x46, 0x8c, 0x3f, 0xba, 0x31, 0x34, 0xbe,
    0x4c, 0x65, 0xdb, 0xbe, 0xf0, 0x54, 0x5e, 0xbe, 0xc1, 0x5d, 0xb3, 0x3f,
    0xf4, 0xe6, 0x15, 0xbf, 0x05, 0xc6, 0x34, 0xbf, 0xc0, 0x37, 0x7e, 0xbd,
    0x6c, 0x35, 0x0b, 0xbf, 0xca, 0x53, 0x26, 0xbf, 0x58, 0xb4, 0x87, 0x3f,
    0x37, 0xee, 0x39, 0xbf, 0xda, 0xfa, 0xf9, 0xbe, 0x97, 0xc1, 0x06, 0xbf,
    0xf9, 0x4e, 0x81, 0x3f, 0xb2, 0x44, 0x85, 0xbf, 0x7f, 0x98, 0x7c, 0x3d,
    0x15, 0x26, 0xbc, 0xbe, 0x5c, 0x48, 0x05, 0x3f, 0xc8, 0xaa, 0xa1, 0xbd,
    0x35, 0xb3, 0x43, 0xbe, 0xeb, 0x46, 0x91, 0x3f, 0x80, 0x71, 0xe3, 0x3c,
    0xd1, 0x98, 0x79, 0x3f, 0x3c, 0xd0, 0x0d, 0xbf, 0x1e, 0x02, 0xd3, 0x3e,
    0x5d, 0x4b, 0xa2, 0xbf, 0x68, 0xac, 0xaa, 0xbd, 0xf8, 0xe1, 0x75, 0x3e,
    0x4a, 0x9c, 0x27, 0xbe, 0xf8, 0xae, 0xb2, 0xbe, 0x7f, 0x9d, 0x91, 0x3f,
    0x1e, 0x8b, 0xa8, 0xbe, 0x35, 0x7e, 0xb2, 0x3f, 0xbe, 0x8c, 0xd3, 0xbe,
    0xf9, 0xcd, 0xb5, 0x3f, 0xa1, 0x50, 0xaa, 0x3f, 0xe4, 0x6d, 0xdd, 0xbe,
    0x0d, 0xce, 0xd3, 0xbe,
}};
const int32_t dnn_hiddenlayer_0_biases_part_0_shape[1] = {10};
const union {
  uint8_t bytes[40];
  float values[10];
} dnn_hiddenlayer_0_biases_part_0 = {{
    0x00, 0x00, 0x00, 0x00, 0xbf, 0x6a, 0x53, 0x3e, 0xd3, 0xc1,
    0xd0, 0x3e, 0x00, 0x00, 0x00, 0x00, 0xb6, 0xd8, 0xc0, 0x3e,
    0xca, 0xe7, 0x35, 0x3e, 0x23, 0xa5, 0x44, 0x3f, 0x61, 0xfd,
    0xd2, 0x3e, 0x00, 0x00, 0x00, 0x00, 0xb6, 0xe0, 0x43, 0x3c,
}};
const int32_t dnn_logits_biases_part_0_shape[1] = {1};
const union {
  uint8_t bytes[4];
  float values[1];
} dnn_logits_biases_part_0 = {{
    0x75,
    0xca,
    0xd7,
    0xbe,
}};
const int32_t dnn_logits_weights_part_0_shape[2] = {10, 1};
const union {
  uint8_t bytes[40];
  float values[10];
} dnn_logits_weights_part_0 = {{
    0x13, 0x12, 0x39, 0x3f, 0xf3, 0xa5, 0xc2, 0xbf, 0x81, 0x7f,
    0xbe, 0x3f, 0xf8, 0x17, 0x26, 0x3e, 0xa4, 0x19, 0xa6, 0x3f,
    0xf0, 0xc9, 0xb7, 0xbf, 0x6a, 0x99, 0xd2, 0x3f, 0x8a, 0x7d,
    0xe9, 0x3f, 0x83, 0x9a, 0x3a, 0xbf, 0xf1, 0x6c, 0x08, 0x3e,
}};

}  // anonymous namespace

// -----------------------------------------------------------------------------
// INFERENCE
// -----------------------------------------------------------------------------

int32_t input0Shape[2] = {1, 4};
int32_t logits_MatMul_merged_with_dnn_logits_BiasAdd0Shape[2] = {1, 1};

void Inference(
    base::span<const float> input0 /* shape: 1,4 */,
    base::span<float> logits_MatMul_merged_with_dnn_logits_BiasAdd0 /* shape:
                                                                       1,1 */
    ,
    FixedAllocations* __restrict fixed) {
  const int32_t input0_shape[] = {1, 4};
  std::array<int32_t, 2> logits_MatMul_merged_with_dnn_logits_BiasAdd0_shape;

  // dnn/hiddenlayer_0/MatMul_merged_with_dnn/hiddenlayer_0/BiasAdd
  FullyConnected<float>(input0_shape, input0,
                        dnn_hiddenlayer_0_weights_part_0_shape,
                        dnn_hiddenlayer_0_weights_part_0.values,
                        dnn_hiddenlayer_0_biases_part_0_shape,
                        dnn_hiddenlayer_0_biases_part_0.values, fixed->alloc0);
  fixed->alloc0_shape[0] = 1;
  fixed->alloc0_shape[1] = 10;

  // dnn/hiddenlayer_0/hiddenlayer_0/Relu
  Relu<float>(2,  // rank
              fixed->alloc0_shape, fixed->alloc0, fixed->alloc1);
  fixed->alloc1_shape[0] = 1;
  fixed->alloc1_shape[1] = 10;

  // dnn/logits/MatMul_merged_with_dnn/logits/BiasAdd
  FullyConnected<float>(
      fixed->alloc1_shape, fixed->alloc1, dnn_logits_weights_part_0_shape,
      dnn_logits_weights_part_0.values, dnn_logits_biases_part_0_shape,
      dnn_logits_biases_part_0.values,
      logits_MatMul_merged_with_dnn_logits_BiasAdd0);
  logits_MatMul_merged_with_dnn_logits_BiasAdd0_shape[0] = 1;
  logits_MatMul_merged_with_dnn_logits_BiasAdd0_shape[1] = 1;
}

}  // namespace darkmode_tfnative_model