/* Copyright (c) 2016 Baidu, Inc. 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 "LstmLayer.h"
#include "paddle/math/Matrix.h"
#include "paddle/math/BaseMatrix.h"

namespace paddle {

class CoordIterator {
public:
  std::vector<int> dims_;
  std::vector<bool> directions_;
  std::vector<int> curPos_;
  bool end_;

  void step(size_t d, bool reversed) {
    if (directions_[d] ^ reversed) {
      if (curPos_[d] == dims_[d] - 1) {
        curPos_[d] = 0;
        if (d) {
          step(d - 1, reversed);
        } else {
          end_ = true;
        }
      } else {
        curPos_[d]++;
      }
    } else {
      if (curPos_[d] == 0) {
        curPos_[d] = dims_[d] - 1;
        if (d) {
          step(d - 1, reversed);
        } else {
          end_ = true;
        }
      } else {
        curPos_[d]--;
      }
    }
  }

public:
  CoordIterator(std::vector<int> dim, std::vector<bool> directions)
      : dims_(dim), directions_(directions), end_(false) {
    CHECK_EQ(dims_.size(), directions_.size());
    for (size_t i = 0; i < dims_.size(); i++) {
      curPos_.push_back(-1);
    }
  }
  CoordIterator& operator++() {
    step(dims_.size() - 1, false);
    return *this;
  }

  CoordIterator& operator--() {
    step(dims_.size() - 1, true);
    return *this;
  }

  std::vector<int>& curPos() { return curPos_; }

  int offset() {
    int offset = curPos_[0];
    for (size_t i = 1; i < dims_.size(); i++) {
      offset = offset * dims_[i] + curPos_[i];
    }
    return offset;
  }

  int offset(const std::vector<int>& pos) {
    int offset = pos[0];
    for (size_t i = 1; i < dims_.size(); i++) {
      offset = offset * dims_[i] + pos[i];
    }
    return offset;
  }

  std::vector<int>& begin() {
    for (size_t i = 0; i < dims_.size(); i++) {
      curPos_[i] = directions_[i] ? 0 : dims_[i] - 1;
    }
    end_ = false;
    return curPos_;
  }

  std::vector<int>& rbegin() {
    for (size_t i = 0; i < dims_.size(); i++) {
      curPos_[i] = directions_[i] ? dims_[i] - 1 : 0;
    }
    end_ = false;
    return curPos_;
  }

  bool end() { return end_; }

  bool getPrePos(const std::vector<int>& delays,
                 int idx,
                 std::vector<int>& prePos) {
    bool isAvial = true;
    prePos.clear();
    prePos.reserve(directions_.size());
    for (size_t i = 0; i < directions_.size(); i++) {
      if (int(i) == idx) {
        prePos.push_back(curPos_[i] + delays[i] * (directions_[i] ? 1 : -1));
        if (prePos[i] < 0) {
          prePos[i] = 0;
          isAvial = false;
        }
        if (prePos[i] >= dims_[i]) {
          prePos[i] = dims_[i] - 1;
          isAvial = false;
        }
      } else {
        prePos.push_back(curPos_[i]);
      }
    }
    return isAvial;
  }

  bool getNextPos(const std::vector<int>& delays,
                  int idx,
                  std::vector<int>& nextPos) {
    bool isAvial = true;
    nextPos.clear();
    nextPos.reserve(directions_.size());
    for (size_t i = 0; i < directions_.size(); i++) {
      if (int(i) == idx) {
        nextPos.push_back(curPos_[i] - delays[i] * (directions_[i] ? 1 : -1));
        if (nextPos[i] < 0) {
          nextPos[i] = 0;
          isAvial = false;
        }
        if (nextPos[i] >= dims_[i]) {
          nextPos[i] = dims_[i] - 1;
          isAvial = false;
        }
      } else {
        nextPos.push_back(curPos_[i]);
      }
    }
    return isAvial;
  }
};
/*
 * MDLstmLayer takes 1 input layer with size * (3+numDims).
 * For each sequence [start, end] it performs the following computation:
 * out_i = actState(state_i) * actGate(outputGate_i)
 *
 * For example the image with 2 dims, we take the scanning order from left-top
 * to right-bottom, then the 2 previous states of the current pixels are the
 * ones located at left and top. And each of them has a independent forget gate.
 *
 * state_i = actInput(input_i) * actGate(inputGate_i) +
 *           \sum{j}(actGate(forgetGate_i_j) * state_prev_i_j)
 *
 * inputGate = input_i * inputW + \sum{j}(output_prev_i_j * recurrInputW_j) +
 *             \sum{j}(state_prev_i_j * inputCheck_j)
 *
 * ouputGate = input_i * outputW + \sum{j}(output_prev_i_j * recurrOutputW_j) +
 *             state_i * outputCheck
 *
 * forgetGate_j = input_i * forgetW_j + \sum{j}(output_prev_i_j *
 *                recurrForgetW_j) + \sum{j}(state_prev_i_j * forgetCheck_j)
 *
 * IG Layer: (Input, InputGate, ForgetGates, OutputGate) * OutputSize
 * */

class MDLstmLayer : public LstmLayer {
public:
  explicit MDLstmLayer(const LayerConfig& config) : LstmLayer(config) {}

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

  void forward(PassType passType);

  void backward(const UpdateCallback& callback);

protected:
  void forwardOneSequence(int start, CoordIterator& coordIter);
  void backwardOneSequence(int start, CoordIterator& coordIter);
  void forwardGate2OutputSequence(int start, CoordIterator& coordIter);
  void backwardGate2OutputSequence(int start, CoordIterator& coordIter);

protected:
  std::vector<Argument> frameInputGate_;
  std::vector<Argument> frameForgetGate_;
  std::vector<Argument> frameOutputGate_;
  std::vector<Argument> frameInputNode_;
  std::vector<Argument> frameGate_;
  std::vector<Argument> frameState_;
  std::vector<Argument> framePreOutput_;
  std::vector<Argument> frameOutput_;

  // Activation
  std::unique_ptr<ActivationFunction> activationGate_;
  std::unique_ptr<ActivationFunction> activationState_;

  int numDims_;
  size_t numBlocks_;
  std::vector<bool> directions_;
  std::vector<int> delays_;
  std::vector<std::vector<int>> dimsV_;
};

REGISTER_LAYER(mdlstmemory, MDLstmLayer);

bool MDLstmLayer::init(const LayerMap& layerMap,
                       const ParameterMap& parameterMap) {
  if (!Layer::init(layerMap, parameterMap)) return false;
  CHECK_EQ(1U, inputLayers_.size());
  CHECK_EQ(1U, parameters_.size());

  numBlocks_ = getSize();
  numDims_ = config_.directions_size();
  CHECK_EQ(numBlocks_ * numBlocks_ * (3 + numDims_), parameters_[0]->getSize());

  // inode(1), ig(1), fg(numDims_), og(1), peepIg(1), peepFg(numDims_),
  // peepOg(1), then size of localBias_ is 3+numDims_
  CHECK_EQ(numBlocks_ * (5 + 2 * numDims_), biasParameter_->getSize());
  weight_.reset(
      new Weight(numBlocks_, numBlocks_ * (3 + numDims_), parameters_[0]));
  if (biasParameter_.get() != NULL) {
    bias_.reset(new Weight(1, numBlocks_ * (5 + 2 * numDims_), biasParameter_));
    localBias_ = Matrix::create(nullptr,
                                /* height= */ 1,
                                numBlocks_ * (3 + numDims_),
                                /* trans= */ false,
                                useGpu_);
    checkIg_ = Matrix::create(nullptr,
                              /* height= */ 1,
                              numBlocks_,
                              /* trans= */ false,
                              useGpu_);
    checkFg_ = Matrix::create(nullptr,
                              /* height= */ numDims_,
                              numBlocks_,
                              /* trans= */ false,
                              useGpu_);
    checkOg_ = Matrix::create(nullptr,
                              /* height= */ 1,
                              numBlocks_,
                              /* trans= */ false,
                              useGpu_);
    localBiasGrad_ = Matrix::create(nullptr,
                                    /* height= */ 1,
                                    numBlocks_ * (3 + numDims_),
                                    /* trans= */ false,
                                    useGpu_);
    checkIgGrad_ = Matrix::create(nullptr,
                                  /* height= */ 1,
                                  numBlocks_,
                                  /* trans= */ false,
                                  useGpu_);
    checkFgGrad_ = Matrix::create(nullptr,
                                  /* height= */ numDims_,
                                  numBlocks_,
                                  /* trans= */ false,
                                  useGpu_);
    checkOgGrad_ = Matrix::create(nullptr,
                                  /* height= */ 1,
                                  numBlocks_,
                                  /* trans= */ false,
                                  useGpu_);

    localBias_->setData(bias_->getW()->getData());
    checkIg_->setData(bias_->getW()->getData() + numBlocks_ * (3 + numDims_));
    checkFg_->setData(bias_->getW()->getData() + numBlocks_ * (4 + numDims_));
    checkOg_->setData(bias_->getW()->getData() +
                      numBlocks_ * (4 + 2 * numDims_));

    if (bias_->getWGrad()) {
      localBiasGrad_->setData(bias_->getWGrad()->getData());
      checkIgGrad_->setData(bias_->getWGrad()->getData() +
                            numBlocks_ * (3 + numDims_));
      checkFgGrad_->setData(bias_->getWGrad()->getData() +
                            numBlocks_ * (4 + numDims_));
      checkOgGrad_->setData(bias_->getWGrad()->getData() +
                            numBlocks_ * (4 + 2 * numDims_));
    }
  } else {
    LOG(FATAL) << "Bias should be here.";
  }
  for (int i = 0; i < numDims_; i++) {
    directions_.push_back(config_.directions(i));
  }
  for (int i = 0; i < numDims_; i++) {
    delays_.push_back(-1);
  }
  activationGate_.reset(ActivationFunction::create(config_.active_gate_type()));
  activationState_.reset(
      ActivationFunction::create(config_.active_state_type()));

  return true;
}

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

  const Argument& input = getInput(0);
  CHECK(input.sequenceStartPositions);
  int batchSize = input.getBatchSize();
  int numSequences = input.getNumSequences();
  resetOutput(batchSize, numBlocks_);
  CHECK_EQ(numBlocks_ * (3 + numDims_), input.value->getWidth());
  const int* starts = input.sequenceStartPositions->getData(false);
  CHECK_EQ(starts[numSequences], batchSize);

  int* dimsData = input.cpuSequenceDims->getData();
  CHECK_EQ(int(input.cpuSequenceDims->getSize()), numDims_ * numSequences);

  for (int i = 0; i < numSequences; i++) {
    std::vector<int> dims;
    for (int j = 0; j < numDims_; j++) {
      dims.push_back(dimsData[i * numDims_ + j]);
    }
    dimsV_.push_back(dims);
  }

  frameInputGate_.reserve(batchSize);
  frameForgetGate_.reserve(batchSize);
  frameOutputGate_.reserve(batchSize);
  frameInputNode_.reserve(batchSize);
  frameGate_.reserve(batchSize);
  frameState_.reserve(batchSize);
  framePreOutput_.reserve(batchSize);
  frameOutput_.reserve(batchSize);

  Matrix::resizeOrCreate(gate_.value,
                         /* height= */ batchSize,
                         numBlocks_ * (3 + numDims_),
                         /* trans= */ false,
                         useGpu_);

  for (int i = frameGate_.size(); i < batchSize; i++) {
    Argument arg;
    arg.value = Matrix::create(nullptr,
                               /* height= */ 1,
                               numBlocks_ * (3 + numDims_),
                               /* trans= */ false,
                               useGpu_);
    arg.grad = Matrix::create(nullptr,
                              /* height= */ 1,
                              numBlocks_ * (3 + numDims_),
                              /* trans= */ false,
                              useGpu_);
    frameGate_.push_back(arg);
  }
  for (int i = frameInputGate_.size(); i < batchSize; i++) {
    Argument arg;
    arg.value = Matrix::create(nullptr,
                               /* height= */ 1,
                               numBlocks_,
                               /* trans= */ false,
                               useGpu_);
    arg.grad = Matrix::create(nullptr,
                              /* height= */ 1,
                              numBlocks_,
                              /* trans= */ false,
                              useGpu_);
    frameInputGate_.push_back(arg);
  }
  for (int i = frameForgetGate_.size(); i < batchSize; i++) {
    Argument arg;
    arg.value = Matrix::create(nullptr,
                               /* height= */ numDims_,
                               numBlocks_,
                               /* trans= */ false,
                               useGpu_);
    arg.grad = Matrix::create(nullptr,
                              /* height= */ numDims_,
                              numBlocks_,
                              /* trans= */ false,
                              useGpu_);
    frameForgetGate_.push_back(arg);
  }
  for (int i = frameOutputGate_.size(); i < batchSize; i++) {
    Argument arg;
    arg.value = Matrix::create(nullptr,
                               /* height= */ 1,
                               numBlocks_,
                               /* trans= */ false,
                               useGpu_);
    arg.grad = Matrix::create(nullptr,
                              /* height= */ 1,
                              numBlocks_,
                              /* trans= */ false,
                              useGpu_);
    frameOutputGate_.push_back(arg);
  }
  for (int i = frameInputNode_.size(); i < batchSize; i++) {
    Argument arg;
    arg.value = Matrix::create(nullptr,
                               /* height= */ 1,
                               numBlocks_,
                               /* trans= */ false,
                               useGpu_);
    arg.grad = Matrix::create(nullptr,
                              /* height= */ 1,
                              numBlocks_,
                              /* trans= */ false,
                              useGpu_);
    frameInputNode_.push_back(arg);
  }
  for (int i = frameState_.size(); i < batchSize; i++) {
    Argument arg;
    arg.value = Matrix::create(
        /* height= */ 1, numBlocks_, /* trans= */ false, useGpu_);
    frameState_.push_back(arg);
  }
  for (int i = framePreOutput_.size(); i < batchSize; i++) {
    Argument arg;
    arg.value = Matrix::create(
        /* height= */ 1, numBlocks_, /* trans= */ false, useGpu_);
    framePreOutput_.push_back(arg);
  }
  for (int i = frameOutput_.size(); i < batchSize; i++) {
    Argument arg;
    arg.value = Matrix::create(nullptr,
                               /* height= */ 1,
                               numBlocks_,
                               /* trans= */ false,
                               useGpu_);
    arg.grad = Matrix::create(nullptr,
                              /* height= */ 1,
                              numBlocks_,
                              /* trans= */ false,
                              useGpu_);
    frameOutput_.push_back(arg);
  }

  for (int i = 0; i < batchSize; i++) {
    frameOutput_[i].value->setData(output_.value->getData() + i * numBlocks_);
    frameGate_[i].value->setData(gate_.value->getData() +
                                 i * numBlocks_ * (3 + numDims_));
    frameInputNode_[i].value->setData(gate_.value->getData() +
                                      i * numBlocks_ * (3 + numDims_) +
                                      numBlocks_ * 0);
    frameInputGate_[i].value->setData(gate_.value->getData() +
                                      i * numBlocks_ * (3 + numDims_) +
                                      numBlocks_ * 1);
    frameForgetGate_[i].value->setData(gate_.value->getData() +
                                       i * numBlocks_ * (3 + numDims_) +
                                       numBlocks_ * 2);
    frameOutputGate_[i].value->setData(gate_.value->getData() +
                                       i * numBlocks_ * (3 + numDims_) +
                                       numBlocks_ * (2 + numDims_));
  }

  AsyncGpuBlock asyncGpuBlock;
  gate_.value->assign(*input.value);

  if (bias_) {
    gate_.value->addBias(*localBias_, 1);
  }

  for (int i = 0; i < numSequences; i++) {
    CoordIterator coordIter(dimsV_[i], directions_);
    forwardOneSequence(starts[i], coordIter);
  }
}

void MDLstmLayer::forwardGate2OutputSequence(int start,
                                             CoordIterator& coordIter) {
  int idxCurr = start + coordIter.offset();
  std::vector<int> preOffsetV;
  preOffsetV.reserve(numDims_);
  for (int i = 0; i < numDims_; i++) {
    std::vector<int> prePos;
    if (coordIter.getPrePos(delays_, i, prePos)) {
      preOffsetV[i] = coordIter.offset(prePos);
    } else {
      preOffsetV[i] = -1;
    }
  }

  for (int i = 0; i < numDims_; i++) {
    if (preOffsetV[i] >= 0) {
      frameInputGate_[idxCurr].value->addDotMul(
          *frameState_[start + preOffsetV[i]].value, *checkIg_, 1.0, 1.0);

      MatrixPtr fgGateOneDim = Matrix::create(
          frameForgetGate_[idxCurr].value->getData() + i * numBlocks_,
          1,
          numBlocks_,
          false,
          useGpu_);
      MatrixPtr checkFgOneDim =
          Matrix::create(checkFg_->getData() + i * numBlocks_,
                         1.0,
                         numBlocks_,
                         false,
                         useGpu_);
      fgGateOneDim->addDotMul(
          *frameState_[start + preOffsetV[i]].value, *checkFgOneDim, 1.0, 1.0);
    }
  }
  activationGate_->forward(frameInputGate_[idxCurr]);
  activationGate_->forward(frameForgetGate_[idxCurr]);
  activation_->forward(frameInputNode_[idxCurr]);

  frameState_[idxCurr].value->zeroMem();
  for (int i = 0; i < numDims_; i++) {
    if (preOffsetV[i] >= 0) {
      MatrixPtr fgGateOneDim = Matrix::create(
          frameForgetGate_[idxCurr].value->getData() + i * numBlocks_,
          1,
          numBlocks_,
          false,
          useGpu_);
      frameState_[idxCurr].value->addDotMul(
          *frameState_[start + preOffsetV[i]].value, *fgGateOneDim, 1.0, 1.0);
    }
  }
  frameState_[idxCurr].value->addDotMul(*frameInputNode_[idxCurr].value,
                                        *frameInputGate_[idxCurr].value,
                                        1.0,
                                        1.0);

  frameOutputGate_[idxCurr].value->addDotMul(
      *frameState_[idxCurr].value, *checkOg_, 1.0, 1.0);
  activationGate_->forward(frameOutputGate_[idxCurr]);

  framePreOutput_[idxCurr].value->copyFrom(*(frameState_[idxCurr].value));
  activationState_->forward(framePreOutput_[idxCurr]);

  frameOutput_[idxCurr].value->dotMul(*framePreOutput_[idxCurr].value,
                                      *frameOutputGate_[idxCurr].value);
}

void MDLstmLayer::forwardOneSequence(int start, CoordIterator& coordIter) {
  for (coordIter.begin(); !coordIter.end(); ++coordIter) {
    int offset = coordIter.offset();
    for (int i = 0; i < numDims_; i++) {
      std::vector<int> prePos;
      if (coordIter.getPrePos(delays_, i, prePos)) {
        int preOffset = coordIter.offset(prePos);
        frameGate_[start + offset].value->mul(
            frameOutput_[start + preOffset].value, weight_->getW(), 1.0, 1.0);
      }
    }
    forwardGate2OutputSequence(start, coordIter);
  }
}

void MDLstmLayer::backward(const UpdateCallback& callback) {
  const Argument& input = getInput(0);
  CHECK(input.sequenceStartPositions);
  int batchSize = input.getBatchSize();
  const int* starts = input.sequenceStartPositions->getData(false);
  size_t numSequences = input.getNumSequences();

  Matrix::resizeOrCreate(gate_.grad,
                         /* height= */ batchSize,
                         numBlocks_ * (3 + numDims_),
                         /* trans= */ false,
                         useGpu_);

  for (int i = 0; i < batchSize; i++) {
    if (frameState_[i].grad == NULL)
      frameState_[i].grad = Matrix::create(
          /* height= */ 1, numBlocks_, /* trans= */ false, useGpu_);
  }
  for (int i = 0; i < batchSize; i++) {
    if (framePreOutput_[i].grad == NULL)
      framePreOutput_[i].grad = Matrix::create(
          /* height= */ 1, numBlocks_, /* trans= */ false, useGpu_);
  }

  for (int i = 0; i < batchSize; i++) {
    frameOutput_[i].grad->setData(output_.grad->getData() + i * numBlocks_);
    frameGate_[i].grad->setData(gate_.grad->getData() +
                                i * numBlocks_ * (3 + numDims_));
    frameInputNode_[i].grad->setData(gate_.grad->getData() +
                                     i * numBlocks_ * (3 + numDims_) +
                                     numBlocks_ * 0);
    frameInputGate_[i].grad->setData(gate_.grad->getData() +
                                     i * numBlocks_ * (3 + numDims_) +
                                     numBlocks_ * 1);
    frameForgetGate_[i].grad->setData(gate_.grad->getData() +
                                      i * numBlocks_ * (3 + numDims_) +
                                      numBlocks_ * 2);
    frameOutputGate_[i].grad->setData(gate_.grad->getData() +
                                      i * numBlocks_ * (3 + numDims_) +
                                      numBlocks_ * (2 + numDims_));
  }

  {
    AsyncGpuBlock asyncGpuBlock;

    for (size_t i = 0; i < numSequences; i++) {
      CoordIterator coordIter(dimsV_[i], directions_);
      backwardOneSequence(starts[i], coordIter);
    }
  }

  if (input.grad) {
    input.grad->add(*gate_.grad);
  }
  if (bias_ && bias_->getWGrad()) {
    localBiasGrad_->collectBias(*gate_.grad, 1);
    bias_->getParameterPtr()->incUpdate(callback);
  }

  weight_->getParameterPtr()->incUpdate(callback);
}

void MDLstmLayer::backwardGate2OutputSequence(int start,
                                              CoordIterator& coordIter) {
  int idxCurr = start + coordIter.offset();
  std::vector<int> preOffsetV;
  std::vector<int> nextOffsetV;
  preOffsetV.reserve(numDims_);
  nextOffsetV.reserve(numDims_);
  for (int i = 0; i < numDims_; i++) {
    std::vector<int> prePos;
    if (coordIter.getPrePos(delays_, i, prePos)) {
      preOffsetV[i] = coordIter.offset(prePos);
    } else {
      preOffsetV[i] = -1;
    }
    std::vector<int> nextPos;
    if (coordIter.getNextPos(delays_, i, nextPos)) {
      nextOffsetV[i] = coordIter.offset(nextPos);
    } else {
      nextOffsetV[i] = -1;
    }
  }

  framePreOutput_[idxCurr].grad->dotMul(*frameOutput_[idxCurr].grad,
                                        *frameOutputGate_[idxCurr].value);
  activationState_->backward(framePreOutput_[idxCurr]);
  frameState_[idxCurr].grad->copyFrom(*(framePreOutput_[idxCurr].grad));

  frameOutputGate_[idxCurr].grad->dotMul(*frameOutput_[idxCurr].grad,
                                         *framePreOutput_[idxCurr].value);
  activationGate_->backward(frameOutputGate_[idxCurr]);

  frameState_[idxCurr].grad->addDotMul(
      *frameOutputGate_[idxCurr].grad, *checkOg_, 1.0, 1.0);
  for (int i = 0; i < numDims_; i++) {
    if (nextOffsetV[i] >= 0) {
      frameState_[idxCurr].grad->addDotMul(
          *frameInputGate_[start + nextOffsetV[i]].grad, *checkIg_, 1.0, 1.0);

      MatrixPtr fgGateOneDimGrad = Matrix::create(
          frameForgetGate_[start + nextOffsetV[i]].grad->getData() +
              i * numBlocks_,
          1,
          numBlocks_,
          false,
          useGpu_);
      MatrixPtr fgGateOneDimVal = Matrix::create(
          frameForgetGate_[start + nextOffsetV[i]].value->getData() +
              i * numBlocks_,
          1,
          numBlocks_,
          false,
          useGpu_);
      MatrixPtr checkFgOneDim = Matrix::create(
          checkFg_->getData() + i * numBlocks_, 1, numBlocks_, false, useGpu_);

      frameState_[idxCurr].grad->addDotMul(
          *fgGateOneDimGrad, *checkFgOneDim, 1.0, 1.0);
      frameState_[idxCurr].grad->addDotMul(
          *frameState_[start + nextOffsetV[i]].grad,
          *fgGateOneDimVal,
          1.0,
          1.0);
    }
  }

  frameInputNode_[idxCurr].grad->dotMul(*frameState_[idxCurr].grad,
                                        *frameInputGate_[idxCurr].value);
  frameInputGate_[idxCurr].grad->dotMul(*frameState_[idxCurr].grad,
                                        *frameInputNode_[idxCurr].value);

  frameForgetGate_[idxCurr].grad->zeroMem();
  for (int i = 0; i < numDims_; i++) {
    if (preOffsetV[i] >= 0) {
      MatrixPtr fgGateOneDimGrad = Matrix::create(
          frameForgetGate_[idxCurr].grad->getData() + i * numBlocks_,
          1,
          numBlocks_,
          false,
          useGpu_);
      fgGateOneDimGrad->addDotMul(*frameState_[idxCurr].grad,
                                  *frameState_[start + preOffsetV[i]].value,
                                  1.0,
                                  1.0);
    }
  }

  activationGate_->backward(frameInputGate_[idxCurr]);
  activationGate_->backward(frameForgetGate_[idxCurr]);
  activation_->backward(frameInputNode_[idxCurr]);

  if (bias_->getWGrad()) {
    for (int i = 0; i < numDims_; i++) {
      if (preOffsetV[i] >= 0) {
        checkIgGrad_->addDotMul(*frameInputGate_[idxCurr].grad,
                                *frameState_[start + preOffsetV[i]].value,
                                1.0,
                                1.0);

        MatrixPtr fgGateOneDimGrad = Matrix::create(
            frameForgetGate_[idxCurr].grad->getData() + i * numBlocks_,
            1,
            numBlocks_,
            false,
            useGpu_);
        MatrixPtr checkFgOneDimGrad =
            Matrix::create(checkFgGrad_->getData() + i * numBlocks_,
                           1,
                           numBlocks_,
                           false,
                           useGpu_);
        checkFgOneDimGrad->addDotMul(*fgGateOneDimGrad,
                                     *frameState_[start + preOffsetV[i]].value,
                                     1.0,
                                     1.0);
      }
    }
    checkOgGrad_->addDotMul(
        *frameOutputGate_[idxCurr].grad, *frameState_[idxCurr].value, 1.0, 1.0);
  }
}

void MDLstmLayer::backwardOneSequence(int start, CoordIterator& coordIter) {
  MatrixPtr weightT = weight_->getW()->getTranspose();
  for (coordIter.rbegin(); !coordIter.end(); --coordIter) {
    int offset = coordIter.offset();
    backwardGate2OutputSequence(start, coordIter);
    for (int i = 0; i < numDims_; i++) {
      std::vector<int> prePos;
      if (coordIter.getPrePos(delays_, i, prePos)) {
        int preOffset = coordIter.offset(prePos);
        frameOutput_[start + preOffset].grad->mul(
            frameGate_[start + offset].grad, weightT, 1.0, 1.0);
        if (weight_->getWGrad()) {
          weight_->getWGrad()->mul(
              frameOutput_[start + preOffset].value->getTranspose(),
              frameGate_[start + offset].grad,
              1.0,
              1.0);
        }
      }
    }
  }
}

}  // namespace paddle