/**
 * Copyright 2020 Huawei Technologies Co., Ltd
 *
 * 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.
 */

#ifndef MINDSPORE_MINDSPORE_CCSRC_PARALLEL_PS_PARAMETER_SERVER_H_
#define MINDSPORE_MINDSPORE_CCSRC_PARALLEL_PS_PARAMETER_SERVER_H_

#include <unistd.h>
#include <unordered_map>
#include <string>
#include <iostream>
#include <memory>
#include <vector>
#include <mutex>
#include <condition_variable>
#include <thread>
#include <cmath>
#include <random>
#include "ir/func_graph.h"
#include "session/session_basic.h"
#include "session/kernel_graph.h"
#include "session/anf_runtime_algorithm.h"
#include "session/session_factory.h"
#include "parallel/ps/common.h"
#include "parallel/ps/optimizer_info.h"
#include "parallel/ps/optimizer_info_builder.h"
#include "parallel/ps/util.h"
#include "device/cpu/kernel_select_cpu.h"
#include "utils/context/ms_context.h"
#include "kernel/kernel.h"
#include "kernel/ps/pserver_kernel.h"
#include "kernel/cpu/cpu_kernel_factory.h"
#include "kernel/ps/sparse_apply_adam_ps_kernel.h"
#include "kernel/ps/sparse_apply_ftrl_ps_kernel.h"
#include "kernel/ps/apply_momentum_ps_kernel.h"
#include "kernel/ps/embedding_look_up_ps_kernel.h"

namespace mindspore {
namespace parallel {
namespace ps {
using mindspore::kernel::ps::PServerKernel;
template <typename T>
class ParameterServer {
 public:
  static ParameterServer &GetInstance() {
    static ParameterServer instance;
    return instance;
  }

  void Run(const FuncGraphPtr &func_graph);

 private:
  ParameterServer()
      : pserver_num_(0),
        worker_num_(0),
        rank_id_(0),
        grad_accum_count_(0),
        ps_(new ::ps::KVServer<T>(0)),
        handler_(nullptr),
        func_graph_(nullptr),
        kernel_graph_(nullptr),
        sess_(nullptr),
        thread_(nullptr) {}
  ~ParameterServer() = default;
  ParameterServer(const ParameterServer &) = delete;
  ParameterServer &operator=(const ParameterServer &) = delete;

  struct ServerHandler {
    explicit ServerHandler(ParameterServer *ps) : ps_(ps) {}
    void operator()(const ::ps::KVMeta &req_meta, const ::ps::KVPairs<T> &req_data, ::ps::KVServer<T> *server);
    void HandlePushReq(const ::ps::KVMeta &req_meta, const ::ps::KVPairs<T> &req_data);
    void HandlePullReq(const ::ps::KVMeta &req_meta, const ::ps::KVPairs<T> &req_data, ::ps::KVPairs<T> *res);
    void HandleInitWeights(const ::ps::KVPairs<T> &req_data);
    void HandleInitWeightToOptimId(const ::ps::KVPairs<T> &req_data);
    void HandleInitInputsShape(const ::ps::KVPairs<T> &req_data);
    void HandleInitEmbeddings(const ::ps::KVPairs<T> &req_data);
    void HandleEmbeddingLookup(const ::ps::KVMeta &req_meta, const ::ps::KVPairs<T> &req_data, ::ps::KVPairs<T> *res);
    ParameterServer *ps_;
  };

  bool Init(const FuncGraphPtr &func_graph);
  void InitOptimInfoBuilders();
  void InitWeightKeyToOptims(const Key &key, const int &optim_id);
  void InitOptimInputsShape(const Keys &keys, const Values &values, const Lengths &lengths);
  void InitWeight(const Key &key, const WeightPtr &weight);
  void InitGrad(const Key &key, const GradPtr &grad);
  void InitEmbeddingTable(const Key &key,
                          const std::shared_ptr<std::vector<std::shared_ptr<std::vector<size_t>>>> &shapes);
  void UpdateWeights();
  void AccumGrad(const Keys &key, const Values &values, const Lengths &lengths);
  WeightPtr weight(const Key &key);
  void DoEmbeddingLookup(Key key, const LookupIds &lookup_ids, ::ps::KVPairs<T> *res);
  int SumOfShapes(const std::vector<int> &shapes) const;
  size_t PreComputeCapacity(const Keys &keys, const Lengths &lens);
  bool ReadyForUpdateWeights();
  bool ReadyForAccumGrads();
  void ResetGradAccumCount();

  size_t pserver_num_;
  size_t worker_num_;
  size_t rank_id_;
  size_t grad_accum_count_;
  std::unique_ptr<::ps::KVServer<T>> ps_;
  std::unique_ptr<ServerHandler> handler_;
  FuncGraphPtr func_graph_;
  std::shared_ptr<session::KernelGraph> kernel_graph_;
  std::shared_ptr<session::SessionBasic> sess_;

  std::unordered_map<std::string, std::shared_ptr<PServerKernel>> optimizers_;
  std::unordered_map<Key, InputsShapePtr> optim_inputs_shape_;
  std::unordered_map<Key, std::shared_ptr<OptimizerInfo>> optim_infos_;
  std::unordered_map<std::string, std::shared_ptr<OptimizerInfoBuilder>> optim_info_builders_;
  std::unordered_map<Key, std::string> weight_key_to_optims_;
  std::unordered_map<Key, WeightPtr> weights_;
  std::unordered_map<Key, WeightPtr> grads_;
  std::unordered_map<Key, size_t> grads_accum_counter_;
  // std::unordered_map<Key, EmbeddingTablePtr> embeddings_;
  std::unordered_map<Key, std::shared_ptr<PServerKernel>> embedding_lookup_ops_;
  std::unordered_map<Key, size_t> embedding_row_lens_;

  T learning_rate_;
  T momentum_;

  std::mutex mutex_;
  std::condition_variable apply_grads_cv_;
  std::condition_variable accum_grads_cv_;

  std::unique_ptr<std::thread> thread_;

  friend struct ServerHandler;
};

class FuncGraph;
template <typename T>
void ParameterServer<T>::ServerHandler::operator()(const ::ps::KVMeta &req_meta, const ::ps::KVPairs<T> &req_data,
                                                   ::ps::KVServer<T> *server) {
  ::ps::KVPairs<T> res;
  if (req_meta.cmd == kInitWeightsCmd) {
    MS_LOG(ERROR) << "handle init weights cmd" << std::endl;
    HandleInitWeights(req_data);
  } else if (req_meta.cmd == kInitWeightToOptimIdCmd) {
    MS_LOG(ERROR) << "handle init weight optim id mapping cmd" << std::endl;
    HandleInitWeightToOptimId(req_data);
  } else if (req_meta.cmd == kInitOptimInputsShapeCmd) {
    MS_LOG(ERROR) << "handle init inputs shape cmd" << std::endl;
    HandleInitInputsShape(req_data);
  } else if (req_meta.cmd == kInitEmbeddingsCmd) {
    MS_LOG(ERROR) << "handle init embedding cmd" << std::endl;
    HandleInitEmbeddings(req_data);
  } else if (req_meta.cmd == kEmbeddingLookupCmd) {
    MS_LOG(ERROR) << "handle embedding lookup cmd" << std::endl;
    HandleEmbeddingLookup(req_meta, req_data, &res);
  } else if (req_meta.push) {
    MS_LOG(ERROR) << "handle push req cmd" << std::endl;
    HandlePushReq(req_meta, req_data);
  } else {
    MS_LOG(ERROR) << "handle pull req cmd" << std::endl;
    HandlePullReq(req_meta, req_data, &res);
  }
  server->Response(req_meta, res);
}

template <typename T>
void ParameterServer<T>::ServerHandler::HandlePushReq(const ::ps::KVMeta &req_meta, const ::ps::KVPairs<T> &req_data) {
  ps_->AccumGrad(req_data.keys, req_data.vals, req_data.lens);
}

template <typename T>
void ParameterServer<T>::ServerHandler::HandlePullReq(const ::ps::KVMeta &req_meta, const ::ps::KVPairs<T> &req_data,
                                                      ::ps::KVPairs<T> *res) {
  res->keys = req_data.keys;
  ::ps::Key key = req_data.keys[0];
  res->vals = *(ps_->weight(key));
}

template <typename T>
void ParameterServer<T>::ServerHandler::HandleInitWeights(const ::ps::KVPairs<T> &req_data) {
  size_t key_num = req_data.keys.size();
  T *data_ptr = req_data.vals.data();
  size_t pos = 0;
  for (size_t i = 0; i < key_num; i++) {
    Key key = req_data.keys[i];
    size_t data_len = req_data.lens.size() != key_num ? req_data.vals.size() / key_num : req_data.lens[i];

    WeightPtr weight_ptr = std::make_shared<::ps::SArray<T>>();
    weight_ptr->CopyFrom(data_ptr + pos, data_len);
    ps_->InitWeight(key, weight_ptr);

    GradPtr grad_ptr = std::make_shared<::ps::SArray<T>>(data_len, 0);
    ps_->InitGrad(key, grad_ptr);
    pos += data_len;
  }
}

template <typename T>
void ParameterServer<T>::ServerHandler::HandleInitWeightToOptimId(const ::ps::KVPairs<T> &req_data) {
  size_t key_num = req_data.keys.size();
  for (size_t i = 0; i < key_num; i++) {
    Key key = req_data.keys[i];
    T val = req_data.vals[i];
    ps_->InitWeightKeyToOptims(key, val);
  }
}

template <typename T>
void ParameterServer<T>::ServerHandler::HandleInitInputsShape(const ::ps::KVPairs<T> &req_data) {
  ps_->InitOptimInputsShape(req_data.keys, req_data.vals, req_data.lens);
}

template <typename T>
void ParameterServer<T>::ServerHandler::HandleInitEmbeddings(const ::ps::KVPairs<T> &req_data) {
  std::shared_ptr<std::vector<std::shared_ptr<std::vector<size_t>>>> shapes =
    std::make_shared<std::vector<std::shared_ptr<std::vector<size_t>>>>();
  std::shared_ptr<std::vector<size_t>> input_shape = std::make_shared<std::vector<size_t>>();
  std::shared_ptr<std::vector<size_t>> indices_shape = std::make_shared<std::vector<size_t>>();
  std::shared_ptr<std::vector<size_t>> output_shape = std::make_shared<std::vector<size_t>>();
  shapes->push_back(input_shape);
  shapes->push_back(indices_shape);
  shapes->push_back(output_shape);

  const Key &key = req_data.keys[0];
  const Lengths &lens = req_data.lens;
  size_t index = 0;
  for (int i = 0; i < lens[0]; i++) {
    input_shape->push_back(static_cast<size_t>(req_data.vals[index++]));
  }
  for (int j = 0; j < lens[1]; j++) {
    indices_shape->push_back(static_cast<size_t>(req_data.vals[index++]));
  }
  for (int k = 0; k < lens[2]; k++) {
    output_shape->push_back(static_cast<size_t>(req_data.vals[index++]));
  }
  ps_->InitEmbeddingTable(key, shapes);
}

template <typename T>
void ParameterServer<T>::ServerHandler::HandleEmbeddingLookup(const ::ps::KVMeta &req_meta,
                                                              const ::ps::KVPairs<T> &req_data, ::ps::KVPairs<T> *res) {
  const Key &key = req_data.keys[0];
  ps_->DoEmbeddingLookup(key, req_data.vals, res);
  for (size_t i = 0; i < req_data.vals.size(); i++) {
    res->keys->push_back(req_data.vals[i]);
  }
}

template <typename T>
bool ParameterServer<T>::Init(const FuncGraphPtr &func_graph) {
  const char *server_num = getenv(kEnvPServerNum);
  const char *worker_num = getenv(kEnvWorkerNum);
  if (server_num != nullptr) {
    pserver_num_ = *server_num - '0';
  }
  if (worker_num != nullptr) {
    worker_num_ = *worker_num - '0';
  }
  func_graph_ = func_graph;
  rank_id_ = ::ps::MyRank();
  handler_.reset(new ServerHandler(this));

  InitOptimInfoBuilders();

  ps_->set_request_handle(*handler_);
  thread_.reset(new std::thread(&ParameterServer::UpdateWeights, this));
  return true;
}

template <typename T>
void ParameterServer<T>::InitOptimInfoBuilders() {
  std::shared_ptr<OptimizerInfoBuilder> momentum_info_builder = std::make_shared<MomentumOptimInfoBuilder>();
  std::shared_ptr<OptimizerInfoBuilder> sparse_adam_info_builder = std::make_shared<SparseAdamOptimInfoBuilder>();
  std::shared_ptr<OptimizerInfoBuilder> sparse_ftrl_info_builder = std::make_shared<SparseFtrlOptimInfoBuilder>();
  optim_info_builders_[kApplyMomentum] = momentum_info_builder;
  optim_info_builders_[kSparseAdam] = sparse_adam_info_builder;
  optim_info_builders_[kSparseFtrl] = sparse_ftrl_info_builder;
}

template <typename T>
void ParameterServer<T>::InitWeightKeyToOptims(const Key &key, const int &optim_id) {
  if (weight_key_to_optims_.count(key) > 0 || Util::optimizer_name(key) == "") {
    return;
  }
  weight_key_to_optims_[key] = Util::optimizer_name(optim_id);
}

template <typename T>
void ParameterServer<T>::InitOptimInputsShape(const Keys &keys, const Values &values, const Lengths &lengths) {
  InputsShapePtr inputs_shape = std::make_shared<InputsShape>();
  int val_idx = 0;
  const Key &key = keys[0];

  if (optim_inputs_shape_.count(key) == 0) {
    optim_inputs_shape_[key] = inputs_shape;
  }
  for (size_t i = 0; i < keys.size(); i++) {
    auto shape = std::make_shared<std::vector<size_t>>();
    inputs_shape->push_back(shape);

    int len = lengths[i];
    for (int j = 0; j < len; j++) {
      shape->push_back(values[val_idx++]);
    }
  }
  if (weight_key_to_optims_.count(key) > 0) {
    const std::string &optim_name = weight_key_to_optims_[key];
    if (optimizers_.count(optim_name) == 0 && optim_inputs_shape_.count(key) > 0) {
      if (optim_name == kSparseAdam) {
        std::shared_ptr<PServerKernel> optimizer =
          std::make_shared<kernel::ps::SparseApplyAdamPSKernel>(rank_id_, pserver_num_);
        optimizer->InitKernel(optim_inputs_shape_[key]);
        optimizers_[optim_name] = optimizer;
      } else if (optim_name == kApplyMomentum) {
        std::shared_ptr<PServerKernel> optimizer =
          std::make_shared<kernel::ps::ApplyMomentumPSKernel>(rank_id_, pserver_num_);
        optimizer->InitKernel(optim_inputs_shape_[key]);
        optimizers_[optim_name] = optimizer;
      } else if (optim_name == kSparseFtrl) {
        std::shared_ptr<PServerKernel> optimizer =
          std::make_shared<kernel::ps::SparseApplyFtrlPSKernel>(rank_id_, pserver_num_);
        optimizer->InitKernel(optim_inputs_shape_[key]);
        optimizers_[optim_name] = optimizer;
      }
    }
  }
}

template <typename T>
void ParameterServer<T>::InitWeight(const Key &key, const WeightPtr &weight) {
  if (weights_.count(key) == 0) {
    weights_[key] = weight;
  }
}

template <typename T>
void ParameterServer<T>::InitGrad(const Key &key, const GradPtr &grad) {
  if (grads_.count(key) == 0) {
    grads_[key] = grad;
    grads_accum_counter_[key] = 0;
  }
}

template <typename T>
void ParameterServer<T>::InitEmbeddingTable(
  const Key &key, const std::shared_ptr<std::vector<std::shared_ptr<std::vector<size_t>>>> &shapes) {
  // Init embedding lookup kernel
  std::shared_ptr<PServerKernel> lookup = std::make_shared<kernel::ps::EmbeddingLookUpPSKernel>(rank_id_, pserver_num_);
  lookup->InitKernel(shapes);
  embedding_lookup_ops_[key] = lookup;

  // Init embedding weight
  const std::vector<size_t> &input_shapes = lookup->input_sizes();
  size_t total_dims = 1;
  for (auto shape : input_shapes) {
    total_dims *= shape;
  }
  WeightPtr embedding = std::make_shared<Weight>(total_dims, 0.01);
  weights_[key] = embedding;

  grads_accum_counter_[key] = 0;
}

template <typename T>
void ParameterServer<T>::UpdateWeights() {
  while (true) {
    std::unique_lock<std::mutex> lock(mutex_);
    apply_grads_cv_.wait(lock, [this] { return this->ReadyForUpdateWeights(); });

    for (auto iter = weights_.begin(); iter != weights_.end(); iter++) {
      Key key = iter->first;
      WeightPtr weight_ptr = iter->second;

      std::shared_ptr<PServerKernel> optimizer = nullptr;
      if (weight_key_to_optims_.count(key) > 0) {
        const std::string &optim_name = weight_key_to_optims_[key];
        optimizer = optimizers_[optim_name];
      }
      MS_EXCEPTION_IF_NULL(optimizer);

      std::shared_ptr<OptimizerInfo> optim_info = optim_infos_[key];
      if (optim_info == nullptr) {
        continue;
      }
      const WeightPtr &weight = weights_[key];
      optim_info->UpdateWeight(weight);
      const std::vector<kernel::AddressPtr> &inputs = optim_info->inputs();
      const std::vector<kernel::AddressPtr> &workspaces = optim_info->workspaces();
      const std::vector<kernel::AddressPtr> &outputs = optim_info->outputs();

      optimizer->Execute(inputs, workspaces, outputs);
      optim_info->Reset();
    }
    ResetGradAccumCount();
    accum_grads_cv_.notify_all();
  }
}

template <typename T>
void ParameterServer<T>::AccumGrad(const Keys &keys, const Values &values, const Lengths &lengths) {
  std::unique_lock<std::mutex> lock(mutex_);
  accum_grads_cv_.wait(lock, [this] { return this->ReadyForAccumGrads(); });

  const Key &key = keys[0];
  std::shared_ptr<OptimizerInfo> optim_info = optim_infos_[key];

  // Create or update the optimizer info
  if (optim_info == nullptr) {
    const std::shared_ptr<OptimizerInfoBuilder> &builder = optim_info_builders_[weight_key_to_optims_[key]];
    std::shared_ptr<kernel::ps::PServerKernel> pserver_kernel = optimizers_[weight_key_to_optims_[key]];
    if (pserver_kernel == nullptr) {
      MS_LOG(EXCEPTION) << "no optimizer found for key " << key << " optim name " << weight_key_to_optims_[key];
    }
    MS_EXCEPTION_IF_NULL(pserver_kernel);
    OptimizerInfo *optim =
      builder->Build(pserver_kernel, weights_[key], keys, values, lengths, optim_inputs_shape_[key], worker_num_);
    optim_info.reset(optim);
    optim_infos_[key] = optim_info;
  } else {
    optim_info->Update(values, lengths);
  }
  MS_EXCEPTION_IF_NULL(optim_info);

  optim_info->Accumulate(values, lengths);

  grads_accum_counter_[key] += 1;
  if (grads_accum_counter_[key] == worker_num_) {
    grad_accum_count_++;
  }
  if (ReadyForUpdateWeights()) {
    apply_grads_cv_.notify_one();
  }
}

template <typename T>
WeightPtr ParameterServer<T>::weight(const Key &key) {
  std::unique_lock<std::mutex> lock(mutex_);

  if (weights_.count(key) == 0) {
    MS_LOG(ERROR) << "Invalid weight key " << key;
    return nullptr;
  }
  WeightPtr weight_ptr = weights_[key];
  WeightPtr copy_weight_ptr = std::make_shared<::ps::SArray<T>>(weight_ptr->size(), 0);
  copy_weight_ptr->CopyFrom(weight_ptr->data(), weight_ptr->size());
  return copy_weight_ptr;
}

template <typename T>
void ParameterServer<T>::DoEmbeddingLookup(Key key, const LookupIds &lookup_ids, ::ps::KVPairs<T> *res) {
  std::unique_lock<std::mutex> lock(mutex_);
  if (weights_.count(key) == 0) {
    MS_LOG(ERROR) << "Invalid embedding table key " << key;
    return;
  }
  if (embedding_lookup_ops_.count(key) == 0) {
    MS_LOG(ERROR) << "Invalid embedding lookup op key " << key;
    return;
  }
  WeightPtr table_ptr = weights_[key];
  std::shared_ptr<PServerKernel> table_lookup_op = embedding_lookup_ops_[key];

  // Update shapes of lookup operator
  std::shared_ptr<std::vector<std::shared_ptr<std::vector<size_t>>>> shapes =
    std::make_shared<std::vector<std::shared_ptr<std::vector<size_t>>>>();
  std::shared_ptr<std::vector<size_t>> indices_shape = std::make_shared<std::vector<size_t>>();
  indices_shape->emplace_back(lookup_ids.size());
  shapes->push_back(indices_shape);
  table_lookup_op->ReInit(shapes);

  const std::vector<size_t> output_shapes = table_lookup_op->output_sizes();
  std::vector<kernel::AddressPtr> inputs;
  AddressPtr embedding_table = std::make_shared<kernel::Address>();
  AddressPtr indices = std::make_shared<kernel::Address>();
  inputs.push_back(embedding_table);
  inputs.push_back(indices);
  embedding_table->addr = table_ptr->data();
  embedding_table->size = table_ptr->size() * sizeof(T);
  indices->addr = lookup_ids.data();
  indices->size = lookup_ids.size() * sizeof(T);

  std::vector<kernel::AddressPtr> workspaces;
  std::vector<kernel::AddressPtr> outputs;
  AddressPtr output = std::make_shared<kernel::Address>();
  std::shared_ptr<Values> addr = std::make_shared<Values>(output_shapes[0] / sizeof(T), 0);

  output->addr = addr->data();
  output->size = output_shapes[0];
  outputs.push_back(output);

  table_lookup_op->Execute(inputs, workspaces, outputs);
  res->vals = *addr;
  res->lens.push_back(res.vals.size());
}

template <typename T>
int ParameterServer<T>::SumOfShapes(const std::vector<int> &shapes) const {
  int sum = 1;
  for (auto shape : shapes) {
    sum *= shape;
  }
  return sum;
}

template <typename T>
size_t ParameterServer<T>::PreComputeCapacity(const Keys &keys, const Lengths &lens) {
  size_t capacity = 0;
  for (size_t i = 0; i < keys.size(); i++) {
    Key key = keys[i];
    if (embedding_row_lens_.count(key) > 0) {
      capacity += embedding_row_lens_[key] * lens[i];
    } else {
      MS_LOG(ERROR) << "Invalid embedding lookup id " << key;
    }
  }
  return capacity;
}

template <typename T>
inline bool ParameterServer<T>::ReadyForUpdateWeights() {
  return grads_accum_counter_.size() > 0 && grad_accum_count_ == grads_accum_counter_.size();
}

template <typename T>
inline bool ParameterServer<T>::ReadyForAccumGrads() {
  return grad_accum_count_ < weights_.size();
}

template <typename T>
inline void ParameterServer<T>::ResetGradAccumCount() {
  grad_accum_count_ = 0;
  for (auto iter = grads_accum_counter_.begin(); iter != grads_accum_counter_.end(); iter++) {
    grads_accum_counter_[iter->first] = 0;
  }
}

template <typename T>
void ParameterServer<T>::Run(const FuncGraphPtr &func_graph) {
  ::ps::Start(0);
  if (!::ps::IsServer()) {
    std::cout << "This is not ther Server" << std::endl;
    return;
  }
  Init(func_graph);
  thread_->join();
}
}  // namespace ps
}  // namespace parallel
}  // namespace mindspore
#endif  // MINDSPORE_MINDSPORE_CCSRC_PARALLEL_PS_PARAMETER_SERVER_H_