
#include <core/graph/onnx_protobuf.h>
#include <core/framework/tensor.h>
#include <core/util/math_cpuonly.h>
#include "core/platform/threadpool.h"
#include <benchmark/benchmark.h>
#include <random>

using namespace onnxruntime;

using namespace onnxruntime;

template <typename T>
void SetRandom(Tensor& input) {
  std::random_device rd;
  std::mt19937 gen(rd());

  std::uniform_real_distribution<T> distr(0, 1);
  auto span = input.MutableDataAsSpan<T>();
  std::generate(span.begin(), span.end(), [&]() { return distr(gen); });
}

static void BM_BatchNormOldEigen(benchmark::State& state) {
  AllocatorPtr alloc = CPUAllocator::DefaultInstance();
  const int64_t batch_size = state.range(0);

  const TensorShape shape = {batch_size, 64, 75, 75};
  using T = float;

  Tensor X(DataTypeImpl::GetType<float>(), shape, alloc);
  SetRandom<T>(X);
  const TensorShape& x_shape = X.Shape();
  Tensor Y(DataTypeImpl::GetType<float>(), shape, alloc);
  Tensor scale(DataTypeImpl::GetType<float>(), {shape[1]}, alloc);
  SetRandom<T>(scale);
  Tensor mean(DataTypeImpl::GetType<float>(), {shape[1]}, alloc);
  SetRandom<T>(mean);

  Tensor B(DataTypeImpl::GetType<float>(), {shape[1]}, alloc);
  SetRandom<T>(B);

  Tensor var(DataTypeImpl::GetType<float>(), {shape[1]}, alloc);
  SetRandom<T>(var);

  bool is_spatial_ = true;
  double epsilon_ = 1e-5;
  const auto& dims_vec = x_shape.GetDims();
  const size_t N = dims_vec[0];
  const size_t C = dims_vec[1];  // assume NCHW as per the spec

  // calculate sample_size (per individual channel)
  size_t sample_size = 1;
  for (size_t i = 2; i < dims_vec.size(); ++i) {
    sample_size *= dims_vec[i];
  }

  // calculate sample_size (including all channels)
  size_t sample_size_incl_all_channels = sample_size * C;
  for (auto _ : state) {
    ConstEigenVectorArrayMap<T> scale_arr(scale.Data<T>(), is_spatial_ ? C : sample_size_incl_all_channels);
    ConstEigenVectorArrayMap<T> bias_arr(B.Data<T>(), is_spatial_ ? C : sample_size_incl_all_channels);

    // Regardless of training or testing, we will apply the estimated mean
    // and standard deviation to the input. For testing, they are
    // specified directly by the input, and for training, they are computed
    // by the op.
    Eigen::Array<T, Eigen::Dynamic, 1> inv_std(is_spatial_ ? C : sample_size_incl_all_channels);
    ConstEigenVectorArrayMap<T> var_arr(var.Data<T>(), is_spatial_ ? C : sample_size_incl_all_channels);
    inv_std = (var_arr + epsilon_).sqrt().inverse();
    ConstEigenVectorArrayMap<T> mean_arr(mean.Data<T>(), is_spatial_ ? C : sample_size_incl_all_channels);
    // We can fuse the output computation as follows:
    //   ((x - est_mean) * (inv_var) * scale + bias
    // to
    //   (x * inv_var * scale) + (bias - est_mean * inv_var * scale)
    Eigen::Array<T, Eigen::Dynamic, 1> new_scale = inv_std * scale_arr;
    Eigen::Array<T, Eigen::Dynamic, 1> new_bias = bias_arr - mean_arr * new_scale;
    EigenArrayMap<T> Y_arr(Y.MutableData<T>(), is_spatial_ ? sample_size : sample_size_incl_all_channels,
                           is_spatial_ ? N * C : N);
    ConstEigenArrayMap<T> X_arr(X.Data<T>(), is_spatial_ ? sample_size : sample_size_incl_all_channels,
                                is_spatial_ ? N * C : N);
    if (is_spatial_) {  // spatial == 1
      for (size_t nc = 0; nc < N * C; ++nc) {
        Y_arr.col(nc) = X_arr.col(nc) * new_scale(nc % C) + new_bias(nc % C);
      }
    } else {  // spatial == 0
      for (size_t n = 0; n < N; ++n) {
        Y_arr.col(n) = X_arr.col(n) * new_scale.col(0) + new_bias.col(0);
      }
    }
  }
}
BENCHMARK(BM_BatchNormOldEigen)
    ->Arg(1)
    ->Arg(16)
    ->Arg(64)
    ->UseRealTime()
    ->Unit(benchmark::TimeUnit::kMicrosecond);
