// Copyright (c) 2022 PaddlePaddle 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 "paddle/phi/kernels/batch_norm_kernel.h"

#include "paddle/phi/backends/cpu/cpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/eigen/common.h"

namespace phi {

template <typename T>
using EigenArrayMap =
    Eigen::Map<Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>>;
template <typename T>
using ConstEigenArrayMap =
    Eigen::Map<const Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>>;
template <typename T>
using EigenVectorArrayMap = Eigen::Map<Eigen::Array<T, Eigen::Dynamic, 1>>;
template <typename T>
using ConstEigenVectorArrayMap =
    Eigen::Map<const Eigen::Array<T, Eigen::Dynamic, 1>>;

template <typename T, typename Context>
void BatchNormKernel(const Context& dev_ctx,
                     const DenseTensor& x,
                     const DenseTensor& mean,
                     const DenseTensor& variance,
                     const optional<DenseTensor>& scale,
                     const optional<DenseTensor>& bias,
                     bool is_test,
                     float momentum,
                     float epsilon,
                     const std::string& data_layout_str,
                     bool use_global_stats,
                     bool trainable_statistics,
                     DenseTensor* y,
                     DenseTensor* mean_out,
                     DenseTensor* variance_out,
                     DenseTensor* saved_mean,
                     DenseTensor* saved_variance,
                     DenseTensor* reserve_space) {
  if (x.numel() == 0) {
    dev_ctx.template Alloc<T>(y);
    if (mean_out) dev_ctx.template Alloc<T>(mean_out);
    if (variance_out) dev_ctx.template Alloc<T>(variance_out);
    if (saved_mean) dev_ctx.template Alloc<T>(saved_mean);
    if (saved_variance) dev_ctx.template Alloc<T>(saved_variance);
    if (reserve_space) {
      // infermeta dim is -1.
      reserve_space->Resize({0});
      dev_ctx.template Alloc<T>(reserve_space);
    }
    return;
  }
  bool test_mode = is_test && (!trainable_statistics);

  bool global_stats = test_mode || use_global_stats;

  auto data_layout = StringToDataLayout(data_layout_str);

  const auto& x_dims = x.dims();
  PADDLE_ENFORCE_GE(
      x_dims.size(),
      2,
      common::errors::InvalidArgument(
          "The size of input X's dimensions should be larger than 1."
          "But received: the size of input X's dimensions is [%d]",
          x_dims.size()));
  PADDLE_ENFORCE_LE(
      x_dims.size(),
      5,
      common::errors::InvalidArgument(
          "The size of input X's dimensions should be less than 6."
          "But received: the size of input X's dimensions is [%d]",
          x_dims.size()));
  const int64_t N = x_dims[0];
  const int64_t C =
      data_layout == DataLayout::NCHW ? x_dims[1] : x_dims[x_dims.size() - 1];
  const int64_t sample_size = x.numel() / N / C;
  const int64_t num_batch_channels = static_cast<int64_t>(N) * C;
  const int64_t num_batch_spatial = static_cast<int64_t>(N) * sample_size;

  // alloc memory
  dev_ctx.template Alloc<T>(y);
  dev_ctx.template Alloc<T>(mean_out);
  dev_ctx.template Alloc<T>(variance_out);
  dev_ctx.template Alloc<T>(saved_mean);
  dev_ctx.template Alloc<T>(saved_variance);
  if (reserve_space != nullptr) {
    reserve_space->Resize({0});
    dev_ctx.template Alloc<T>(reserve_space);
  }

  // input dimension is 2 and the format is NCHW. The input can be regarded
  // as NHWC format
  if (x_dims.size() == 2 && data_layout == DataLayout::NCHW) {
    data_layout = DataLayout::NHWC;
  }

  if (!global_stats) {
    // saved_xx is use just in this batch of data
    EigenVectorArrayMap<T> saved_mean_e(dev_ctx.template Alloc<T>(saved_mean),
                                        C);
    EigenVectorArrayMap<T> saved_variance_e(
        dev_ctx.template Alloc<T>(saved_variance), C);
    saved_mean_e.setZero();
    saved_variance_e.setZero();
    EigenVectorArrayMap<uint8_t> reserve_space_e(
        dev_ctx.template Alloc<uint8_t>(reserve_space), 0);
    reserve_space_e.setZero();

    EigenVectorArrayMap<T> running_mean_arr(dev_ctx.template Alloc<T>(mean_out),
                                            C);
    EigenVectorArrayMap<T> running_var_arr(
        dev_ctx.template Alloc<T>(variance_out), C);

    if (num_batch_spatial == 1) {
      // Only 1 element in normalization dimension,
      // we skip the batch norm calculation, let y = x.
      Copy(dev_ctx, x, dev_ctx.GetPlace(), false, y);
      return;
    }

    switch (data_layout) {
      case DataLayout::NCHW: {
        ConstEigenArrayMap<T> x_arr(
            x.data<T>(), sample_size, num_batch_channels);
        for (int64_t nc = 0; nc < num_batch_channels; ++nc) {
          saved_mean_e(nc % C) += x_arr.col(nc).sum();
        }
        saved_mean_e /= num_batch_spatial;
        for (int64_t nc = 0; nc < num_batch_channels; ++nc) {
          saved_variance_e(nc % C) +=
              (x_arr.col(nc) - saved_mean_e(nc % C)).matrix().squaredNorm();
        }
        saved_variance_e /= num_batch_spatial;
        break;
      }
      case DataLayout::NHWC: {
        ConstEigenArrayMap<T> x_arr(x.data<T>(), C, num_batch_spatial);
        for (int64_t i = 0; i < num_batch_spatial; ++i) {
          saved_mean_e += x_arr.col(i);
        }
        saved_mean_e /= num_batch_spatial;
        for (int64_t i = 0; i < num_batch_spatial; ++i) {
          saved_variance_e +=
              (x_arr.col(i) - saved_mean_e) * (x_arr.col(i) - saved_mean_e);
        }
        saved_variance_e /= num_batch_spatial;
        break;
      }
      default:
        PADDLE_THROW(common::errors::InvalidArgument(
            "Unknown storage order: %s", data_layout_str));
    }

    // if MomentumTensor is set, use MomentumTensor value, momentum
    // is only used in this training branch

    running_mean_arr =
        running_mean_arr * momentum + saved_mean_e * (1. - momentum);
    running_var_arr =
        running_var_arr * momentum + saved_variance_e * (1. - momentum);
  } else {
    const auto* est_mean = &mean;
    const auto* est_var = &variance;
    PADDLE_ENFORCE_EQ(
        est_mean->dims().size(),
        1UL,
        common::errors::InvalidArgument(
            "The size of mean's dimensions must equal to 1."
            "But received: the size of mean's dimensions mean is [%d],"
            "the dimensions of mean is [%s].",
            est_mean->dims().size(),
            est_mean->dims()));
    PADDLE_ENFORCE_EQ(
        est_var->dims().size(),
        1UL,
        common::errors::InvalidArgument(
            "The size of variance's dimensions must equal to 1."
            "But received: the size of variance's dimensions is [%d],"
            "the dimensions of variance is [%s].",
            est_var->dims().size(),
            est_var->dims()));
    PADDLE_ENFORCE_EQ(
        est_mean->dims()[0],
        C,
        common::errors::InvalidArgument(
            "The first dimension of mean must equal to the number of "
            "Channels, which is [%d]. But received: the first dimension "
            "of mean is [%d], the dimensions of mean is [%s].",
            C,
            est_mean->dims()[0],
            est_mean->dims()));
    PADDLE_ENFORCE_EQ(
        est_var->dims()[0],
        C,
        common::errors::InvalidArgument(
            "The first dimension of variance must equal to the number "
            "of Channels, which is [%d]. But received: the first dimension of "
            "variance is [%d], the dimensions of variance is [%s].",
            C,
            est_var->dims()[0],
            est_var->dims()));
  }

  // use SavedMean and SavedVariance to do normalize
  Eigen::Array<T, Eigen::Dynamic, 1> inv_std(C);
  if (global_stats) {  // NOLINT
    ConstEigenVectorArrayMap<T> var_arr(variance.data<T>(), C);
    inv_std = (var_arr + epsilon).sqrt().inverse();
  } else {
    EigenVectorArrayMap<T> saved_inv_std(saved_variance->data<T>(), C);
    // inverse SavedVariance first, gradient will use it too.
    saved_inv_std = (saved_inv_std + epsilon).inverse().sqrt();
    inv_std = saved_inv_std;
  }
  ConstEigenVectorArrayMap<T> mean_arr(
      global_stats ? mean.data<T>() : saved_mean->data<T>(), C);

  //   ((x - est_mean) * (inv_var) * scale + bias
  //   formula transform ====>
  //   (x * inv_var * scale) + (bias - est_mean * inv_var * scale)
  auto* Scale = scale.get_ptr();
  auto* Bias = bias.get_ptr();
  Eigen::Array<T, Eigen::Dynamic, 1> new_scale(C);
  Eigen::Array<T, Eigen::Dynamic, 1> new_bias(C);
  if (Scale && Bias) {  // NOLINT
    ConstEigenVectorArrayMap<T> scale_arr(Scale->data<T>(), C);
    ConstEigenVectorArrayMap<T> bias_arr(Bias->data<T>(), C);
    new_scale = inv_std * scale_arr;
    new_bias = bias_arr - mean_arr * inv_std * scale_arr;
  } else if (Scale) {
    ConstEigenVectorArrayMap<T> scale_arr(Scale->data<T>(), C);
    new_scale = inv_std * scale_arr;
    new_bias = -(mean_arr * inv_std * scale_arr);
  } else if (Bias) {
    ConstEigenVectorArrayMap<T> bias_arr(Bias->data<T>(), C);
    new_scale = inv_std;
    new_bias = bias_arr - mean_arr * inv_std;
  } else {
    new_scale = inv_std;
    new_bias = -(mean_arr * inv_std);
  }

  switch (data_layout) {
    case DataLayout::NCHW: {
      EigenArrayMap<T> y_arr(
          dev_ctx.template Alloc<T>(y), sample_size, num_batch_channels);
      ConstEigenArrayMap<T> x_arr(x.data<T>(), sample_size, num_batch_channels);
      for (int64_t nc = 0; nc < num_batch_channels; ++nc) {
        y_arr.col(nc) = x_arr.col(nc) * new_scale(nc % C) + new_bias(nc % C);
      }
      break;
    }
    case DataLayout::NHWC: {
      EigenArrayMap<T>(dev_ctx.template Alloc<T>(y), C, num_batch_spatial) =
          (ConstEigenArrayMap<T>(x.data<T>(), C, num_batch_spatial).colwise() *
           new_scale)
              .colwise() +
          new_bias;
      break;
    }
    default:
      PADDLE_THROW(common::errors::InvalidArgument("Unknown storage order: %d",
                                                   data_layout));
  }
}

}  // namespace phi

PD_REGISTER_KERNEL(
    batch_norm, CPU, ALL_LAYOUT, phi::BatchNormKernel, float, double) {
  kernel->OutputAt(5).SetDataType(phi::DataType::UINT8);
}