// nnet/nnet-various.h
  
  // Copyright 2012-2016  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_VARIOUS_H_
  #define KALDI_NNET_NNET_VARIOUS_H_
  
  #include <string>
  #include <vector>
  #include <algorithm>
  #include <sstream>
  
  #include "nnet/nnet-component.h"
  #include "nnet/nnet-utils.h"
  #include "cudamatrix/cu-math.h"
  #include "util/text-utils.h"
  
  namespace kaldi {
  namespace nnet1 {
  
  /**
   * Splices the time context of the input features
   * in N, out k*N, FrameOffset o_1,o_2,...,o_k
   * FrameOffset example 11frames: -5 -4 -3 -2 -1 0 1 2 3 4 5
   */
  class Splice: public Component {
   public:
    Splice(int32 dim_in, int32 dim_out):
      Component(dim_in, dim_out)
    { }
  
    ~Splice()
    { }
  
    Component* Copy() const { return new Splice(*this); }
    ComponentType GetType() const { return kSplice; }
  
    void InitData(std::istream &is) {
      // define options,
      std::vector<std::vector<int32> > build_vector;
      // parse config,
      std::string token;
      while (is >> std::ws, !is.eof()) {
        ReadToken(is, false, &token);
        /**/ if (token == "<ReadVector>") {
          frame_offsets_.Read(is, false);
        } else if (token == "<BuildVector>") {
          // Parse the list of 'matlab-like' indices:
          // <BuildVector> 1:1:1000 1 2 3 1:10 </BuildVector>
          while (is >> std::ws, !is.eof()) {
            std::string colon_sep_list_or_end;
            ReadToken(is, false, &colon_sep_list_or_end);
            if (colon_sep_list_or_end == "</BuildVector>") break;
            std::vector<int32> v;
            SplitStringToIntegers(colon_sep_list_or_end, ":", false, &v);
            build_vector.push_back(v);
          }
        } else {
          KALDI_ERR << "Unknown token " << token << ", a typo in config?"
                    << " (ReadVector|BuildVector)";
        }
      }
  
      if (build_vector.size() > 0) {
        // build the vector, using <BuildVector> ... </BuildVector> inputs,
        BuildIntegerVector(build_vector, &frame_offsets_);
      }
  
      // check dim
      KALDI_ASSERT(frame_offsets_.Dim()*InputDim() == OutputDim());
    }
  
    void ReadData(std::istream &is, bool binary) {
      frame_offsets_.Read(is, binary);
      KALDI_ASSERT(frame_offsets_.Dim() * InputDim() == OutputDim());
    }
  
    void WriteData(std::ostream &os, bool binary) const {
      frame_offsets_.Write(os, binary);
    }
  
    std::string Info() const {
      std::ostringstream ostr;
      ostr << "
    frame_offsets " << frame_offsets_;
      std::string str = ostr.str();
      str.erase(str.end()-1);
      return str;
    }
  
    void PropagateFnc(const CuMatrixBase<BaseFloat> &in,
                      CuMatrixBase<BaseFloat> *out) {
      cu::Splice(in, frame_offsets_, out);
    }
  
    void BackpropagateFnc(const CuMatrixBase<BaseFloat> &in,
                          const CuMatrixBase<BaseFloat> &out,
                          const CuMatrixBase<BaseFloat> &out_diff,
                          CuMatrixBase<BaseFloat> *in_diff) {
      // WARNING!!! WARNING!!! WARNING!!!
      // THIS BACKPROPAGATION CAN BE USED ONLY WITH 'PER-UTTERANCE' TRAINING!
      // IN MINI-BATCH TRAINING, THIS <Splice> COMPONENT HAS TO BE PART OF THE
      // 'feature_transform' SO WE DON'T BACKPROPAGATE THROUGH IT...
  
      // dims,
      int32 input_dim = in.NumCols(),
            num_frames = out_diff.NumRows();
      // Copy offsets to 'host',
      std::vector<int32> offsets(frame_offsets_.Dim());
      frame_offsets_.CopyToVec(&offsets);
      // loop over the offsets,
      for (int32 i = 0; i < offsets.size(); i++) {
        int32 o_i = offsets.at(i);
        int32 n_rows = num_frames - abs(o_i),
              src_row = std::max(-o_i, 0),
              tgt_row = std::max(o_i, 0);
        const CuSubMatrix<BaseFloat> src = out_diff.Range(src_row, n_rows, i*input_dim, input_dim);
        CuSubMatrix<BaseFloat> tgt = in_diff->RowRange(tgt_row, n_rows);
        tgt.AddMat(1.0, src, kNoTrans);
      }
    }
  
   protected:
    CuArray<int32> frame_offsets_;
  };
  
  
  /**
   * Rearrange the matrix columns according to the indices in copy_from_indices_
   */
  class CopyComponent: public Component {
   public:
    CopyComponent(int32 dim_in, int32 dim_out):
      Component(dim_in, dim_out)
    { }
  
    ~CopyComponent()
    { }
  
    Component* Copy() const { return new CopyComponent(*this); }
    ComponentType GetType() const { return kCopy; }
  
    void InitData(std::istream &is) {
      // define options,
      std::vector<std::vector<int32> > build_vector;
      // parse config,
      std::string token;
      while (is >> std::ws, !is.eof()) {
        ReadToken(is, false, &token);
        /**/ if (token == "<ReadVector>") {
          copy_from_indices_.Read(is, false);
        } else if (token == "<BuildVector>") {
          // <BuildVector> 1:1:1000 1:1:1000 1 2 3 1:10 </BuildVector>
          // 'matlab-line' indexing, read the colon-separated-lists:
          while (is >> std::ws, !is.eof()) {
            std::string colon_sep_list_or_end;
            ReadToken(is, false, &colon_sep_list_or_end);
            if (colon_sep_list_or_end == "</BuildVector>") break;
            std::vector<int32> v;
            SplitStringToIntegers(colon_sep_list_or_end, ":", false, &v);
            build_vector.push_back(v);
          }
        } else {
          KALDI_ERR << "Unknown token " << token << ", a typo in config?"
                    << " (ReadVector|BuildVector)";
        }
      }
  
      if (build_vector.size() > 0) {
        // build the vector, using <BuildVector> ... </BuildVector> inputs,
        BuildIntegerVector(build_vector, &copy_from_indices_);
      }
  
      // decrease by 1,
      copy_from_indices_.Add(-1);
  
      // check range,
      KALDI_ASSERT(copy_from_indices_.Min() >= 0);
      KALDI_ASSERT(copy_from_indices_.Max() < InputDim());
      // check dim,
      KALDI_ASSERT(copy_from_indices_.Dim() == OutputDim());
    }
  
    void ReadData(std::istream &is, bool binary) {
      copy_from_indices_.Read(is, binary);
      KALDI_ASSERT(copy_from_indices_.Dim() == OutputDim());
      copy_from_indices_.Add(-1);  // -1 from each element,
    }
  
    void WriteData(std::ostream &os, bool binary) const {
      CuArray<int32> tmp(copy_from_indices_);
      tmp.Add(1);  // +1 to each element,
      tmp.Write(os, binary);
    }
  
    std::string Info() const {
      return std::string("
    min ") + ToString(copy_from_indices_.Min()) +
                           ", max "  + ToString(copy_from_indices_.Max());
    }
  
    void PropagateFnc(const CuMatrixBase<BaseFloat> &in,
                      CuMatrixBase<BaseFloat> *out) {
      cu::Copy(in, copy_from_indices_,out);
    }
  
    void BackpropagateFnc(const CuMatrixBase<BaseFloat> &in,
                          const CuMatrixBase<BaseFloat> &out,
                          const CuMatrixBase<BaseFloat> &out_diff,
                          CuMatrixBase<BaseFloat> *in_diff) {
      static bool warning_displayed = false;
      if (!warning_displayed) {
        KALDI_WARN << Component::TypeToMarker(GetType()) << " : "
                   << __func__ << "() Not implemented!";
  
        warning_displayed = true;
      }
      in_diff->SetZero();
    }
  
   protected:
    CuArray<int32> copy_from_indices_;
  };
  
  
  
  /**
   * Rescale the matrix-rows to have unit length (L2-norm).
   */
  class LengthNormComponent: public Component {
   public:
    LengthNormComponent(int32 dim_in, int32 dim_out):
      Component(dim_in, dim_out)
    { }
  
    ~LengthNormComponent()
    { }
  
    Component* Copy() const { return new LengthNormComponent(*this); }
    ComponentType GetType() const { return kLengthNormComponent; }
  
    void PropagateFnc(const CuMatrixBase<BaseFloat> &in,
                      CuMatrixBase<BaseFloat> *out) {
      // resize vector when needed,
      if (row_scales_.Dim() != in.NumRows()) {
        row_scales_.Resize(in.NumRows());
      }
      // get the normalization scalars,
      l2_aux_ = in;
      l2_aux_.MulElements(l2_aux_);  // x^2,
      row_scales_.AddColSumMat(1.0, l2_aux_, 0.0);  // sum_of_cols(x^2),
      row_scales_.ApplyPow(0.5);  // L2norm = sqrt(sum_of_cols(x^2)),
      row_scales_.InvertElements();  // 1/L2norm,
      // compute the output,
      out->CopyFromMat(in);
      out->MulRowsVec(row_scales_);  // re-normalize,
    }
  
    void BackpropagateFnc(const CuMatrixBase<BaseFloat> &in,
                          const CuMatrixBase<BaseFloat> &out,
                          const CuMatrixBase<BaseFloat> &out_diff,
                          CuMatrixBase<BaseFloat> *in_diff) {
      in_diff->CopyFromMat(out_diff);
      in_diff->MulRowsVec(row_scales_);  // diff_by_x(s * x) = s,
    }
  
   private:
    CuMatrix<BaseFloat> l2_aux_;  ///< auxiliary matrix for L2 norm computation,
    CuVector<BaseFloat> row_scales_;  ///< normalization scale of each row,
  };
  
  
  /**
   * Adds shift to all the lines of the matrix
   * (can be used for global mean normalization)
   */
  class AddShift : public UpdatableComponent {
   public:
    AddShift(int32 dim_in, int32 dim_out):
      UpdatableComponent(dim_in, dim_out),
      shift_data_(dim_in)
    { }
  
    ~AddShift()
    { }
  
    Component* Copy() const { return new AddShift(*this); }
    ComponentType GetType() const { return kAddShift; }
  
    void InitData(std::istream &is) {
      // define options
      float init_param = 0.0;
      // parse config
      std::string token;
      while (is >> std::ws, !is.eof()) {
        ReadToken(is, false, &token);
        /**/ if (token == "<InitParam>") ReadBasicType(is, false, &init_param);
        else if (token == "<LearnRateCoef>") ReadBasicType(is, false, &learn_rate_coef_);
        else KALDI_ERR << "Unknown token " << token << ", a typo in config?"
                       << " (InitParam)";
      }
      // initialize
      shift_data_.Resize(InputDim(), kSetZero);  // set to zero
      shift_data_.Set(init_param);
    }
  
    void ReadData(std::istream &is, bool binary) {
      // optional learning-rate coef,
      if ('<' == Peek(is, binary)) {
        ExpectToken(is, binary, "<LearnRateCoef>");
        ReadBasicType(is, binary, &learn_rate_coef_);
      }
      // read the shift data
      shift_data_.Read(is, binary);
    }
  
    void WriteData(std::ostream &os, bool binary) const {
      WriteToken(os, binary, "<LearnRateCoef>");
      WriteBasicType(os, binary, learn_rate_coef_);
      shift_data_.Write(os, binary);
    }
  
    int32 NumParams() const { return shift_data_.Dim(); }
  
    void GetGradient(VectorBase<BaseFloat>* gradient) const {
      KALDI_ASSERT(gradient->Dim() == NumParams());
      shift_data_grad_.CopyToVec(gradient);
    }
  
    void GetParams(VectorBase<BaseFloat>* params) const {
      KALDI_ASSERT(params->Dim() == NumParams());
      shift_data_.CopyToVec(params);
    }
  
    void SetParams(const VectorBase<BaseFloat>& params) {
      KALDI_ASSERT(params.Dim() == NumParams());
      shift_data_.CopyFromVec(params);
    }
  
    std::string Info() const {
      return std::string("
    shift_data") +
        MomentStatistics(shift_data_) +
        ", lr-coef " + ToString(learn_rate_coef_);
    }
  
    std::string InfoGradient() const {
      return std::string("
    shift_data_grad") +
        MomentStatistics(shift_data_grad_) +
        ", lr-coef " + ToString(learn_rate_coef_);
    }
  
    void PropagateFnc(const CuMatrixBase<BaseFloat> &in,
                      CuMatrixBase<BaseFloat> *out) {
      // copy, add the shift,
      out->CopyFromMat(in);
      out->AddVecToRows(1.0, shift_data_, 1.0);
    }
  
    void BackpropagateFnc(const CuMatrixBase<BaseFloat> &in,
                          const CuMatrixBase<BaseFloat> &out,
                          const CuMatrixBase<BaseFloat> &out_diff,
                          CuMatrixBase<BaseFloat> *in_diff) {
      // the derivative of additive constant is zero...
      in_diff->CopyFromMat(out_diff);
    }
  
    void Update(const CuMatrixBase<BaseFloat> &input,
                const CuMatrixBase<BaseFloat> &diff) {
      // we use following hyperparameters from the option class,
      const BaseFloat lr = opts_.learn_rate;
      // gradient,
      shift_data_grad_.Resize(InputDim(), kSetZero);  // reset to zero,
      shift_data_grad_.AddRowSumMat(1.0, diff, 0.0);
      // update,
      shift_data_.AddVec(-lr * learn_rate_coef_, shift_data_grad_);
    }
  
    void SetLearnRateCoef(BaseFloat c) { learn_rate_coef_ = c; }
  
   protected:
    CuVector<BaseFloat> shift_data_;
    CuVector<BaseFloat> shift_data_grad_;
  };
  
  
  /**
   * Rescale the data column-wise by a vector
   * (can be used for global variance normalization)
   */
  class Rescale : public UpdatableComponent {
   public:
    Rescale(int32 dim_in, int32 dim_out):
      UpdatableComponent(dim_in, dim_out),
      scale_data_(dim_in)
    { }
  
    ~Rescale()
    { }
  
    Component* Copy() const { return new Rescale(*this); }
    ComponentType GetType() const { return kRescale; }
  
    void InitData(std::istream &is) {
      // define options
      float init_param = 0.0;
      // parse config
      std::string token;
      while (is >> std::ws, !is.eof()) {
        ReadToken(is, false, &token);
        /**/ if (token == "<InitParam>") ReadBasicType(is, false, &init_param);
        else if (token == "<LearnRateCoef>") ReadBasicType(is, false, &learn_rate_coef_);
        else KALDI_ERR << "Unknown token " << token << ", a typo in config?"
                       << " (InitParam)";
      }
      // initialize
      scale_data_.Resize(InputDim(), kSetZero);
      scale_data_.Set(init_param);
    }
  
    void ReadData(std::istream &is, bool binary) {
      // optional learning-rate coef,
      if ('<' == Peek(is, binary)) {
        ExpectToken(is, binary, "<LearnRateCoef>");
        ReadBasicType(is, binary, &learn_rate_coef_);
      }
      // read the shift data
      scale_data_.Read(is, binary);
    }
  
    void WriteData(std::ostream &os, bool binary) const {
      WriteToken(os, binary, "<LearnRateCoef>");
      WriteBasicType(os, binary, learn_rate_coef_);
      scale_data_.Write(os, binary);
    }
  
    int32 NumParams() const { return scale_data_.Dim(); }
  
    void GetGradient(VectorBase<BaseFloat>* gradient) const {
      KALDI_ASSERT(gradient->Dim() == NumParams());
      scale_data_grad_.CopyToVec(gradient);
    }
  
    void GetParams(VectorBase<BaseFloat>* params) const {
      KALDI_ASSERT(params->Dim() == NumParams());
      scale_data_.CopyToVec(params);
    }
  
    void SetParams(const VectorBase<BaseFloat>& params) {
      KALDI_ASSERT(params.Dim() == NumParams());
      scale_data_.CopyFromVec(params);
    }
  
    std::string Info() const {
      return std::string("
    scale_data") +
        MomentStatistics(scale_data_) +
        ", lr-coef " + ToString(learn_rate_coef_);
    }
  
    std::string InfoGradient() const {
      return std::string("
    scale_data_grad") +
        MomentStatistics(scale_data_grad_) +
        ", lr-coef " + ToString(learn_rate_coef_);
    }
  
    void PropagateFnc(const CuMatrixBase<BaseFloat> &in,
                      CuMatrixBase<BaseFloat> *out) {
      // copy, rescale the data,
      out->CopyFromMat(in);
      out->MulColsVec(scale_data_);
    }
  
    void BackpropagateFnc(const CuMatrixBase<BaseFloat> &in,
                          const CuMatrixBase<BaseFloat> &out,
                          const CuMatrixBase<BaseFloat> &out_diff,
                          CuMatrixBase<BaseFloat> *in_diff) {
      // derivatives are scaled with the scale_data_,
      in_diff->CopyFromMat(out_diff);
      in_diff->MulColsVec(scale_data_);
    }
  
    void Update(const CuMatrixBase<BaseFloat> &input,
                const CuMatrixBase<BaseFloat> &diff) {
      // we use following hyperparameters from the option class,
      const BaseFloat lr = opts_.learn_rate;
      // gradient,
      scale_data_grad_.Resize(InputDim(), kSetZero);  // reset,
      CuMatrix<BaseFloat> gradient_aux(diff);
      gradient_aux.MulElements(input);
      scale_data_grad_.AddRowSumMat(1.0, gradient_aux, 0.0);
      // update,
      scale_data_.AddVec(-lr * learn_rate_coef_, scale_data_grad_);
    }
  
    void SetLearnRateCoef(BaseFloat c) { learn_rate_coef_ = c; }
  
   protected:
    CuVector<BaseFloat> scale_data_;
    CuVector<BaseFloat> scale_data_grad_;
  };
  
  }  // namespace nnet1
  }  // namespace kaldi
  
  #endif  // KALDI_NNET_NNET_VARIOUS_H_