/*
Copyright 2020 The OneFlow 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.
*/
#ifndef ONEFLOW_CORE_FRAMEWORK_ATTR_VALUE_H_
#define ONEFLOW_CORE_FRAMEWORK_ATTR_VALUE_H_

#include "oneflow/core/framework/user_op_attr.pb.h"
#include "oneflow/core/common/util.h"
#include "oneflow/core/common/shape.h"
#include "oneflow/core/common/stride.h"
#include "oneflow/core/common/data_type.h"
#include "oneflow/core/common/protobuf.h"

namespace oneflow {

template<typename T>
size_t HashTypedAttrVal(const T& val);

namespace user_op {

// SEQ
#define BASIC_ATTR_SEQ                                         \
  OF_PP_MAKE_TUPLE_SEQ(at_int32, int32_t, AttrType::kAtInt32)  \
  OF_PP_MAKE_TUPLE_SEQ(at_int64, int64_t, AttrType::kAtInt64)  \
  OF_PP_MAKE_TUPLE_SEQ(at_bool, bool, AttrType::kAtBool)       \
  OF_PP_MAKE_TUPLE_SEQ(at_float, float, AttrType::kAtFloat)    \
  OF_PP_MAKE_TUPLE_SEQ(at_double, double, AttrType::kAtDouble) \
  OF_PP_MAKE_TUPLE_SEQ(at_string, std::string, AttrType::kAtString)

#define ENUM_ATTR_SEQ OF_PP_MAKE_TUPLE_SEQ(at_data_type, DataType, AttrType::kAtDataType)

#define MESSAGE_ATTR_SEQ                                    \
  OF_PP_MAKE_TUPLE_SEQ(at_shape, Shape, AttrType::kAtShape) \
  OF_PP_MAKE_TUPLE_SEQ(at_stride, Stride, AttrType::kAtStride)

#define LIST_BASIC_ATTR_SEQ                                                         \
  OF_PP_MAKE_TUPLE_SEQ(at_list_int32, std::vector<int32_t>, AttrType::kAtListInt32) \
  OF_PP_MAKE_TUPLE_SEQ(at_list_int64, std::vector<int64_t>, AttrType::kAtListInt64) \
  OF_PP_MAKE_TUPLE_SEQ(at_list_float, std::vector<float>, AttrType::kAtListFloat)

#define LIST_ENUM_ATTR_SEQ \
  OF_PP_MAKE_TUPLE_SEQ(at_list_data_type, std::vector<DataType>, AttrType::kAtListDataType)

#define LIST_MESSAGE_ATTR_SEQ                                                     \
  OF_PP_MAKE_TUPLE_SEQ(at_list_shape, std::vector<Shape>, AttrType::kAtListShape) \
  OF_PP_MAKE_TUPLE_SEQ(at_list_stride, std::vector<Stride>, AttrType::kAtListStride)

#define LIST_STRING_ATTR_SEQ \
  OF_PP_MAKE_TUPLE_SEQ(at_list_string, std::vector<std::string>, AttrType::kAtListString)

#define ATTR_SEQ        \
  BASIC_ATTR_SEQ        \
  ENUM_ATTR_SEQ         \
  MESSAGE_ATTR_SEQ      \
  LIST_BASIC_ATTR_SEQ   \
  LIST_ENUM_ATTR_SEQ    \
  LIST_MESSAGE_ATTR_SEQ \
  LIST_STRING_ATTR_SEQ

// Type Trait: GetAttrType, GetCppType

template<typename T>
struct GetAttrType;

template<AttrType AttrT>
struct GetCppType;

#define SPECIALIZE_GET_ATTR_TYPE(field, type_cpp, type_proto)                     \
  template<>                                                                      \
  struct GetAttrType<type_cpp> : std::integral_constant<AttrType, type_proto> {}; \
  template<>                                                                      \
  struct GetCppType<type_proto> {                                                 \
    typedef type_cpp type;                                                        \
  };
OF_PP_FOR_EACH_TUPLE(SPECIALIZE_GET_ATTR_TYPE, ATTR_SEQ);
#undef SPECIALIZE_GET_ATTR_TYPE

class AttrVal {
 public:
  AttrVal() = default;
  virtual ~AttrVal() = default;

  virtual AttrType type() const = 0;
  virtual size_t hash_value() const = 0;

  virtual const void* Ptr() const = 0;
  virtual bool operator==(const AttrVal& other) const = 0;
  bool operator!=(const AttrVal& other) const { return !(*this == other); }

 private:
  OF_DISALLOW_COPY_AND_MOVE(AttrVal);
};

template<typename T>
class TypedAttrValIf : public AttrVal {
 public:
  virtual const T& val() const = 0;

  AttrType type() const override { return GetAttrType<T>::value; }

  bool operator==(const AttrVal& other) const override {
    if (other.type() != GetAttrType<T>::value) { return false; }
    return *static_cast<const T*>(Ptr()) == *static_cast<const T*>(other.Ptr());
  }
};

template<typename T>
class TypedAttrVal final : public TypedAttrValIf<T> {
 public:
  TypedAttrVal(T v) : val_(v) {}
  ~TypedAttrVal() = default;

  const T& val() const override { return val_; }
  const void* Ptr() const override { return static_cast<const void*>(&val_); }

  size_t hash_value() const override { return std::hash<T>()(val_); }

 private:
  OF_DISALLOW_COPY_AND_MOVE(TypedAttrVal);

  T val_;
};

template<typename T>
class TypedAttrValRef final : public TypedAttrValIf<T> {
 public:
  TypedAttrValRef(const T* v) : val_(v) {}
  ~TypedAttrValRef() = default;

  const T& val() const override { return *val_; }
  const void* Ptr() const override { return static_cast<const void*>(val_); }

  size_t hash_value() const override { return std::hash<T>()(*val_); }

 private:
  OF_DISALLOW_COPY_AND_MOVE(TypedAttrValRef);

  const T* val_;
};

}  // namespace user_op

template<typename T>
const T& AttrValueCast(const user_op::AttrVal& val);

template<typename T>
std::shared_ptr<user_op::AttrVal> CastAttrValue(const T& attr_val);

template<typename T>
std::shared_ptr<user_op::AttrVal> CastAttrValue(const T* attr_val);

}  // namespace oneflow

#endif  // ONEFLOW_CORE_FRAMEWORK_ATTR_VALUE_H_
