/* 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 calSelectedCols(const MatrixPtr scores,
                       const int* seqStartPos,
                       const int* subSeqStartPos);
  void buildOutputSeqInfo();

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

  MatrixPtr scoreOverInputSeq_;

  // 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::calSelectedCols(const MatrixPtr selected_indices,
                                             const int* seqStartPos,
                                             const int* subSeqStartPos) {
  selectedRows_.clear();
  outSubSeqStartInfo_.resize(1, 0);
  outSeqStartInfo_.resize(1, 0);
}

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);
  const MatrixPtr selected_indices = getInputValue(1);
  CHECK(inputSeq.hasSubseq()) << "The first input of SubNestSequence layer "
                              << "must be a nested sequence.";
  CHECK_EQ(inputSeq.getNumSequences(), selected_indices->getHeight());

  calSelectedCols(selected_indices,
                  inputSeq.sequenceStartPositions->getMutableData(false),
                  inputSeq.subSequenceStartPositions->getMutableData(false));

  resetOutput(selectedRows_.size(), getSize());
  buildOutputSeqInfo();

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

  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