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

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 <algorithm>
#include <map>
#include <string>
#include <vector>

#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/operators/beam_search_op.h"

namespace paddle {
namespace operators {

void BeamSearch::operator()(const framework::LoDTensor &pre_ids,
                            const framework::LoDTensor &pre_scores,
                            framework::LoDTensor *selected_ids,
                            framework::LoDTensor *selected_scores) {
  auto abs_lod = framework::ToAbsOffset(ids_->lod());
  auto &high_level = abs_lod[lod_level_];

  auto items = SelectTopBeamSizeItems(pre_ids, pre_scores);
  auto selected_items = ToMap(items, high_level.back());
  VLOG(3) << "selected_items:";
  for (size_t i = 0; i < selected_items.size(); ++i) {
    VLOG(3) << "offset:" << i;
    for (auto &item : selected_items[i]) {
      VLOG(3) << ItemToString(item);
    }
  }

  PruneEndBeams(pre_ids, &selected_items);
  // calculate the output tensor's height
  size_t num_instances = std::accumulate(
      std::begin(selected_items), std::end(selected_items), 0,
      [](size_t a, std::vector<Item> &b) { return a + b.size(); });
  // the output tensor shape should be [num_instances, 1]
  auto dims = framework::make_ddim(
      std::vector<int64_t>({static_cast<int>(num_instances), 1}));
  selected_ids->Resize(dims);
  selected_scores->Resize(dims);

  std::map<size_t /*offset*/, std::vector<Item>> hash;
  framework::LoD new_lod;
  auto *ids_data = selected_ids->mutable_data<int64_t>(platform::CPUPlace());
  auto *scores_data =
      selected_scores->mutable_data<float>(platform::CPUPlace());

  // fill in data
  std::vector<size_t> low_level;
  size_t low_offset = 0;
  for (auto &items : selected_items) {
    low_level.push_back(low_offset);
    for (auto &item : items) {
      ids_data[low_offset] = item.id;
      scores_data[low_offset] = item.score;
      low_offset++;
    }
  }
  low_level.push_back(low_offset);

  // fill lod
  framework::LoD lod(2);
  lod[0].assign(high_level.begin(), high_level.end());
  lod[1].assign(low_level.begin(), low_level.end());
  if (!framework::CheckLoD(lod)) {
    PADDLE_THROW("lod %s is not right", framework::LoDToString(lod));
  }
  selected_ids->set_lod(lod);
  selected_scores->set_lod(lod);
}

void BeamSearch::PruneEndBeams(const framework::LoDTensor &pre_ids,
                               std::vector<std::vector<Item>> *items) {
  auto *pre_ids_data = pre_ids.data<int64_t>();
  auto abs_lod = framework::ToAbsOffset(ids_->lod());
  auto &high_level = abs_lod[lod_level_];
  for (size_t src_idx = 0; src_idx < high_level.size(); ++src_idx) {
    size_t src_prefix_start = high_level[src_idx];
    size_t src_prefix_end = high_level[src_idx + 1];
    bool finish_flag = true;
    for (size_t offset = src_prefix_start; offset < src_prefix_end; offset++) {
      for (auto &item : items->at(offset)) {
        if (item.id != static_cast<size_t>(end_id_) ||
            pre_ids_data[offset] != end_id_) {
          finish_flag = false;
          break;
        }
      }
      if (!finish_flag) break;
    }
    if (finish_flag) {  // all branchs of the beam (source sentence) end and
                        // prune this beam
      for (size_t offset = src_prefix_start; offset < src_prefix_end; offset++)
        items->at(offset).clear();
    }
  }
}

std::vector<std::vector<BeamSearch::Item>> BeamSearch::ToMap(
    const std::vector<std::vector<Item>> &items, size_t element_num) {
  std::vector<std::vector<Item>> result;
  result.resize(element_num);
  for (auto &entries : items) {
    for (const auto &item : entries) {
      result[item.offset].push_back(item);
    }
  }
  return result;
}

std::vector<std::vector<BeamSearch::Item>> BeamSearch::SelectTopBeamSizeItems(
    const framework::LoDTensor &pre_ids,
    const framework::LoDTensor &pre_scores) {
  std::vector<std::vector<Item>> result;
  std::vector<Item> items;
  // for each source sentence, select the top beam_size items across all
  // candidate sets.
  while (NextItemSet(pre_ids, pre_scores, &items)) {
    std::nth_element(std::begin(items), std::begin(items) + beam_size_,
                     std::end(items), [](const Item &a, const Item &b) {
                       // TODO(superjom) make score's comparation customizable.
                       // partial sort in descending order
                       return a.score > b.score;
                     });
    // prune the top beam_size items.
    if (items.size() > beam_size_) {
      items.resize(beam_size_);
    }
    result.emplace_back(items);
  }
  VLOG(3) << "SelectTopBeamSizeItems result size " << result.size();
  for (auto &items : result) {
    VLOG(3) << "item set:";
    for (auto &item : items) {
      VLOG(3) << ItemToString(item);
    }
  }

  return result;
}

// the candidates of a source
bool BeamSearch::NextItemSet(const framework::LoDTensor &pre_ids,
                             const framework::LoDTensor &pre_scores,
                             std::vector<BeamSearch::Item> *items) {
  if (sent_offset_ >= ids_->NumElements(lod_level_)) {
    return false;
  }
  // find the current candidates
  auto ids = *ids_;
  auto scores = *scores_;

  auto abs_lod = framework::ToAbsOffset(ids.lod());

  auto *ids_data = ids.data<int64_t>();
  auto *scores_data = scores.data<float>();

  size_t instance_dim = 1;
  for (int i = 1; i < ids.dims().size(); i++) {
    instance_dim *= ids.dims()[i];
  }

  auto *pre_ids_data = pre_ids.data<int64_t>();
  auto *pre_scores_data = pre_scores.data<float>();
  items->clear();
  items->reserve(framework::product(ids.dims()));
  for (size_t offset = abs_lod[lod_level_][sent_offset_];
       offset < abs_lod[lod_level_][sent_offset_ + 1]; offset++) {
    auto pre_id = pre_ids_data[offset];
    auto pre_score = pre_scores_data[offset];
    if (pre_id == end_id_) {
      // Allocate all probability mass to eos_id for finished branchs and the
      // other candidate ids can be ignored.
      items->emplace_back(offset, end_id_, pre_score);
    } else {
      for (size_t d = 0; d < instance_dim; d++) {
        const size_t dim_offset = offset * instance_dim + d;
        items->emplace_back(offset, ids_data[dim_offset],
                            scores_data[dim_offset]);
      }
    }
  }

  sent_offset_++;
  return true;
}

std::ostream &operator<<(std::ostream &os, const BeamSearch::Item &item) {
  os << "{";
  os << "offset: " << item.offset << ", ";
  os << "id: " << item.id << ", ";
  os << "score: " << item.score << "";
  os << "}";

  return os;
}

std::string ItemToString(const BeamSearch::Item &item) {
  std::ostringstream stream;
  stream << item;
  return stream.str();
}

class BeamSearchOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
  void Make() override {
    // inputs and outputs stored in proto
    AddInput("pre_ids", "ids in the previous step");
    AddInput("pre_scores", "accumulated scores in the previous step");
    AddInput("ids", "a LoDTensor of shape of [None,k]");
    AddInput("scores",
             "a LoDTensor that has the same shape and LoD with `ids`");
    AddOutput("selected_ids",
              "a LoDTensor that stores the IDs selected by beam search");
    AddOutput(
        "selected_scores",
        "a LoDTensor that has the same shape and LoD with `selected_ids`");

    // Attributes stored in AttributeMap
    AddAttr<int>("level", "the level of LoDTensor");
    AddAttr<int>("beam_size", "beam size for beam search");
    AddAttr<int>("end_id",
                 "the token id which indicates the end of a sequence");

    AddComment(
        "This is a beam search operator that help to generate sequences.");
  }
};

class BeamSearchOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

 protected:
  void InferShape(framework::InferShapeContext *ctx) const override {
    for (const std::string &arg :
         std::vector<std::string>({"pre_ids", "ids", "scores"})) {
      PADDLE_ENFORCE(ctx->HasInput(arg), "BeamSearch need input argument '%s'",
                     arg);
    }
    for (const std::string &arg :
         std::vector<std::string>({"selected_ids", "selected_scores"})) {
      PADDLE_ENFORCE(ctx->HasOutput(arg),
                     "BeamSearch need output argument '%s'", arg);
    }
  }

  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
    framework::OpKernelType kt = framework::OpKernelType(
        framework::ToDataType(
            ctx.Input<framework::LoDTensor>("pre_ids")->type()),
        platform::CPUPlace());
    return kt;
  }
};

class BeamSearchInferVarType : public framework::VarTypeInference {
 public:
  void operator()(const framework::OpDesc &op_desc,
                  framework::BlockDesc *block) const override {
    for (auto &o : op_desc.Output("selected_ids")) {
      auto &selected_ids = block->FindRecursiveOrCreateVar(o);
      selected_ids.SetType(framework::proto::VarType::LOD_TENSOR);
    }
    for (auto &o : op_desc.Output("selected_scores")) {
      auto &selected_scores = block->FindRecursiveOrCreateVar(o);
      selected_scores.SetType(framework::proto::VarType::LOD_TENSOR);
    }
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

REGISTER_OPERATOR(beam_search, ops::BeamSearchOp, ops::BeamSearchOpMaker,
                  ops::BeamSearchInferVarType);
REGISTER_OP_CPU_KERNEL(
    beam_search,
    ops::BeamSearchOpKernel<paddle::platform::CPUDeviceContext, float>,
    ops::BeamSearchOpKernel<paddle::platform::CPUDeviceContext, double>,
    ops::BeamSearchOpKernel<paddle::platform::CPUDeviceContext, int>,
    ops::BeamSearchOpKernel<paddle::platform::CPUDeviceContext, int64_t>);