/* 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 "paddle/operators/cond_op.h"

#include <cstring>
#include <sstream>

#include "paddle/framework/op_registry.h"
#include "paddle/operators/gather.h"
#include "paddle/operators/net_op.h"
#include "paddle/operators/scatter.h"

namespace paddle {
namespace operators {

using Scope = framework::Scope;
using Variable = framework::Variable;
using Tensor = framework::Tensor;
using DDim = framework::DDim;

void CondOp::CreateScope(const Scope& scope) const {
  auto sub_scopes_var = scope.FindVar("SubScopes");
  PADDLE_ENFORCE(sub_scopes_var != nullptr, "");
  auto sub_scopes = sub_scopes_var->GetMutable<std::vector<Scope*>>();
  auto& sub_scope = scope.NewScope();
  sub_scopes->push_back(&sub_scope);
}

void CondOp::CreateIndexTensor(const Scope& scope) const {
  auto index_tensors_var = scope.FindVar("IndexTensors");
  PADDLE_ENFORCE(index_tensors_var != nullptr, "");
  auto& index_tensors = *index_tensors_var->GetMutable<std::vector<Tensor>>();
  index_tensors.push_back(Tensor());
}

void CondOp::InferShape(const Scope& scope) const {
  auto sub_scopes_var = scope.FindVar("SubScopes");
  PADDLE_ENFORCE_NOT_NULL(sub_scopes_var);
  auto& sub_scopes = *sub_scopes_var->GetMutable<std::vector<Scope*>>();

  for (int i = 0; i < 2; ++i) {
    // Create two sub scopes for true and false branches
    // sub_scopes[0] for the true branch and sub_scopes[1] for the false
    // branch
    CreateScope(scope);

    // Create two tensors for true and false indices
    // index_tensors[0] for the true branch and index_tensors[1] for the false
    // branch
    CreateIndexTensor(scope);

    PADDLE_ENFORCE(!Inputs("Xs").empty(), "Inputs can't be empty");
    for (auto& input : Inputs("Xs")) {
      // Create a new tensor in sub-scope for input-type tensor
      Variable* v = sub_scopes[i]->NewVar(input);
      Tensor* sub_input = v->GetMutable<Tensor>();
      sub_input->Resize(scope.FindVar(input)->GetMutable<Tensor>()->dims());
    }

    for (auto& output : (*sub_net_op_[i]).Outputs()) {
      for (auto& var_name : output.second) {
        sub_scopes[i]->NewVar(var_name);
      }
    }

    // each net calls InferShape
    sub_net_op_[i]->InferShape(*sub_scopes[i]);
  }

  for (auto& output : Outputs("Outs")) {
    Tensor* tensor_t_out = sub_scopes[0]->FindVar(output)->GetMutable<Tensor>();
    PADDLE_ENFORCE_NOT_NULL(tensor_t_out, "True output should be NULL");
    Tensor* tensor_f_out = sub_scopes[1]->FindVar(output)->GetMutable<Tensor>();
    PADDLE_ENFORCE_NOT_NULL(tensor_f_out, "True output should be NULL");

    auto* tensor_out_var = scope.FindVar(output);
    PADDLE_ENFORCE_NOT_NULL(tensor_out_var, "Output not found");
    Tensor* tensor_out = tensor_out_var->GetMutable<Tensor>();
    PADDLE_ENFORCE_NOT_NULL(tensor_t_out, "True output should be NULL");
    // check output size should be same
    PADDLE_ENFORCE_EQ(tensor_t_out->dims(), tensor_f_out->dims(),
                      "Outputs not of the same shape");
    tensor_out->Resize(tensor_t_out->dims());
    tensor_out->mutable_data<float>(tensor_out->dims(), platform::CPUPlace());
  }
}

void CondOp::Run(const Scope& scope,
                 const platform::DeviceContext& dev_ctx) const {
  auto sub_scopes = scope.FindVar("SubScopes")->Get<std::vector<Scope*>>();
  auto index_tensors =
      scope.FindVar("IndexTensors")->Get<std::vector<Tensor>>();

  std::string cond_name = Input("Cond");
  Variable* cond_var = scope.FindVar(cond_name);
  PADDLE_ENFORCE_NOT_NULL(cond_var);
  const Tensor* cond = cond_var->GetMutable<Tensor>();

  // Step 1: get the true/false index at runtime
  // index_[0]: vector<int>, contains all index for cond[i] == true
  // index_[1]: vector<int>, contains all index for cond[i] == false
  for (int i = 0; i < 2; ++i) index_[i].clear();

  const int* cond_data = cond->data<int>();
  for (int i = 0; i < cond->dims()[0]; ++i) {
    if (cond_data[i])
      index_[0].push_back(i);
    else
      index_[1].push_back(i);
  }

  // put index_[0] and index_[1] into two tensors:
  // index_tensor_[0] and index_tensor_[1]
  DDim dim = paddle::framework::make_ddim({0});
  for (int i = 0; i < 2; ++i) {
    dim[0] = index_[i].size();
    int* tmp_ptr =
        index_tensors[i].mutable_data<int>(dim, platform::CPUPlace());
    index_tensors[i].Resize(dim);
    memcpy(tmp_ptr, index_[i].data(), dim[0] * sizeof(int));
  }

  // Step 2: collect data by calling gather
  for (int i = 0; i < 2; ++i) {
    // i= 0/i for True and False branches respectively
    for (auto& input : Inputs("Xs")) {
      // find Tensor
      Variable* v = scope.FindVar(input);
      PADDLE_ENFORCE_NOT_NULL(v);
      Tensor* tensor_parent = v->GetMutable<Tensor>();

      v = sub_scopes[i]->FindVar(input);
      PADDLE_ENFORCE_NOT_NULL(v);
      Tensor* tensor_child = v->GetMutable<Tensor>();

      // Resize child
      DDim dim = tensor_child->dims();
      dim[0] = index_[i].size();
      tensor_child->Resize(dim);
      tensor_child->mutable_data<float>(dim, platform::CPUPlace());

      Gather<float>(dev_ctx.GetPlace(), tensor_parent, &index_tensors[i],
                    tensor_child);
    }
  }

  // Step 3: run
  for (int i = 0; i < 2; ++i) sub_net_op_[i]->Run(*sub_scopes[i], dev_ctx);

  // Step 4: merge output results
  for (int i = 0; i < 2; ++i) {
    // i= 0/i for True and False branches respectively
    for (auto& output : Outputs("Outs")) {
      // find Tensor
      Variable* v = scope.FindVar(output);
      PADDLE_ENFORCE_NOT_NULL(v);
      Tensor* tensor_parent = v->GetMutable<Tensor>();

      v = sub_scopes[i]->FindVar(output);
      PADDLE_ENFORCE_NOT_NULL(v);
      Tensor* tensor_child = v->GetMutable<Tensor>();

      ScatterUpdate<float>(dev_ctx.GetPlace(), tensor_child, &index_tensors[i],
                           tensor_parent);
    }
  }
}

class CondOpProtoAndCheckerMaker : public framework::OpProtoAndCheckerMaker {
 public:
  CondOpProtoAndCheckerMaker(framework::OpProto* proto,
                             framework::OpAttrChecker* op_checker)
      : OpProtoAndCheckerMaker(proto, op_checker) {
    AddInput("Cond", "The condition, which is a bool vector");
    AddInput("Xs", "Inputs of Subnets").AsDuplicable();
    AddOutput("Outs", "Outputs of Cond_Op after merge").AsDuplicable();

    AddOutput("SubScopes", "sub scopes for true and false branches");
    AddOutput("IndexTensors", "Index Tensors contains indices for true/false");

    AddComment(R"DOC(
Sample dependent Cond Operator:
The equation is: Out[i] = subnet_t[i], if Cond[i] == true
Out[i] = subnet_t[i], if Cond[i] == false
)DOC");
  }
};

}  // namespace operators
}  // namespace paddle

REGISTER_OP_WITHOUT_GRADIENT(cond, paddle::operators::CondOp,
                             paddle::operators::CondOpProtoAndCheckerMaker);