/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

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 "Layer.h"
#include "paddle/math/Matrix.h"
#include "paddle/math/Vector.h"
#include "paddle/utils/Logging.h"
#include "paddle/utils/Stat.h"

namespace paddle {

class SubNestedSequenceLayer : public Layer {
public:
  explicit SubNestedSequenceLayer(const LayerConfig& config) : Layer(config) {}

  bool init(const LayerMap& layerMap,
            const ParameterMap& parameterMap) override;

  void forward(PassType passType) override;
  void backward(const UpdateCallback& callback = nullptr) override;

private:
  void reorganizeSeqInfo(const ICpuGpuVectorPtr seqStartPos,
                         const ICpuGpuVectorPtr subSeqStartPos);
  void calSelectedCols(const MatrixPtr selectedIndices,
                       const std::vector<std::vector<int>> inputSeqInfo);
  void buildOutputSeqInfo();

  std::vector<int> outSeqStartInfo_;
  std::vector<int> outSubSeqStartInfo_;

  // if the second input of this layer is on GPU memory, copy it to CPU memory.
  MatrixPtr selIdsCpu_;
  // reorganize sequenceStartPositions and subSequenceStartPositions altogether
  // into a 2d vector to facilitate the sequence selection process.
  std::vector<std::vector<int>> inputSeqInfo_;

  // the final seleted row indices in a batch,
  // rowIdx_ and selectedRows_ actually share a same memory.
  IVectorPtr rowIndice_;
  std::vector<int> selectedRows_;
};

REGISTER_LAYER(sub_nested_seq, SubNestedSequenceLayer);

bool SubNestedSequenceLayer::init(const LayerMap& layerMap,
                                  const ParameterMap& parameterMap) {
  /* Initialize the basic parent class */
  Layer::init(layerMap, parameterMap);
  CHECK_EQ(2U, inputLayers_.size());
  setNeedSequenceInfo(false);
  return true;
}

void SubNestedSequenceLayer::reorganizeSeqInfo(
    const ICpuGpuVectorPtr seqStartPos, const ICpuGpuVectorPtr subSeqStartPos) {
  int* seqStarts = seqStartPos->getMutableData(false);
  int* subSeqStarts = subSeqStartPos->getMutableData(false);

  int seqNum = seqStartPos->getSize() - 1;
  inputSeqInfo_.resize(seqNum, std::vector<int>());
  int seqIdx = 0;
  for (size_t i = 0; i < subSeqStartPos->getSize(); ++i) {
    inputSeqInfo_[seqIdx].push_back(subSeqStarts[i]);
    if (subSeqStarts[i] == seqStarts[seqIdx + 1]) {
      seqIdx++;
      if (seqIdx == seqNum) return;
      inputSeqInfo_[seqIdx].push_back(subSeqStarts[i]);
    }
  }
}

void SubNestedSequenceLayer::calSelectedCols(
    const MatrixPtr selectedIndices,
    const std::vector<std::vector<int>> inputSeqInfo) {
  selectedRows_.clear();
  outSubSeqStartInfo_.resize(1, 0);
  outSeqStartInfo_.resize(1, 0);

  size_t seqNum = selectedIndices->getHeight();
  size_t beamSize = selectedIndices->getWidth();
  for (size_t i = 0; i < seqNum; ++i) {
    for (size_t j = 0; j < beamSize; ++j) {
      if (selectedIndices->getElement(i, j) == -1.) break;
      int selSubSeqIdx = selectedIndices->getElement(i, j);
      CHECK_GT(inputSeqInfo_[i].size() - 1, selSubSeqIdx);

      size_t subSeqLen =
          inputSeqInfo_[i][selSubSeqIdx + 1] - inputSeqInfo_[i][selSubSeqIdx];
      for (size_t k = 0; k < subSeqLen; ++k)
        selectedRows_.push_back(inputSeqInfo_[i][selSubSeqIdx] + k);
      outSubSeqStartInfo_.push_back(outSubSeqStartInfo_.back() + subSeqLen);
    }
    outSeqStartInfo_.push_back(outSubSeqStartInfo_.back());
  }
}

void SubNestedSequenceLayer::buildOutputSeqInfo() {
  Argument& output = getOutput();

  ICpuGpuVector::resizeOrCreate(
      output.sequenceStartPositions, outSeqStartInfo_.size(), false);
  output.sequenceStartPositions->copyFrom(
      outSeqStartInfo_.data(), outSeqStartInfo_.size(), false);

  ICpuGpuVector::resizeOrCreate(
      output.subSequenceStartPositions, outSubSeqStartInfo_.size(), false);
  output.subSequenceStartPositions->copyFrom(
      outSubSeqStartInfo_.data(), outSubSeqStartInfo_.size(), false);
}

void SubNestedSequenceLayer::forward(PassType passType) {
  Layer::forward(passType);

  const Argument& inputSeq = getInput(0);
  CHECK(inputSeq.hasSubseq()) << "The first input of SubNestSequence layer "
                              << "must be a nested sequence.";
  const MatrixPtr selectedIndices = getInputValue(1);
  CHECK_EQ(inputSeq.getNumSequences(), selectedIndices->getHeight());

  if (dynamic_cast<GpuMatrix*>(selectedIndices.get())) {
    /*
     * Currently, the second input for this layer generated by
     * kmax_sequence_score_layer whose output is always stored on CPU,
     * or a data_layer which canbe on GPU.
     *
     * If the second input is on GPU, copy it to CPU memory, because this
     * input always uses very few memory, and operations related to it are
     * all logic control, not computations.
     */
    Matrix::resizeOrCreate(selIdsCpu_,
                           selectedIndices->getHeight(),
                           selectedIndices->getWidth(),
                           false /* trans */,
                           false /* useGpu */);
    selIdsCpu_->copyFrom(*selectedIndices);
  } else {
    selIdsCpu_ = selectedIndices;
  }

  reorganizeSeqInfo(inputSeq.sequenceStartPositions,
                    inputSeq.subSequenceStartPositions);
  calSelectedCols(selIdsCpu_, inputSeqInfo_);
  resetOutput(selectedRows_.size(), getSize());

  if (useGpu_) {
    rowIndice_ = IVector::create(selectedRows_.size(), useGpu_);
    rowIndice_->copyFrom(selectedRows_.data(), selectedRows_.size());
  } else {
    rowIndice_ =
        IVector::create(selectedRows_.data(), selectedRows_.size(), useGpu_);
  }

  buildOutputSeqInfo();
  getOutputValue()->selectRows(*getInputValue(0), *rowIndice_);
}

void SubNestedSequenceLayer::backward(const UpdateCallback& callback) {
  MatrixPtr inputSeqGrad = getInputGrad(0);
  MatrixPtr outputGrad = getOutputGrad();

  if (inputSeqGrad) outputGrad->addToRows(*inputSeqGrad, *rowIndice_);
}

}  // namespace paddle