// nnet/nnet-max-pooling-component.h

// Copyright 2014  Brno University of Technology (author: Karel Vesely)

// See ../../COPYING for clarification regarding multiple authors
//
// 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
//
// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
// MERCHANTABLITY OR NON-INFRINGEMENT.
// See the Apache 2 License for the specific language governing permissions and
// limitations under the License.


#ifndef KALDI_NNET_NNET_MAX_POOLING_COMPONENT_H_
#define KALDI_NNET_NNET_MAX_POOLING_COMPONENT_H_

#include <string>
#include <vector>

#include "nnet/nnet-component.h"
#include "nnet/nnet-utils.h"
#include "cudamatrix/cu-math.h"

namespace kaldi {
namespace nnet1 {

/**
 * MaxPoolingComponent :
 * The input/output matrices are split to submatrices with width 'pool_stride_'.
 * The pooling is done over 3rd axis, of the set of 2d matrices.
 * Our pooling supports overlaps, overlaps occur when (pool_step_ < pool_size_).
 */
class MaxPoolingComponent : public Component {
 public:
  MaxPoolingComponent(int32 dim_in, int32 dim_out):
    Component(dim_in, dim_out),
    pool_size_(0),
    pool_step_(0),
    pool_stride_(0)
  { }

  ~MaxPoolingComponent()
  { }

  Component* Copy() const { return new MaxPoolingComponent(*this); }
  ComponentType GetType() const { return kMaxPoolingComponent; }

  void InitData(std::istream &is) {
    // parse config
    std::string token;
    while (is >> std::ws, !is.eof()) {
      ReadToken(is, false, &token);
      /**/ if (token == "<PoolSize>") ReadBasicType(is, false, &pool_size_);
      else if (token == "<PoolStep>") ReadBasicType(is, false, &pool_step_);
      else if (token == "<PoolStride>") ReadBasicType(is, false, &pool_stride_);
      else KALDI_ERR << "Unknown token " << token << ", a typo in config?"
                     << " (PoolSize|PoolStep|PoolStride)";
    }
    // check
    KALDI_ASSERT(pool_size_ != 0 && pool_step_ != 0 && pool_stride_ != 0);
  }

  void ReadData(std::istream &is, bool binary) {
    // pooling hyperparameters
    ExpectToken(is, binary, "<PoolSize>");
    ReadBasicType(is, binary, &pool_size_);
    ExpectToken(is, binary, "<PoolStep>");
    ReadBasicType(is, binary, &pool_step_);
    ExpectToken(is, binary, "<PoolStride>");
    ReadBasicType(is, binary, &pool_stride_);

    //
    // Sanity checks:
    //
    // number of patches:
    KALDI_ASSERT(input_dim_ % pool_stride_ == 0);
    int32 num_patches = input_dim_ / pool_stride_;
    // number of pools:
    KALDI_ASSERT((num_patches - pool_size_) % pool_step_ == 0);
    int32 num_pools = 1 + (num_patches - pool_size_) / pool_step_;
    // check output dim:
    KALDI_ASSERT(output_dim_ == num_pools * pool_stride_);
    //
  }

  void WriteData(std::ostream &os, bool binary) const {
    // pooling hyperparameters
    WriteToken(os, binary, "<PoolSize>");
    WriteBasicType(os, binary, pool_size_);
    WriteToken(os, binary, "<PoolStep>");
    WriteBasicType(os, binary, pool_step_);
    WriteToken(os, binary, "<PoolStride>");
    WriteBasicType(os, binary, pool_stride_);
  }

  void PropagateFnc(const CuMatrixBase<BaseFloat> &in,
                    CuMatrixBase<BaseFloat> *out) {
    // useful dims
    int32 num_patches = input_dim_ / pool_stride_;
    int32 num_pools = 1 + (num_patches - pool_size_) / pool_step_;

    // do the max-pooling (pools indexed by q)
    for (int32 q = 0; q < num_pools; q++) {
      // get output buffer of the pool
      CuSubMatrix<BaseFloat> pool(out->ColRange(q*pool_stride_, pool_stride_));
      pool.Set(-1e20);  // reset (large negative value)
      for (int32 r = 0; r < pool_size_; r++) {  // max
        int32 p = r + q * pool_step_;  // p = input patch
        pool.Max(in.ColRange(p*pool_stride_, pool_stride_));
      }
    }
  }

  void BackpropagateFnc(const CuMatrixBase<BaseFloat> &in,
                        const CuMatrixBase<BaseFloat> &out,
                        const CuMatrixBase<BaseFloat> &out_diff,
                        CuMatrixBase<BaseFloat> *in_diff) {
    // useful dims
    int32 num_patches = input_dim_ / pool_stride_;
    int32 num_pools = 1 + (num_patches - pool_size_) / pool_step_;

    //
    // here we note how many diff matrices are summed for each input patch,
    std::vector<int32> patch_summands(num_patches, 0);
    // this metainfo will be used to divide diff of patches
    // used in more than one pool.
    //

    in_diff->SetZero();  // reset

    for (int32 q = 0; q<num_pools; q++) {  // sum
      for (int32 r = 0; r<pool_size_; r++) {
        int32 p = r + q * pool_step_;  // patch number
        //
        CuSubMatrix<BaseFloat> in_p(in.ColRange(p*pool_stride_, pool_stride_));
        CuSubMatrix<BaseFloat> out_q(out.ColRange(q*pool_stride_, pool_stride_));
        //
        CuSubMatrix<BaseFloat> tgt(in_diff->ColRange(p*pool_stride_, pool_stride_));
        CuMatrix<BaseFloat> src(out_diff.ColRange(q*pool_stride_, pool_stride_));

        // Only the pool-inputs with 'max-values' are used to back-propagate into,
        // the rest of derivatives is zeroed-out by a mask.
        CuMatrix<BaseFloat> mask;
        in_p.EqualElementMask(out_q, &mask);
        src.MulElements(mask);
        tgt.AddMat(1.0, src);

        patch_summands[p] += 1;
      }
    }

    // divide diff by #summands (compensate for patches used in more pools)
    for (int32 p = 0; p < num_patches; p++) {
      CuSubMatrix<BaseFloat> tgt(in_diff->ColRange(p*pool_stride_, pool_stride_));
      KALDI_ASSERT(patch_summands[p] > 0);  // patch at least in one pool
      tgt.Scale(1.0/patch_summands[p]);
    }
  }

 private:
  int32 pool_size_,    // input patches used for pooling
        pool_step_,    // shift used for pooling (allow overlapping pools)
        pool_stride_;  // stride used to slice input to a vector of matrices
};

}  // namespace nnet1
}  // namespace kaldi

#endif  // KALDI_NNET_NNET_MAX_POOLING_COMPONENT_H_