// @file   bnorm_cpu.cpp
// @brief  Batch normalization implementation (CPU)
// @author Sebastien Ehrhardt
// @author Andrea Vedaldi

/*
Copyright (C) 2015-16 Sebastien Ehrhardt and Andrea Vedaldi.
All rights reserved.

This file is part of the VLFeat library and is made available under
the terms of the BSD license (see the COPYING file).
*/

#include "bnorm.hpp"
#include "../data.hpp"
#include <math.h>
#include <memory.h>
#include <cstdlib>
#include <algorithm>
#include <limits>
#include <cassert>

/* ---------------------------------------------------------------- */
/*          compute_moments, compute_ders, compute_ders_and_moments	*/
/* ---------------------------------------------------------------- */

// Compute moments (means and sigmas) from the batch data
// WH is the product of the data width and height
// moments is a 2 x depth array with means and sigmas

template<typename T> inline void
compute_moments(T * moments,
                T const * data,
                int WH,
                int depth,
                int num,
                T epsilon)
{
  int mass = WH * num ;
  for(int channel = 0; channel < depth; ++channel) {
    for(int element = 0; element < num; ++element) {
      for(int wh = 0; wh < WH; ++wh){
        T x = data[wh + channel*WH + element*(depth*WH)] ;
        moments[channel] += x ; // mean
        moments[channel + depth] += x * x; // sigma
      }
    }
  }
  for(int i = 0; i < depth; ++i) {
    moments[i] /= mass;
    moments[i + depth] = sqrt(moments[i + depth]/mass
                              - moments[i]*moments[i] + epsilon);
  }
}

// this version assumes that moments is precomputed
template<typename T> inline void
compute_ders(T * derMultipliers,
             T * derBiases,
             T const * moments,
             T const * data,
             T const * derOutput,
             int WH, int depth, int num,
             T epsilon)
{
  memset(derMultipliers, 0, sizeof(T) * depth) ;
  memset(derBiases, 0, sizeof(T) * depth) ;
  for(int channel = 0; channel < depth; ++channel){
    for(int element = 0; element < num; ++element ){
      for(int wh = 0; wh < WH; ++wh){
        int offset = wh + channel*WH + element * (WH*depth) ;
        derMultipliers[channel] += derOutput[offset] * data[offset];
        derBiases[channel] += derOutput[offset];
      }
    }
  }

  T mass = WH*num;
  for(int i = 0; i < depth; ++i) {
    T mean = moments[i] ;
    T sigma = moments[i + depth] ;
    derMultipliers[i] = (derMultipliers[i] - mean*derBiases[i]) / sigma;
  }
}

template<typename T> inline void
compute_ders_and_moments(T * derMultipliers,
                         T * derBiases,
                         T * moments,
                         T const * data,
                         T const * derOutput,
                         int WH, int depth, int num,
                         T epsilon)
{
  memset(derMultipliers, 0, sizeof(T) * depth) ;
  memset(derBiases, 0, sizeof(T) * depth) ;
  for(int channel = 0; channel < depth; ++channel){
    for(int element = 0; element < num; ++element ){
      for(int wh = 0; wh < WH; ++wh){
        int offset = wh + channel*WH + element * (WH*depth) ;
        moments[channel] += data[offset] ;
        moments[channel + depth] += data[offset] * data[offset];
        derMultipliers[channel] += derOutput[offset] * data[offset];
        derBiases[channel] += derOutput[offset];
      }
    }
  }

  T mass = WH*num;
  for(int i = 0; i < depth; ++i) {
    T mean = moments[i] /= mass ;
    T sigma = sqrt(moments[i + depth]/mass - mean*mean + epsilon);
    moments[i] = mean ;
    moments[i + depth] = sigma ;
    derMultipliers[i] = (derMultipliers[i] - mean*derBiases[i]) / sigma;
  }
}

/* ---------------------------------------------------------------- */
/*                                         batch_normalize_backward	*/
/* ---------------------------------------------------------------- */

template<typename T> inline void
batch_normalize_backward(T * derData,
                         T const * moments,
                         T const * data,
                         T const * multipliers,
                         T const * derMultipliers,
                         T const * derBiases,
                         T const * derOutput,
                         int WH,
                         int depth,
                         int num)
{
  T mass = WH*num;
  for(int channel = 0; channel < depth; ++channel ) {
    T mean = moments[channel] ;
    T sigma = moments[channel + depth] ;

    T muz = derBiases[channel]/mass;
    T G1 = multipliers[channel]/sigma ;
    T G2 = G1 * derMultipliers[channel]/(mass*sigma);

    for(int element = 0; element < num; ++element){
      for(int wh = 0; wh < WH; ++wh){
        int offset = wh + channel*WH + element * (WH*depth) ;
        derData[offset] = G1 * (derOutput[offset] - muz) - G2 * (data[offset]-mean) ;
      }
    }
  }
}

/* ---------------------------------------------------------------- */
/*                                                           driver */
/* ---------------------------------------------------------------- */

namespace vl { namespace impl {

  template<typename T>
  struct bnorm<vl::CPU,T>
  {

    /* ------------------------------------------------------------ */
    /*                                                      forward */
    /* ------------------------------------------------------------ */

    static vl::Error
    forward_given_moments(Context& context,
                          T* output,
                          T const* moments,
                          T const* data,
                          T const* multipliers,
                          T const* biases,
                          size_t height, size_t width, size_t depth, size_t num)
    {
      int WH = height * width ;
      for(int channel = 0; channel < depth; ++channel) {
        T mean = moments[channel] ;
        T sigma = moments[channel + depth] ;
        T bias = biases[channel];
        T coefficient = multipliers[channel] / sigma ;

        for(int element = 0; element < num; ++element) {
          for(int wh = 0; wh < WH; ++wh){
            int offset = wh + channel*WH + element * (depth*WH) ;
            output[offset] = coefficient * (data[offset] - mean) + bias ;
          }
        }
      }
      return vlSuccess;
    }

    static vl::Error
    forward(Context& context,
            T* output,
            T* moments,
            T const* data,
            T const* multipliers,
            T const* biases,
            size_t height, size_t width, size_t depth, size_t size,
            T epsilon)
    {
      vl::Error error = vlSuccess ;
      bool ownMoments = false ;
      if (moments == NULL) {
        moments = (T*)calloc(sizeof(T),2*depth);
        if (!moments) {
          error = vlErrorOutOfMemory ;
          goto done ;
        }
        ownMoments = true ;
      } else {
        memset(moments, 0, sizeof(T) * 2*depth) ;
      }
      compute_moments<T>(moments,
                         data, width*height, depth, size,
                         epsilon) ;

      error = bnorm<vl::CPU,T>::forward_given_moments
      (context,
       output,
       moments, data,
       multipliers, biases,
       height, width, depth, size) ;

      // Delete intermediate variable
    done:
      if (ownMoments)  { free(moments) ; }
      return error ;
    }

    /*------------------------------------------------------------- */
    /*                                                     backward */
    /* ------------------------------------------------------------ */

    static vl::Error
    backward_given_moments(Context& context,
                           T* derData,
                           T* derMultipliers,
                           T* derBiases,
                           T const* moments,
                           T const* data,
                           T const* multipliers,
                           T const* biases,
                           T const* derOutput,
                           size_t height, size_t width, size_t depth, size_t size,
                           T epsilon)
    {
      vl::Error error = vlSuccess ;
      T * muz;
      int WH = width * height;

      // Allocate muz
      muz = (T*)calloc(sizeof(T),depth);
      if (!muz) {
        error = vlErrorOutOfMemory ;
        goto done ;
      }

      // Compute derMultipliers, derBiases, muz, and moments
      compute_ders<T>(derMultipliers, derBiases,
                      moments, data, derOutput,
                      WH, depth, size,
                      epsilon);

      // Compute derData
      batch_normalize_backward<T>(derData,
                                  moments, data, muz,
                                  multipliers, derMultipliers, derOutput,
                                  WH, depth, size);

    done:;
      if (muz) { free(muz) ; }
      return error ;
    }

    static vl::Error
    backward(Context& context,
             T* derData,
             T* derMultipliers,
             T* derBiases,
             T* moments,
             T const* data,
             T const* multipliers,
             T const* biases,
             T const* derOutput,
             size_t height, size_t width, size_t depth, size_t size,
             T epsilon)
    {
      vl::Error error = vlSuccess ;
      T* muz = NULL ;
      bool ownMoments = false ;
      int WH = width * height;

      // Allocate or reuse moments
      if (moments == NULL) {
        moments = (T*)calloc(sizeof(T),2*depth);
        if (!moments) {
          error = vlErrorOutOfMemory ;
          goto done ;
        }
        ownMoments = true ;
      } else {
        memset(moments, 0, sizeof(T) * 2*depth) ;
      }

      // Compute derMultipliers, derBiases, and moments
      compute_ders_and_moments<T>(derMultipliers, derBiases, moments,
                                  data, derOutput,
                                  WH, depth, size,
                                  epsilon);

      // Compute derData
      batch_normalize_backward<T>(derData,
                                  moments, data,
                                  multipliers,
                                  derMultipliers, derBiases, derOutput,
                                  WH, depth, size);

      // Delete intermediate variable
    done:;
      if (ownMoments) { free(moments) ; }
      return error ;
    }
  } ;

} } // namespace vl::impl

template struct vl::impl::bnorm<vl::CPU, float> ;

#ifdef ENABLE_DOUBLE
template struct vl::impl::bnorm<vl::CPU, double> ;
#endif

