#pragma once

#include <ATen/core/TensorBody.h>
#include <ATen/core/functional.h>
#include <ATen/core/interned_strings.h>
#include <ATen/core/ivalue.h>
#include <ATen/core/qualified_name.h>
#include <c10/util/TypeList.h>

#include <c10/util/Optional.h>

#include <iostream>
#include <memory>
#include <type_traits>
#include <array>

struct ClassType;
namespace torch {
namespace jit {
struct Function;
namespace script {
struct CompilationUnit;
}
} // namespace jit
} // namespace torch

namespace c10 {

struct FunctionSchema;
using OptNameList = c10::optional<std::vector<std::string>>;

#define C10_FORALL_TYPES(_) \
  _(AnyType)                \
  _(TensorType)             \
  _(TupleType)              \
  _(ListType)               \
  _(DictType)               \
  _(NumberType)             \
  _(FloatType)              \
  _(FutureType)             \
  _(IntType)                \
  _(NoneType)               \
  _(StringType)             \
  _(GeneratorType)          \
  _(BoolType)               \
  _(OptionalType)           \
  _(VarType)                \
  _(DeviceObjType)          \
  _(FunctionType)           \
  _(ClassType)              \
  _(CapsuleType)            \
  _(InterfaceType)

enum class TypeKind {
#define DEFINE_TYPE(T) T,
  C10_FORALL_TYPES(DEFINE_TYPE)
#undef DEFINE_TYPE
};

CAFFE2_API const char* typeKindToString(TypeKind kind);

struct Type;
using TypePtr = std::shared_ptr<Type>;

struct CAFFE2_API Type : std::enable_shared_from_this<Type> {
 private:
  TypeKind kind_;

 protected:
  Type(TypeKind kind) : kind_(kind) {}

 public:
  virtual bool operator==(const Type& rhs) const = 0;

  // subtyping relation. By default, we return true for the case
  // when the type is exactly equal or if this <: T where rhs = Optional[T]

  // if this returns false and the why_not stream is non-null, it contains
  // additional details that describe why this is not a subtype of 'rhs'.
  // This additional information should only contain details that are not obvious
  // from the python_str() that describes the type. For instance it is clear that `int <: str` is false
  // but not clear why `Foo <: InterfaceBar` might be false.
  virtual bool isSubtypeOfExt(const TypePtr rhs, std::ostream* why_not) const;
  virtual bool is_module() const;
  bool isSubtypeOf(const TypePtr rhs) const {
    return isSubtypeOfExt(rhs, nullptr);
  }

  // How this type will appear in FunctionSchema declarations
  virtual std::string str() const = 0;

  // How this type will appear as if it were a type annotation in Python
  // which is sometimes different than how it appears in declarations (e.g.
  // int[] vs List[int])
  virtual std::string python_str() const {
    return str();
  }

  TypeKind kind() const {
    return kind_;
  }

  virtual bool requires_grad() const {
    for (const auto& ct : containedTypes()) {
      if (ct->requires_grad()) {
        return true;
      }
    }
    return false;
  }

  // Dynamically cast this object to the subclass indicated by the
  // template variable, returning nullptr if the cast is invalid.
  template <typename T>
  std::shared_ptr<T> cast() {
    if (T::Kind == kind()) {
      return std::static_pointer_cast<T>(shared_from_this());
    }
    return nullptr;
  }
  template <typename T>
  std::shared_ptr<const T> cast() const {
    if (T::Kind == kind()) {
      return std::static_pointer_cast<const T>(shared_from_this());
    }
    return nullptr;
  }
  template <typename T>
  std::shared_ptr<T> expect() {
    auto r = cast<T>();
    AT_ASSERT(r);
    return r;
  }
  template <typename T>
  std::shared_ptr<const T> expect() const {
    auto r = cast<const T>();
    AT_ASSERT(r);
    return r;
  }
  virtual ~Type() = default;
  virtual bool hasFreeVariables() const {
    return false;
  }
  // list of types this type contains, e.g. for a List then element type of a
  // list for a tuple, the types of the tuple elements
  virtual at::ArrayRef<TypePtr> containedTypes() const {
    return {};
  }
  // create a new version of this type, replacing its contained types with
  // contained_types
  TypePtr withContained(std::vector<TypePtr> contained_types) {
    auto current_contained = containedTypes();
    AT_ASSERT(current_contained.size() == contained_types.size());
    if (current_contained.equals(contained_types)) {
      return shared_from_this();
    }
    return createWithContained(std::move(contained_types));
  }
  // per-type constructor, you only need to override this if the
  // containedTypes() is not empty
  virtual TypePtr createWithContained(
      std::vector<TypePtr> contained_types) const {
    AT_ERROR(
        "type with contained types did not overload createWithContained: ",
        str());
  }
};

struct AnyType;
using AnyTypePtr = std::shared_ptr<AnyType>;
// Any is the top of the type hierarchy, all other types are subtypes
// T <: Any, forall T
struct CAFFE2_API AnyType : public Type {
  static AnyTypePtr create() {
    return AnyTypePtr(
        new AnyType()); // NOLINT(modernize-make-shared)
  }
  bool operator==(const Type& rhs) const override {
    return rhs.kind() == kind();
  }
  std::string str() const override {
    return "Any";
  }
  static const TypeKind Kind = TypeKind::AnyType;
  // global singleton
  static AnyTypePtr get();

 private:
  AnyType() : Type(TypeKind::AnyType) {}
};

inline std::string toString(TypePtr typePtr) {
  return typePtr->str();
}

inline bool operator!=(const Type& lhs, const Type& rhs) {
  return !(lhs == rhs);
}

// common base for all types that have a single sub element
// e.g. Future[T], Option[T], List[T]
template <TypeKind K, typename T>
struct SingleElementType : public Type {
  static const TypeKind Kind = K;

  TypePtr getElementType() const {
    return elem;
  }

  bool hasFreeVariables() const override {
    return getElementType()->hasFreeVariables();
  }

  at::ArrayRef<TypePtr> containedTypes() const override {
    return elem;
  }

  bool operator==(const Type& rhs) const override {
    if (auto rhs_ = rhs.cast<T>()) {
      return *getElementType() == *rhs_->getElementType();
    }
    return false;
  }

 protected:
  SingleElementType(TypePtr elem) : Type(Kind), elem(std::move(elem)) {}

 private:
  TypePtr elem;
};

struct OptionalType;
using OptionalTypePtr = std::shared_ptr<OptionalType>;
// This type represents an optional type, for each element type.
// Optional[T] can accept both T and None(nullopt in C++)
// Subtype hierarchy for Optional:
// 1. Optional[T] <: Optional[R] iff T <: R
// 2. T <: Optional[R] if T <: R
// 3. None <: Optional[T] for all T
struct CAFFE2_API OptionalType
    : public SingleElementType<TypeKind::OptionalType, OptionalType> {
  static OptionalTypePtr create(TypePtr element) {
    // Optional is a union of [None, T], so Optional[[Optional[T]]] ->
    // Optional[T]
    if (auto opt_ptr = element->cast<OptionalType>()) {
      return opt_ptr;
    }
    return OptionalTypePtr(
        new OptionalType(std::move(element))); // NOLINT(modernize-make-shared)
  }

  std::string str() const override {
    std::stringstream ss;
    ss << getElementType()->str() << "?";
    return ss.str();
  }
  std::string python_str() const override {
    std::stringstream ss;
    ss << "Optional[" << getElementType()->python_str() << "]";
    return ss.str();
  }

  TypePtr createWithContained(
      std::vector<TypePtr> contained_types) const override {
    AT_ASSERT(contained_types.size() == 1);
    return create(contained_types[0]);
  }

  bool isSubtypeOfExt(const TypePtr rhs, std::ostream* why_not) const override {
    if (Type::isSubtypeOfExt(rhs, why_not)) {
      return true;
    }
    if (auto rhs_ = rhs->cast<OptionalType>()) {
      return getElementType()->isSubtypeOfExt(rhs_->getElementType(), why_not);
    }
    return false;
  }
  // common cast Optional[Tensor] for undefined tensor type
  static OptionalTypePtr ofTensor();

 private:
  OptionalType(TypePtr elem) : SingleElementType(elem) {}
};

template <typename T>
inline c10::optional<T> merge_primitive(
    const c10::optional<T>& a,
    const c10::optional<T>& b) {
  if (a.has_value() && b.has_value() && a.value() == b.value()) {
    return a;
  }
  return c10::optional<T>{};
}

// `VaryingShape` tracks if individual dimensions or a rank vary across
// profiled runs. A *varying* or *dynamic* dimension is expressed as
// an empty c10::optional in `sizes_`. If a rank is dynamic, the entire
// `sizes_` becomes the empty optional.
struct CAFFE2_API VaryingShape {
  using ListOfOptionalInts = std::vector<c10::optional<int64_t>>;
  VaryingShape(const std::vector<int64_t>& vec)
      : VaryingShape(ListOfOptionalInts(vec.begin(), vec.end())) {}

  VaryingShape(c10::ArrayRef<int64_t> vec)
      : VaryingShape(ListOfOptionalInts(vec.begin(), vec.end())){}

  VaryingShape(c10::optional<size_t> size = c10::nullopt) : dims_(c10::nullopt) {
    if (size) {
      dims_ = ListOfOptionalInts(*size);
    }
  }

  VaryingShape(ListOfOptionalInts dims)
  : dims_(std::move(dims)) {}

  VaryingShape(size_t size) : VaryingShape(c10::optional<size_t>(size)) {}

  bool operator==(const VaryingShape& other) const {
    return dims_ == other.dims_;
  }

  const c10::optional<int64_t>& operator[](int i) const {
    if (!dims_) {
      throw std::runtime_error("Rank isn't fixed");
    }
    return (*dims_).at(i);
  }

  c10::optional<size_t> size() const {
    if (!dims_) {
      return c10::nullopt;
    }
    const auto& dims = dims_.value();
    return dims.size();
  }

  const c10::optional<ListOfOptionalInts>& sizes() const {
    return dims_;
  }

  VaryingShape merge(const VaryingShape& other) const;

  c10::optional<std::vector<int64_t>> concrete_sizes() const {
    if (!dims_) {
      return c10::nullopt;
    }
    std::vector<int64_t> sizes;
    for (auto d : *dims_) {
      if (!d) {
        return c10::nullopt;
      }
      sizes.push_back(d.value());
    }
    return sizes;
  }

  bool isComplete() const {
    if (!dims_) {
      return false;
    }
    for (auto d : *dims_) {
      if(!d) {
        return false;
      }
    }
    return true;
  }

 private:
  c10::optional<ListOfOptionalInts> dims_;
};

using VaryingStrides = VaryingShape;

struct TensorType;
using TensorTypePtr = std::shared_ptr<TensorType>;
// This type represents a single Tensor with a specific size
struct CAFFE2_API TensorType : public Type {
  static TensorTypePtr create(const at::Tensor& t) {
    return TensorTypePtr(new TensorType(t));
  }

  static TensorTypePtr create(c10::optional<at::ScalarType> scalar_type,
                              c10::optional<Device> device,
                              const VaryingShape &sizes,
                              const VaryingStrides &strides,
                              c10::optional<bool> requires_grad,
                              c10::optional<bool> undefined = false) {
    return TensorTypePtr(new TensorType(scalar_type, device, sizes, strides,
                                        requires_grad, undefined));
  }

  static TensorTypePtr create(
      c10::optional<at::ScalarType> scalar_type,
      c10::optional<Device> device,
      c10::optional<size_t> dim,
      c10::optional<bool> requires_grad) {
    return TensorType::create(
        scalar_type,
        device,
        VaryingShape(dim),
        VaryingShape(dim),
        requires_grad);
  }

  // overloaded create variadic template argument as it could not distinguish
  // initializer list
  static TensorTypePtr createContiguous(
      at::ScalarType scalar_type,
      at::Device device,
      at::IntArrayRef sizes) {
    return create(
        scalar_type,
        device,
        VaryingShape(sizes),
        VaryingShape(contiguousStridesOf(sizes)),
        c10::nullopt);
  }
  static TensorTypePtr create(
      at::ScalarType scalar_type,
      at::Device device,
      at::IntArrayRef sizes,
      at::IntArrayRef strides) {
    return create(
        scalar_type,
        device,
        VaryingShape(sizes),
        c10::VaryingShape(strides),
        c10::nullopt);
  }
  static TypePtr fromNumberType(TypePtr typ);
  static TypePtr fromBoolType();

  c10::optional<size_t> dim() const {
    return sizes().size();
  }

  const VaryingShape& sizes() const {
    return sizes_;
  }
  const VaryingStrides& strides() const {
    return strides_;
  }
  c10::optional<at::Device> device() const {
    return device_;
  }
  c10::optional<at::ScalarType> scalarType() const {
    return scalar_type_;
  }
  c10::optional<bool> requiresGrad() const {
    return requires_grad_;
  }
  bool requires_grad() const override {
    return requires_grad_ ? *requires_grad_ : true;
  }


  bool operator==(const Type& rhs) const override {
    if (rhs.kind() != kind()) {
      return false;
    }

    auto rt = rhs.expect<TensorType>();
    return scalar_type_ == rt->scalarType() && sizes() == rt->sizes() &&
           strides() == rt->strides() && device() == rt->device() &&
           requiresGrad() == rt->requiresGrad() &&
           undefined() == rt->undefined();
  }
  bool isSubtypeOfExt(const TypePtr rhs, std::ostream* why_not) const override;

  std::string str() const override;

  c10::optional<size_t> numel() const {
    size_t prod = 1;
    const auto& shape = sizes();

    for (size_t i = 0; i < shape.size(); i++) {
      if (!shape[i]) {
        return c10::optional<size_t>{};
      }
      prod *= shape[i].value();
    }
    return prod;
  }

  TensorTypePtr withRequiresGrad(c10::optional<bool> s) {
    auto copy = clone();
    copy->requires_grad_ = s;
    return copy;
  }

  TensorTypePtr withScalarType(c10::optional<ScalarType> st) {
    auto copy = clone();
    copy->scalar_type_ = st;
    return copy;
  }


  TensorTypePtr withDim(c10::optional<size_t> d) {
    auto copy = clone();
    copy->sizes_ = VaryingShape(d);
    copy->strides_ = VaryingShape(d);
    return copy;
  }

  TensorTypePtr withSizesStrides(
      at::IntArrayRef sizes,
      at::IntArrayRef strides) const {
    auto cloned = clone();
    cloned->sizes_ = VaryingShape(sizes);
    cloned->strides_ = VaryingStrides(strides);
    return cloned;
  }

  TensorTypePtr withSizes(at::IntArrayRef sizes) const {
    return withSizesStrides(
        sizes, contiguousStridesOf(sizes));
  }

  TensorTypePtr dimensionedOnly() const {
    auto copy = clone();
    copy->sizes_ = VaryingShape(sizes().size());
    copy->strides_ = VaryingShape(sizes().size());
    return copy;
  }

  TensorTypePtr contiguous() const {
    auto cloned = clone();
    if (auto concrete_sizes = sizes().concrete_sizes()) {
      cloned->strides_ = VaryingShape(contiguousStridesOf(*concrete_sizes));
    } else  {
      cloned->strides_ = VaryingShape(sizes().size());
    }
    return cloned;
  }

  TensorTypePtr merge(TensorTypePtr other) const;

  // is all information about the type specified except for autograd?
  // This replaces the notion of a 'CompleteTensorType' that used to exist
  // in the type-hierarchy. Excluding require_grad and undefined allows
  // this to match the old behavior.
  bool isComplete() const {
    return scalar_type_ && device_ && sizes_.isComplete() && strides_.isComplete();
  }

  // this property is used by GuardElimination
  // please see `checkInputs` for more details
  bool isSummarized() const {
    return !(isComplete() && requiresGrad().has_value() &&
             undefined().has_value());
  }

  TensorTypePtr withUndefined() {
    auto r = clone();
    r->undefined_ = true;
    return r;
  }

  TensorTypePtr withPossiblyUndefined() {
    auto r = clone();
    r->undefined_ = c10::nullopt;
    return r;
  }

  c10::optional<bool> undefined() const { return undefined_; }

  static TensorTypePtr get();

  static const TypeKind Kind = TypeKind::TensorType;

 private:
  TensorType(const at::Tensor& tensor)
      : Type(TypeKind::TensorType),
        scalar_type_(tensor.scalar_type()),
        device_(tensor.device()),
        sizes_(tensor.sizes().size()),
        strides_(tensor.sizes().size()),
        requires_grad_(tensor.requires_grad()),
        undefined_(!tensor.defined()) {
    if (!tensor.is_mkldnn() && !tensor.is_sparse()) {
      sizes_ = tensor.sizes().vec();
      strides_ = tensor.strides().vec();
    }
  }
  TensorType(
      c10::optional<at::ScalarType> scalar_type,
      c10::optional<Device> device,
      const VaryingShape& sizes,
      const VaryingStrides& strides,
      c10::optional<bool> requires_grad,
      c10::optional<bool> undefined = false)
      : Type(TypeKind::TensorType),
        scalar_type_(scalar_type),
        device_(device),
        sizes_(sizes),
        strides_(strides),
        requires_grad_(requires_grad),
        undefined_(undefined) {}

  TensorTypePtr clone() const {
    return TensorTypePtr(new TensorType(
        scalar_type_, device_, sizes_, strides_, requires_grad_, undefined_));
  }

  static std::vector<int64_t> contiguousStridesOf(at::IntArrayRef sizes) {
    std::vector<int64_t> strides(sizes.size());
    if (sizes.empty()) // zero-dim case
      return strides;
    strides.back() = 1;
    for (size_t i = strides.size() - 1; i > 0; i--) {
      strides[i - 1] = strides[i] * sizes[i];
    }
    return strides;
  }

  c10::optional<at::ScalarType> scalar_type_;
  c10::optional<at::Device> device_;
  VaryingShape sizes_;
  VaryingStrides strides_;
  c10::optional<bool> requires_grad_;
  // we exploit the fact certain tensors must be zero in the autograd to
  // optimize gradient computation. Such zero tensors are currently implemented
  // with `UndefinedTensorImpl.` They can be handled only by special operators
  // (e.g. `AutogradAdd`) and their `Tensor::defined()` property returns false.
  // Normally, `undefined_` is set to false, unless a type was created
  // with `withUndefined`
  // This will also mean that `undefined` tensors will fail
  // `subtypeOf(TensorType::get())` check
  // undefined_ may become `c10::nullopt` if the tensor was observed to be both
  // defined and undefined. However, no tensor type starts out with
  // `undefined_` set to `c10::nullopt`
  c10::optional<bool> undefined_;
};

struct ListType;
using ListTypePtr = std::shared_ptr<ListType>;
struct CAFFE2_API ListType
    : public SingleElementType<TypeKind::ListType, ListType> {
  // It's not exactly a singleton, but there should be exactly one instance of
  // List[T] for every T
  friend struct Type;
  template <typename... T>
  static ListTypePtr create(T&&... all) {
    return ListTypePtr(
        new ListType(std::forward<T>(all)...)); // NOLINT(modernize-make-shared)
  }

  std::string str() const override {
    std::stringstream ss;
    ss << getElementType()->str() << "[]";
    return ss.str();
  }
  std::string python_str() const override {
    std::stringstream ss;
    ss << "List[" << getElementType()->python_str() << "]";
    return ss.str();
  }
  TypePtr createWithContained(
      std::vector<TypePtr> contained_types) const override {
    return create(contained_types.at(0));
  }
  // common cast List[Tensor]
  static ListTypePtr ofTensors();
  static ListTypePtr ofInts();
  static ListTypePtr ofFloats();
  static ListTypePtr ofBools();

 private:
  ListType(TypePtr elem) : SingleElementType(elem) {}
};

struct DictType;
using DictTypePtr = std::shared_ptr<DictType>;
struct CAFFE2_API DictType : public Type {
  friend struct Type;
  static const TypeKind Kind = TypeKind::DictType;

  static DictTypePtr create(TypePtr key, TypePtr value) {
    switch (key->kind()) {
      case TypeKind::AnyType:
      case TypeKind::IntType:
      case TypeKind::FloatType:
      case TypeKind::StringType:
      case TypeKind::TensorType:
        return DictTypePtr(new DictType(key, value));
      default:
        AT_ERROR(
            "Cannot create dict for key type '",
            key->str(),
            "', only int, float, Tensor and string keys are supported");
    }
  }

  // aligned with the format in FunctionSchema
  std::string str() const override {
    std::stringstream ss;
    ss << "Dict(" << getKeyType()->str() << ", " << getValueType()->str()
       << ")";
    return ss.str();
  }

  std::string python_str() const override {
    std::stringstream ss;
    ss << "Dict[" << getKeyType()->python_str() << ", "
       << getValueType()->python_str() << "]";
    return ss.str();
  }

  TypePtr createWithContained(
      std::vector<TypePtr> contained_types) const override {
    if (contained_types.size() != 2) {
      throw std::runtime_error("Expected 2 contained types");
    }
    return create(contained_types.at(0), contained_types.at(1));
  }

  TypePtr getKeyType() const {
    return types.at(0);
  }

  TypePtr getValueType() const {
    return types.at(1);
  }

  bool hasFreeVariables() const override {
    return has_free_variables;
  }

  at::ArrayRef<TypePtr> containedTypes() const override {
    return types;
  }

  bool operator==(const Type& rhs) const override {
    if (auto dict_rhs = rhs.cast<DictType>()) {
      return *getKeyType() == *(dict_rhs->getKeyType()) &&
          *getValueType() == *(dict_rhs->getValueType());
    }
    return false;
  }

 private:
  DictType(TypePtr key, TypePtr value)
      : Type(TypeKind::DictType),
        types({key, value}),
        has_free_variables(
            key->hasFreeVariables() || value->hasFreeVariables()) {}
  std::vector<TypePtr> types;
  bool has_free_variables;
};

struct FutureType;
using FutureTypePtr = std::shared_ptr<FutureType>;

struct CAFFE2_API FutureType
    : public SingleElementType<TypeKind::FutureType, FutureType> {
  friend struct Type;
  template <typename... T>
  static FutureTypePtr create(TypePtr elem) {
    return FutureTypePtr(
        new FutureType(std::move(elem))); // NOLINT(modernize-make-shared)
  }

  std::string str() const override {
    std::stringstream ss;
    ss << "Future(" << getElementType()->str() << ")";
    return ss.str();
  }
  std::string python_str() const override {
    std::stringstream ss;
    ss << "Future[" << getElementType()->python_str() << "]";
    return ss.str();
  }
  TypePtr createWithContained(
      std::vector<TypePtr> contained_types) const override {
    return create(contained_types.at(0));
  }

 private:
  FutureType(TypePtr elem) : SingleElementType(elem) {}
};

using ::torch::jit::Function;
struct NamedType;
using NamedTypePtr = std::shared_ptr<NamedType>;

struct CAFFE2_API NamedType : public Type {
  NamedType(TypeKind tk, c10::optional<QualifiedName> name)
      : Type(tk), name_(std::move(name)) {}

  // Fully qualified name of type
  // Looks like: "foo.bar.Baz".
  const c10::optional<QualifiedName>& name() const {
    return name_;
  }
private:
  c10::optional<QualifiedName> name_;
};

// Any should never appear in a named type like a class, namedtuple or
// interface. If it does, then dynamic type information will be lost in the
// Pickler, leading to hard-to-track-down bugs that will only occur
// after saving or loading a model. This is because we rely on the
// static types in named types to reconstruct type tags of loaded
// values. Lifting this restriction requires solving the serialization
// problem first.
CAFFE2_API void checkNoAny(
    const Type& base,
    const char* what,
    const std::string& attrname,
    const TypePtr& attrtype);

struct TupleType;
using TupleTypePtr = std::shared_ptr<TupleType>;
using NameList = std::vector<std::string>;
// This type represents a Tuple
struct CAFFE2_API TupleType : public NamedType {
  static TupleTypePtr createNamed(const c10::optional<c10::QualifiedName>& name,
      const std::vector<std::string>& field_names,
      const std::vector<TypePtr>& types);
  static TupleTypePtr create(
      std::vector<TypePtr> types) {
    return TupleTypePtr(new TupleType(
        std::move(types),
        c10::nullopt,
        nullptr)); // NOLINT(modernize-make-shared)
  }

  at::ArrayRef<TypePtr> elements() const {
    return elements_;
  }

  bool operator==(const Type& rhs) const override;
  bool isSubtypeOfExt(const TypePtr rhs_, std::ostream* why_not) const override;

  std::string str() const override;
  std::string python_str() const override;
  bool hasFreeVariables() const override {
    return has_free_variables_;
  }
  at::ArrayRef<TypePtr> containedTypes() const override {
    return elements_;
  }
  TypePtr createWithContained(
      std::vector<TypePtr> contained_types) const override {
    return std::shared_ptr<TupleType>(
        new TupleType(std::move(contained_types), name(), schema()));
  }
  const std::shared_ptr<FunctionSchema>& schema() const {
    return schema_;
  }

  static const TypeKind Kind = TypeKind::TupleType;

 private:
  TupleType(
      std::vector<TypePtr> elements_,
      c10::optional<c10::QualifiedName> name,
      std::shared_ptr<FunctionSchema> schema);

  bool compare(
      const Type& rhs,
      std::function<bool(const TypePtr, const TypePtr)> fn) const {
    if (rhs.kind() != kind()) {
      return false;
    }

    const auto& l_elements = elements();
    const auto& r_elements = rhs.cast<TupleType>()->elements();
    if (l_elements.size() != r_elements.size())
      return false;
    for (size_t i = 0; i < l_elements.size(); ++i) {
      if (!fn(l_elements[i], r_elements[i]))
        return false;
    }
    return true;
  }

  std::vector<TypePtr> elements_;
  bool has_free_variables_;
  std::shared_ptr<FunctionSchema> schema_;
};

struct NumberType;
using NumberTypePtr = std::shared_ptr<NumberType>;
// This type represents a Python number
// Subtype hierarchy for Number Types (NumberType as the base type):
// IntType <: NumberType
// FloatType <: NumberType
struct CAFFE2_API NumberType : public Type {
  static NumberTypePtr create() {
    return NumberTypePtr(new NumberType()); // NOLINT(modernize-make-shared)
  }
  bool operator==(const Type& rhs) const override {
    return rhs.kind() == kind();
  }
  std::string str() const override {
    return "Scalar"; // match what PythonArgParser says for clarity
  }
  std::string python_str() const override {
    return "number"; // technically not a valid python type, but
                     // we need to use it when parsing back in annotations
                     // for implicit conversions
  }
  static const TypeKind Kind = TypeKind::NumberType;
  // global singleton
  static NumberTypePtr get();

 protected:
  NumberType(TypeKind kind = TypeKind::NumberType) : Type(kind) {}
};

struct FloatType;
using FloatTypePtr = std::shared_ptr<FloatType>;
// This type represents a Python float number
struct CAFFE2_API FloatType : public NumberType {
  static FloatTypePtr create() {
    return FloatTypePtr(new FloatType()); // NOLINT(modernize-make-shared)
  }
  bool operator==(const Type& rhs) const override {
    return rhs.kind() == kind();
  }
  std::string str() const override {
    return "float";
  }
  std::string python_str() const override {
    return "float";
  }
  bool isSubtypeOfExt(const TypePtr rhs, std::ostream* why_not) const override {
    return rhs->kind() == TypeKind::NumberType || NumberType::isSubtypeOfExt(rhs, why_not);
  }
  static const TypeKind Kind = TypeKind::FloatType;
  // global singleton
  static FloatTypePtr get();

 private:
  FloatType() : NumberType(TypeKind::FloatType) {}
};

struct IntType;
using IntTypePtr = std::shared_ptr<IntType>;
// This type represents a Python int number
struct CAFFE2_API IntType : public NumberType {
  static IntTypePtr create() {
    return IntTypePtr(new IntType()); // NOLINT(modernize-make-shared)
  }
  bool operator==(const Type& rhs) const override {
    return rhs.kind() == kind();
  }
  std::string str() const override {
    return "int";
  }
  std::string python_str() const override {
    return "int";
  }
  bool isSubtypeOfExt(const TypePtr rhs, std::ostream* why_not) const override {
    return rhs->kind() == TypeKind::NumberType || NumberType::isSubtypeOfExt(rhs, why_not);
  }
  static const TypeKind Kind = TypeKind::IntType;
  // global singleton
  static IntTypePtr get();

 private:
  IntType() : NumberType(TypeKind::IntType) {}
};

struct BoolType;
using BoolTypePtr = std::shared_ptr<BoolType>;
// This node represents a Python bool value
struct CAFFE2_API BoolType : public Type {
  static BoolTypePtr create() {
    return BoolTypePtr(new BoolType());
  }
  bool operator==(const Type& rhs) const override {
    return rhs.kind() == kind();
  }
  std::string str() const override {
    return "bool";
  }
  static const TypeKind Kind = TypeKind::BoolType;
  // global singleton
  static BoolTypePtr get();

 private:
  BoolType() : Type(TypeKind::BoolType) {}
};

struct StringType;
using StringTypePtr = std::shared_ptr<StringType>;
// This type represents a Python string
struct CAFFE2_API StringType : public Type {
  static StringTypePtr create() {
    return StringTypePtr(new StringType()); // NOLINT(modernize-make-shared)
  }
  bool operator==(const Type& rhs) const override {
    return rhs.kind() == kind();
  }
  std::string str() const override {
    // we only use "str" (not "string") in both FunctionSchema and script
    return python_str();
  }
  std::string python_str() const override {
    return "str";
  }
  static const TypeKind Kind = TypeKind::StringType;
  // global singleton
  static StringTypePtr get();

 private:
  StringType() : Type(TypeKind::StringType) {}
};

struct FunctionType;
using FunctionTypePtr = std::shared_ptr<FunctionType>;
using ::torch::jit::Function;
struct CAFFE2_API FunctionType : public NamedType {
  static FunctionTypePtr create(Function* function) {
    return FunctionTypePtr(
        new FunctionType(function)); // NOLINT(modernize-make-shared)
  }
  bool operator==(const Type& rhs) const override {
    if (auto func_type = rhs.cast<FunctionType>()) {
      return func_type->function_ == function_;
    }

    return false;
  }
  std::string str() const override {
    return "Function";
  }
  std::string python_str() const override {
    throw "Function";
  }
  Function* function() const {
    return function_;
  }
  static const TypeKind Kind = TypeKind::FunctionType;

 private:
  FunctionType(Function* function);
  Function* function_;
};

struct NoneType;
using NoneTypePtr = std::shared_ptr<NoneType>;
// This type represents a Python None
struct CAFFE2_API NoneType : public Type {
  static NoneTypePtr create() {
    return NoneTypePtr(new NoneType()); // NOLINT(modernize-make-shared)
  }
  bool operator==(const Type& rhs) const override {
    return rhs.kind() == kind();
  }
  std::string str() const override {
    return "None";
  }
  bool isSubtypeOfExt(const TypePtr rhs, std::ostream *why_not) const override {
    if (rhs->kind() == OptionalType::Kind) {
      return true;
    }
    return Type::isSubtypeOfExt(rhs, why_not);
  }
  static const TypeKind Kind = TypeKind::NoneType;
  // global singleton
  static NoneTypePtr get();

 private:
  NoneType() : Type(TypeKind::NoneType) {}
};

struct GeneratorType;
using GeneratorTypePtr = std::shared_ptr<GeneratorType>;
// This type represents a Generator
struct CAFFE2_API GeneratorType : public Type {
  static GeneratorTypePtr create() {
    return GeneratorTypePtr(
        new GeneratorType()); // NOLINT(modernize-make-shared)
  }
  bool operator==(const Type& rhs) const override {
    return rhs.kind() == kind();
  }
  std::string str() const override {
    return "Generator";
  }
  static const TypeKind Kind = TypeKind::GeneratorType;
  // global singleton
  static GeneratorTypePtr get();

 private:
  GeneratorType() : Type(TypeKind::GeneratorType) {}
};

struct DeviceObjType;
using DeviceObjTypePtr = std::shared_ptr<DeviceObjType>;
// This type represents a Generator
struct CAFFE2_API DeviceObjType : public Type {
  static DeviceObjTypePtr create() {
    return DeviceObjTypePtr(
        new DeviceObjType()); // NOLINT(modernize-make-shared)
  }
  bool operator==(const Type& rhs) const override {
    return rhs.kind() == kind();
  }
  std::string str() const override {
    return "Device";
  }
  static const TypeKind Kind = TypeKind::DeviceObjType;
  // global singleton
  static DeviceObjTypePtr get();

 private:
  DeviceObjType() : Type(TypeKind::DeviceObjType) {}
};

struct VarType;
using VarTypePtr = std::shared_ptr<VarType>;
// This type represents a type variable, used in FunctionSchema
struct VarType : public Type {
  static VarTypePtr create(std::string name_) {
    return VarTypePtr(new VarType(std::move(name_)));
  }
  bool operator==(const Type& rhs) const override {
    return rhs.kind() == kind();
  }
  std::string str() const override {
    return name();
  }
  const std::string& name() const {
    return name_;
  }
  bool hasFreeVariables() const override {
    return true;
  }
  static const TypeKind Kind = TypeKind::VarType;

 private:
  VarType(std::string name_)
      : Type(TypeKind::VarType), name_(std::move(name_)) {}
  std::string name_;
};

struct CapsuleType;
using CapsuleTypePtr = std::shared_ptr<CapsuleType>;
// This type represents a Python Capsule
struct CAFFE2_API CapsuleType : public Type {
  static CapsuleTypePtr create() {
    return CapsuleTypePtr(new CapsuleType()); // NOLINT(modernize-make-shared)
  }
  bool operator==(const Type& rhs) const override {
    return rhs.kind() == kind();
  }
  std::string str() const override {
    return "Capsule";
  }
  static const TypeKind Kind = TypeKind::CapsuleType;
  // global singleton
  static CapsuleTypePtr get();
private:
  CapsuleType()
  : Type(TypeKind::CapsuleType) {}
};

CAFFE2_API std::ostream& operator<<(std::ostream& out, const Type& t);
CAFFE2_API std::ostream& operator<<(std::ostream& out, const VaryingShape& t);
// what is the type, ignoring extra size/shape information?
// e.g. Tensor(2x3) -> Dynamic, and Tuple(Tensor(2x3),...) -> Tuple(Dynamic,...)

inline TypePtr unshapedType(const TypePtr& type) {
  if (type->isSubtypeOf(TensorType::get())) {
    return TensorType::get();
  }
  return type->withContained(fmap(type->containedTypes(), unshapedType));
}

inline TypePtr TensorType::fromNumberType(TypePtr typ) {
  if (typ->isSubtypeOf(IntType::get())) {
    return TensorType::createContiguous(at::kLong, at::kCPU, {});
  } else if (typ->isSubtypeOf(FloatType::get())) {
    return TensorType::createContiguous(at::kFloat, at::kCPU, {});
  } else if (typ->isSubtypeOf(BoolType::get())) {
    return TensorType::createContiguous(at::kLong, at::kCPU, {});
  }
  TORCH_CHECK(false, "Unknown number type: ", typ->str());
}
inline TypePtr TensorType::fromBoolType() {
  return TensorType::createContiguous(at::kLong, at::kCPU, {});
}

inline c10::optional<c10::ScalarType> tryScalarTypeFromJitType(const c10::TypePtr & type) {
  if (type == FloatType::get()) {
    return at::ScalarType::Double;
  } else if (type == IntType::get()) {
    return at::ScalarType::Long;
  } else if (type == BoolType::get()) {
    return at::ScalarType::Bool;
  }
  return c10::nullopt;
}

inline at::ScalarType scalarTypeFromJitType(const c10::TypePtr& type) {
  auto result = tryScalarTypeFromJitType(type);
  AT_ASSERTM(
      result,
      "Add new condition, expected Float, Int, or Bool but got",
      type->str());
  return *result;
}

// Attempt to find the correct supertype of t1 and t2. If none is found then
// nullopt will be returned. If t1 == t2, or t1 is a type refinement of t2,
// then t2 will be returned (and vice versa).
// Two different tensortypes will return dynamic.
// Currently we chose not to support returning a NumberType for a float & int
// input because of a lack of operator support for NumberType
CAFFE2_API c10::optional<TypePtr> unifyTypes(
    const TypePtr& t1,
    const TypePtr& t2);

CAFFE2_API c10::optional<TypePtr> unifyTypeList(at::ArrayRef<TypePtr> elements);

namespace detail {
template <typename T>
struct getTypePtr_ final {
  static TypePtr call() {
    if (!isCustomClassRegistered<T>()) {
      throw c10::Error("Type could not be converted to any of the known types.", "");
    }
    auto res = getCustomClassType<T>();
    return std::dynamic_pointer_cast<Type>(res.type_);
  }
};

template <>
struct getTypePtr_<at::Tensor> final {
  static TypePtr call() {
    return TensorType::get();
  }
};
template <>
struct getTypePtr_<double> final {
  static TypePtr call() {
    return FloatType::get();
  }
};
template <>
struct getTypePtr_<int64_t> final {
  static TypePtr call() {
    return IntType::get();
  }
};
template <>
struct getTypePtr_<bool> final {
  static TypePtr call() {
    return BoolType::get();
  }
};
template <>
struct getTypePtr_<at::Scalar> final {
  static TypePtr call() {
    return NumberType::get();
  }
};
template <>
struct getTypePtr_<at::Generator*> final {
  static TypePtr call() {
    return OptionalType::create(GeneratorType::get());
  }
};
template <>
struct getTypePtr_<std::string> final {
  static TypePtr call() {
    return StringType::get();
  }
};
template <class T>
struct getTypePtr_<std::vector<T>> final {
  static TypePtr call() {
    static auto type = ListType::create(getTypePtr_<T>::call());
    return type;
  }
};
template <class T>
struct getTypePtr_<c10::ArrayRef<T>> final {
  static TypePtr call() {
    static auto type = ListType::create(getTypePtr_<T>::call());
    return type;
  }
};
template <class T>
struct getTypePtr_<c10::List<T>> final {
  static TypePtr call() {
    static auto type = ListType::create(getTypePtr_<T>::call());
    return type;
  }
};
template <class T, size_t N>
struct getTypePtr_<std::array<T, N>> final {
  static TypePtr call() {
    static auto type = ListType::create(getTypePtr_<T>::call());
    return type;
  }
};
template <class K, class V>
struct getTypePtr_<std::unordered_map<K, V>> final {
  static TypePtr call() {
    static auto type =
        DictType::create(getTypePtr_<K>::call(), getTypePtr_<V>::call());
    return type;
  }
};
template <class K, class V>
struct getTypePtr_<c10::Dict<K, V>> final {
  static TypePtr call() {
    static auto type =
        DictType::create(getTypePtr_<K>::call(), getTypePtr_<V>::call());
    return type;
  }
};
template <class T>
struct getTypePtr_<at::optional<T>> final {
  static TypePtr call() {
    static auto type = OptionalType::create(getTypePtr_<T>::call());
    return type;
  }
};
} // namespace detail
template <class T>
inline TypePtr getTypePtr() {
  // TODO: static_assert that a templated function exists, and throw a friendy
  // error message if not
  return detail::getTypePtr_<T>::call();
}

using TypeEnv = std::unordered_map<std::string, TypePtr>;
struct MatchTypeReturn {
  MatchTypeReturn(std::string reason) : reason_(std::move(reason)) {}
  static MatchTypeReturn Success() {
    return MatchTypeReturn();
  }
  bool success() const {
    return !reason_.has_value();
  }
  const std::string& reason() const {
    return reason_.value();
  }

 private:
  MatchTypeReturn()
  : reason_(c10::nullopt) {}
  c10::optional<std::string> reason_; // is there is no match, this contains the reason
};

// attempt to match the type variables in formal to actual, adding them to type_env.
// If no match is possible this returns a MatchTypeReturn with r.success() == false
// and a r.reason() that describes why it could not match.
// note: It is possible to successfully match a formal, but for type variables
// in the formal to still not be defined. In particular, None matches Optional[T]
// but does not define the value of T.
CAFFE2_API MatchTypeReturn
matchTypeVariables(TypePtr formal, TypePtr actual, TypeEnv& type_env);

// replace type variables appearing in `type` with the values in
// `type_env`. Returns nullptr if a variable used in `type`
// does not appear in `type_env`
CAFFE2_API TypePtr tryEvalTypeVariables(TypePtr type, TypeEnv& type_env);


/**
 * User Defined Types
 */

struct ClassType;
using ClassTypePtr = std::shared_ptr<ClassType>;
using ::torch::jit::script::CompilationUnit;

// This represents a class in TorchScript.
struct CAFFE2_API ClassType : public NamedType {
  // Create a class type with name `name` and its methods stored in `cu`.
  static ClassTypePtr create(
      c10::optional<QualifiedName> qualifiedName,
      std::weak_ptr<CompilationUnit> cu,
      bool is_module = false);

  bool operator==(const Type& rhs) const override {
    if (auto user_rhs = rhs.cast<ClassType>()) {
      const auto& lhs_name = name().value();
      const auto& rhs_name = user_rhs->name().value();

      return lhs_name == rhs_name;
    }
    return false;
  }

  std::string str() const override {
    return python_str();
  }

  std::string python_str() const override {
    const auto& n = name().value();
    return n.qualifiedName();
  }

  const std::vector<Function*>& methods() const;

  TypePtr getAttribute(const std::string& name) const {
    AT_ASSERT(attributeNames_.size() == attributeTypes_.size());
    size_t pos = 0;
    for (const auto& attr : attributeNames_) {
      if (name == attr) {
        break;
      }
      ++pos;
    }

    if (pos >= attributeNames_.size()) {
      return nullptr;
    }
    return attributeTypes_[pos];
  }

  size_t numAttributes() const {
    AT_ASSERT(attributeNames_.size() == attributeTypes_.size());
    return attributeNames_.size();
  }

  const TypePtr& getAttribute(size_t slot) const {
    AT_ASSERT(attributeNames_.size() == attributeTypes_.size());
    AT_ASSERT(slot < attributeTypes_.size());
    return attributeTypes_[slot];
  }

  const std::string& getAttributeName(size_t slot) const {
    AT_ASSERT(attributeNames_.size() == attributeTypes_.size());
    AT_ASSERT(slot < attributeTypes_.size());
    return attributeNames_[slot];
  }

  // Attributes are stored in a specific slot at runtime for effiency.
  // When emitting instructions we specify the slot so that attribute access is
  // a constant lookup
  c10::optional<size_t> findAttributeSlot(const std::string& name) const {
    AT_ASSERT(attributeNames_.size() == attributeTypes_.size());
    size_t slot = 0;
    for (const auto& attr : attributeNames_) {
      if (name == attr) {
        return slot;
      }
      slot++;
    }
    return c10::nullopt;
  }
  size_t getAttributeSlot(const std::string& name) const {
    if (auto r = findAttributeSlot(name)) {
      return *r;
    }
    TORCH_CHECK(
        false,
        python_str(),
        " does not have a field with the name '",
        name,
        "'");
  }

  bool hasAttribute(const std::string& name) const {
    return std::find_if(
               attributeNames_.cbegin(),
               attributeNames_.cend(),
               [&](const std::string& attr) { return attr == name; }) !=
        attributeNames_.cend();
  }

  size_t addAttribute(
      const std::string& name,
      TypePtr type,
      bool is_parameter = false);

  // [Internal Only] Remove attribute from the ClassType,
  // caller is responsible to make sure the modification is safe:
  // it is unsafe to having existing allocations
  // of this object around anymore, and any code that works on
  // the attribute is now invalid. Only newly created code is
  // valid again.
  void unsafeRemoveAttribute(const std::string& name);

  // Add attribute \p NAME if it doesn't exist or verify that it has a
  // compatible type otherwise.
  size_t addOrCheckAttribute(
      const std::string& name,
      TypePtr ty,
      bool is_parameter = false) {
    auto slot_idx = findAttributeSlot(name);
    if (!slot_idx) {
      return addAttribute(name, ty, is_parameter);
    }

    TORCH_CHECK(
        is_parameter == this->is_parameter(*slot_idx),
        "Parameter field mismatch for the field '",
        name,
        "'");
    TypePtr atype = getAttribute(*slot_idx);
    TORCH_CHECK(
      ty->isSubtypeOf(atype),
      ty->python_str(),
      " is not compatible with the type ",
      atype->python_str(),
      " for the field '",
      name,
      "'");
    return *slot_idx;
  }

  at::ArrayRef<std::string> attributeNames() const {
    return attributeNames_;
  }

  at::ArrayRef<TypePtr> containedTypes() const override {
    return attributeTypes_;
  }

  TypePtr createWithContained(std::vector<TypePtr> contained_types) const override {
    auto ptr = ClassType::create(name(), compilation_unit_);
    AT_ASSERT(numAttributes() == contained_types.size());
    for(size_t i = 0; i < attributeNames_.size(); ++i) {
      AT_ASSERT(attributeTypes_[i]->isSubtypeOf(contained_types[i]));
      ptr->addAttribute(attributeNames_[i], contained_types[i]);
    }
    // Copy methods over
    for (const auto& method : methods()) {
      ptr->addMethod(method);
    }
    return ptr;
  }

  bool is_module() const override {
    return bool(parameterSlots_);
  }
  bool is_parameter(size_t slot) const {
    TORCH_INTERNAL_ASSERT(
        is_module(), "asking for parameterSlots of non-Module");
    return parameterSlots_->at(slot);
  }

  void addMethod(Function* method);
  Function* getMethod(const std::string& name) const;

  std::shared_ptr<CompilationUnit> compilation_unit();
  std::shared_ptr<const CompilationUnit> compilation_unit() const;

  // generate a refined version of this class.
  // It has the same name but the slot Types are subtypes of
  // the original slots. It is only valid to refine a class type in a context
  // where it is know that there are not assignments to the objects slots
  // that would invalidate the refinement.
  // These variants are not registered in the global class table.
  ClassTypePtr refine(at::ArrayRef<TypePtr> refined_slots) const;

  bool isSubtypeOfExt(const TypePtr rhs, std::ostream* why_not) const override;

  static const TypeKind Kind = TypeKind::ClassType;

 private:
  ClassType(
      c10::optional<QualifiedName> name,
      std::weak_ptr<CompilationUnit> cu,
      bool is_module);

  // Mapping of attribute names -> their type.
  // NOTE: this does not contain methods, which are stored in the module
  // TODO: once modules support arbitrary ivalue attributes, we don't need this
  // anymore.
  // TODO: This is better represented as an OrderedDict, but alas it is not yet
  // available from c10
  std::vector<std::string> attributeNames_;
  std::vector<TypePtr> attributeTypes_;
  // Holds method attributes
  std::weak_ptr<CompilationUnit> compilation_unit_;

  // if present, this class inherits from torch.nn.Module
  // and these are the indices of the attributes which are parameters
  std::shared_ptr<std::vector<bool>> parameterSlots_;

  // List of methods associated with this class.
  std::vector<Function*> methods_;

};

struct InterfaceType;
using InterfaceTypePtr = std::shared_ptr<InterfaceType>;
using ::torch::jit::script::CompilationUnit;
using ::torch::jit::Function;

// Interfaces are a list of abstract methods that a class might meet.
// If a class provides those methods, it implicitly meets the interface.

// Subtype relations for Interface with ClassType:
// lhs (ClassType or InterfaceType) is a subtype of rhs if:
// 1. lhs methods are a superset of rhs methods
// 2. if rhs is module interface, the lhs must be module interface or module itself
struct CAFFE2_API InterfaceType : public NamedType {
  static InterfaceTypePtr create(
      QualifiedName qualifiedName, bool is_module=false);

  bool operator==(const Type& rhs) const override {
    if (auto user_rhs = rhs.cast<InterfaceType>()) {
      return name() == user_rhs->name();
    }
    return false;
  }

  std::string str() const override {
    return std::string("InterfaceType<") + name()->name() + ">";
  }

  std::string python_str() const override {
    return name()->qualifiedName();
  }

  bool isSubtypeOfExt(const TypePtr rhs, std::ostream* why_not) const override;

  // try to find a method of this interface,
  // returns nullptr if not found.
  const FunctionSchema* getMethod(const std::string& name) const;
  void addMethod(FunctionSchema schema);
  const std::vector<FunctionSchema>& methods() {
    return *methods_;
  }

  bool is_module() const override{
    return is_module_;
  }
  static const TypeKind Kind = TypeKind::InterfaceType;
  ~InterfaceType() override;
 private:
  InterfaceType(QualifiedName name, bool is_module);

  // shared_ptr so that this header does not have to depend on
  // FunctionSchema.h
  std::shared_ptr<std::vector<FunctionSchema>> methods_;
  // flag to distinguish if it's an interface type from a module or not
  bool is_module_;
};

} // namespace c10
