#ifndef CAFFE2_OPERATORS_H_SOFTMAX_OP_H_
#define CAFFE2_OPERATORS_H_SOFTMAX_OP_H_

#include <c10/util/Optional.h>
#include "caffe2/core/context.h"
#include "caffe2/core/logging.h"
#include "caffe2/core/operator.h"
#include "caffe2/proto/hsm.pb.h"
#include "caffe2/utils/math.h"

namespace caffe2 {

template <typename T, typename Context>
class HSoftmaxOpBase : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit HSoftmaxOpBase(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...) {
    HierarchyProto hierarchy;
    CAFFE_ENFORCE(hierarchy.ParseFromString(
        this->template GetSingleArgument<string>("hierarchy", "")));
    for (const auto& path : hierarchy.paths()) {
      hierarchy_all_map_.emplace(path.word_id(), path);
    }
  }

 protected:
  std::unordered_map<int, PathProto> hierarchy_all_map_;
  c10::optional<Tensor> scale_;
  c10::optional<Tensor> sum_multiplier_;
  c10::optional<Tensor> bias_multiplier_;
  static constexpr T kLOG_THRESHOLD() {
    return 1e-20f;
  }
  static std::unordered_map<int, PathProto> getHierarchyForLabels(
      int M,
      const int* labels,
      const std::unordered_map<int, PathProto>& hierarchy_all_map) {
    std::unordered_map<int, PathProto> hierarchy_map;
    std::set<int> label_set = std::set<int>(labels, labels + M);
    for (const auto& label : label_set) {
      auto search = hierarchy_all_map.find(label);
      CAFFE_ENFORCE(search != hierarchy_all_map.end(), "incorrect label.");
      hierarchy_map.emplace(search->first, search->second);
    }
    return hierarchy_map;
  }
  int getIntermediateOutputSize(
      const int* labels,
      int M,
      std::unordered_map<int, PathProto>& hierarchy) const {
    int size = 0;
    for (int label = 0; label < M; ++label) {
      int word_id = labels[label];
      const auto& path = hierarchy[word_id];
      size += std::accumulate(
          path.path_nodes().begin(),
          path.path_nodes().end(),
          0,
          // Output of FC + Output of Softmax
          [](int sz, PathNodeProto node) {
            return sz + 2 * node.length();
          });
    }
    return size;
  }
};

template <typename T, class Context>
class HSoftmaxOp : public HSoftmaxOpBase<T, Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  using HSoftmaxOpBase<T, Context>::HSoftmaxOpBase;

  bool RunOnDevice() override;

 protected:
  float RunForwardSingle(
      const float* X,
      const float* W,
      const float* b,
      int target,
      float* output,
      const float* bias_multiplier,
      int w_length,
      int K,
      int& output_offset);
};

template <typename T, class Context>
class HSoftmaxGradientOp final : public HSoftmaxOpBase<T, Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  using HSoftmaxOpBase<T, Context>::HSoftmaxOpBase;
  bool RunOnDevice() override;

 private:
  void RunBackwardSingle(
      const float* X,
      const float* dY,
      const float* W,
      int target,
      const float* int_output,
      float* dX,
      float* dW,
      float* db,
      float* dOutput,
      int dim_in,
      int w_length,
      int& output_offset);
};

template <typename T, class Context>
class HSoftmaxSearchOp final : public HSoftmaxOp<T, Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit HSoftmaxSearchOp(Args&&... args)
      : HSoftmaxOp<T, Context>(std::forward<Args>(args)...),
        top_n_(this->template GetSingleArgument<int>("topN", 5)),
        beam_(this->template GetSingleArgument<float>("beam", 0.01f)) {
    CAFFE_ENFORCE(tree_.ParseFromString(
        this->template GetSingleArgument<string>("tree", "")));
  }
  bool RunOnDevice() override;

 private:
  int top_n_;
  float beam_;
  TreeProto tree_;
  bool pruning(
      const float* X,
      int sample,
      int K,
      const float* W,
      const float* b,
      const NodeProto& src_node,
      NodeProto& dst_node,
      float parent_score,
      float beam);
  bool extractNodes(
      const NodeProto& node,
      std::vector<std::pair<string, float>>& info);
};

template <typename T, class Context>
class HuffmanTreeHierarchyOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  template <class... Args>
  explicit HuffmanTreeHierarchyOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...),
        num_classes_(this->template GetSingleArgument<int>("num_classes", -1)) {
  }
  bool RunOnDevice() override;

 private:
  // Internal huffman tree data.
  struct Node {
    Node(T l, int count)
        : label(l), count(count), left_ch_index(-1), right_ch_index(-1) {}
    T label;
    int count;
    int left_ch_index;
    int right_ch_index;
  };

  struct NodeComparator {
    bool operator()(const Node& node_a, const Node& node_b) {
      return node_a.count > node_b.count;
    }
  };

  int num_classes_;
};

} // namespace caffe2

#endif // CAFFE2_OPERATORS_SOFTMAX_OP_H_