//   Copyright (c) 2020 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.

#pragma once
#include <memory>
#include <string>
#include <vector>
#include "paddle/fluid/framework/fleet/fleet_wrapper.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/tensor.h"

namespace paddle {
namespace operators {

template <typename T>
void PullSparseFunctor(const framework::ExecutionContext& ctx) {
  auto inputs = ctx.MultiInput<framework::LoDTensor>("Ids");
  auto outputs = ctx.MultiOutput<framework::LoDTensor>("Out");
  uint32_t fea_dim = static_cast<uint32_t>(ctx.Attr<int>("EmbeddingDim"));
  uint64_t padding_id = static_cast<uint64_t>(ctx.Attr<int>("PaddingId"));
  auto table_id = static_cast<uint32_t>(ctx.Attr<int>("TableId"));
  // note: GetInstance() is not thread-safe
  // we assume FleetWrapper has been already initialized
  auto fleet_ptr = framework::FleetWrapper::GetInstance();
  fleet_ptr->PullSparseToTensorSync(table_id, fea_dim, padding_id,
                                    ctx.GetPlace(), &inputs, &outputs);
}

template <typename T>
void PushSparseFunctor(const framework::ExecutionContext& ctx) {
  auto inputs = ctx.MultiInput<framework::LoDTensor>("Ids");
  auto grads =
      ctx.MultiInput<framework::LoDTensor>(framework::GradVarName("Out"));
  uint32_t fea_dim = static_cast<uint32_t>(ctx.Attr<int>("EmbeddingDim"));
  std::string accesor = ctx.Attr<std::string>("AccessorClass");
  bool scale_sparse = ctx.Attr<bool>("ScaleSparseGrad");
  uint64_t padding_id = static_cast<uint64_t>(ctx.Attr<int>("PaddingId"));
  const std::string& label_name = ctx.Attr<std::string>("CtrLabelName");
  const framework::Scope& scope = ctx.scope();
  auto input_names = ctx.Attr<std::vector<std::string>>("InputNames");
  auto table_id = static_cast<uint32_t>(ctx.Attr<int>("TableId"));
  // note: GetInstance() is not thread-safe
  // we assume FleetWrapper has been already initialized
  auto fleet_ptr = framework::FleetWrapper::GetInstance();
  fleet_ptr->PushSparseFromTensorWithLabelAsync(
      scope, table_id, fea_dim, padding_id, scale_sparse, accesor, label_name,
      ctx.GetPlace(), input_names, &inputs, &grads);
}

template <typename T>
class PullSparseCPUKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
    PullSparseFunctor<T>(ctx);
  }
};

template <typename T>
class PushSparseCPUKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
    PushSparseFunctor<T>(ctx);
  }
};
}  // namespace operators
}  // namespace paddle