text: #pragma once
  //------------------------------------------------------------------------------
  //
  //   Copyright 2018-2019 Fetch.AI Limited
  //
  //   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 "ml/ops/ops.hpp"
  
  #include <cassert>
  #include <vector>
  
  namespace fetch {
  namespace ml {
  namespace ops {
  
  template <class T>
  class AvgPool2D : public Ops<T>
  {
  public:
    using TensorType    = T;
    using SizeType      = fetch::math::SizeType;
    using DataType      = typename TensorType::Type;
    using ArrayPtrType  = std::shared_ptr<TensorType>;
    using VecTensorType = typename Ops<T>::VecTensorType;
    using SPType        = OpAvgPool2DSaveableParams<T>;
    using MyType        = AvgPool2D<TensorType>;
  
    AvgPool2D(SizeType const kernel_size, SizeType const stride_size)
      : kernel_size_{kernel_size}
      , stride_size_{stride_size}
    {}
  
    explicit AvgPool2D(SPType const &sp)
      : Ops<T>(sp)
    {
      kernel_size_ = sp.kernel_size;
      stride_size_ = sp.stride_size;
    }
  
    ~AvgPool2D() override = default;
  
    std::shared_ptr<OpsSaveableParams> GetOpSaveableParams() override
    {
      SPType sp{};
      sp.kernel_size = kernel_size_;
      sp.stride_size = stride_size_;
      return std::make_shared<SPType>(sp);
    }
  
    std::shared_ptr<fetch::ml::ops::Ops<TensorType>> MakeSharedCopy(
        std::shared_ptr<fetch::ml::ops::Ops<TensorType>> me) override
    {
      FETCH_UNUSED(me);
      assert(me.get() == this);
  
      auto copyshare = std::make_shared<MyType>(*this);  // calls default copy constructor of MyType
  
      return copyshare;
    }
    /**
     * Applies 2D avg pooling of kernel_size_ x kernel_size_ for each channel described here:
     * http://ais.uni-bonn.de/papers/icann2010_maxpool.pdf
     * @param inputs vector of tensor references where at:
     * inputs[0] = input_data[input_channels x input_height x input_width]
     * @param output tensor of size [input_channels=output_channels x
     * number_of_stride_sized_steps_over_input_height x number_of_stride_sized_steps_over_input_width]
     * @return: output tensor parameter
     */
    void Forward(VecTensorType const &inputs, TensorType &output) override
    {
      assert(inputs.size() == 1);
      // Input must be a 4D tensor [C x W x H x N]
      assert(inputs.at(0)->shape().size() == 4);
      assert(output.shape() == ComputeOutputShape(inputs));
  
      SizeType iterw;
      SizeType iterh;
      DataType sum;
      auto     cnt = static_cast<DataType>(kernel_size_ * kernel_size_);
      auto     oit = output.begin();
  
      for (SizeType n_i{0}; n_i < output.shape().at(2); n_i++)  // iterate over batch
      {
        for (SizeType ih{0}; ih < output.shape().at(2); ih++)  // Iterate height over kernel stride
        {
          iterh = ih * stride_size_;
  
          for (SizeType iw{0}; iw < output.shape().at(1); iw++)  // Iterate width over kernel stride
          {
            iterw = iw * stride_size_;
  
            for (SizeType c{0}; c < output.shape().at(0); ++c)  // Iterate over output channels
            {
              sum = static_cast<DataType>(0);
  
              // Sum all values on kernel_size_ x kernel_size_ window
              for (SizeType jw{0}; jw < kernel_size_; jw++)  // Iterate over kernel width
              {
                for (SizeType jh{0}; jh < kernel_size_; jh++)  // Iterate over kernel width
                {
                  sum += inputs.at(0)->At(c, iterw + jw, iterh + jh, n_i);
                }
              }
  
              // Set average value for each [kernel_size_ x kernel_size_] window to output
              *oit = sum / cnt;
              ++oit;
            }
          }
        }
      }
    }
  
    /**
     * Computes gradient of 2D avg pooling of kernel_size_ x kernel_size for each channel described
     * here: http://ais.uni-bonn.de/papers/icann2010_maxpool.pdf Error signal
     * of avg pool is averaged to all node
     * @param inputs vector of tensor references where at:
     * inputs[0] = input_data[input_channels x input_height x input_width]
     * @param error_signal tensor of size  [output_channels x
     * number_of_stride_sized_steps_over_input_height x number_of_stride_sized_steps_over_input_width]
     * @return: output vector of tensors with back propagated error signal
     * output[0]=input_error[inputs[0].shape]
     */
    std::vector<TensorType> Backward(VecTensorType const &inputs,
                                     TensorType const &   error_signal) override
    {
      assert(inputs.size() == 1);
      assert(error_signal.shape() == ComputeOutputShape(inputs));
      TensorType return_signal{inputs.at(0)->shape()};
  
      SizeType iterh;
      SizeType iterw;
      auto     cnt = static_cast<DataType>(kernel_size_ * kernel_size_);
  
      auto erit = error_signal.cbegin();
      for (SizeType n_i{0}; n_i < error_signal.shape().at(2); n_i++)  // iterate over batch
      {
        // Iterate width over kernel stride
        for (SizeType iw{0}; iw < error_signal.shape().at(1); iw++)
        {
          iterw = iw * stride_size_;
          // Iterate height over kernel stride
          for (SizeType ih{0}; ih < error_signal.shape().at(2); ih++)
          {
            iterh = ih * stride_size_;
            // Iterate over output channels
            for (SizeType c{0}; c < error_signal.shape().at(0); ++c)
            {
              // Add error signal to return_signal nodes divided by number of kernel nodes
              for (SizeType jw{0}; jw < kernel_size_; jw++)  // Iterate over kernel width
              {
                for (SizeType jh{0}; jh < kernel_size_; jh++)  // Iterate over kernel width
                {
                  return_signal(c, iterw + jw, iterh + jh, n_i) += *erit / cnt;
                }
              }
  
              ++erit;
            }
          }
        }
      }
  
      return {return_signal};
    }
  
    std::vector<SizeType> ComputeOutputShape(VecTensorType const &inputs) const override
    {
      std::vector<SizeType> output_shape;
  
      // output_shape_[0]=number of output channels
      output_shape.emplace_back(inputs.at(0)->shape().at(0));
      // output_shape_[1]=number of stride_size steps over input height
      output_shape.emplace_back((inputs.at(0)->shape().at(1) - (kernel_size_ - stride_size_)) /
                                stride_size_);
      // output_shape_[2]=number of stride_size steps over input width
      output_shape.emplace_back((inputs.at(0)->shape().at(2) - (kernel_size_ - stride_size_)) /
                                stride_size_);
      // output_shape_[3]=batch dimension
      output_shape.emplace_back(inputs.at(0)->shape().at(3));
      return output_shape;
    }
  
    static constexpr OpType OpCode()
    {
      return OpType::OP_AVG_POOL_2D;
    }
    static constexpr char const *DESCRIPTOR = "AvgPool2D";
  
  private:
    SizeType kernel_size_;
    SizeType stride_size_;
  };
  
  }  // namespace ops
  }  // namespace ml
  }  // namespace fetch
  