#ifndef CAFFE2_CORE_TENSOR_H_
#define CAFFE2_CORE_TENSOR_H_

#include <c10/macros/Macros.h>
#include "caffe2/core/storage.h"

#include <ATen/core/UndefinedTensorImpl.h>
#include <c10/core/TensorOptions.h>
#include <c10/util/intrusive_ptr.h>

#if defined(EXPOSE_C2_OPS) || \
    !defined(CAFFE2_IS_XPLAT_BUILD) && !defined(C10_MOBILE)
namespace at {
class Tensor;
};
#endif
namespace caffe2 {

using at::UndefinedTensorImpl;

/**
 * @brief Tensor class holds a shared pointer to the implementation TensorImpl,
 * redirects API calls to TensorImpl;
 * Copying of Tensor results in sharing the same underlying implementation
 * object
 *
 * NB: See TensorImpl for documentation on these methods.
 */
class TORCH_API Tensor final {
 private:
  enum Unsafe { IDoWantAliasing };
  Tensor(const Tensor& other, Unsafe _) : impl_(other.getIntrusivePtr()) {}

 protected:
  using TensorImplPtr = c10::intrusive_ptr<TensorImpl, UndefinedTensorImpl>;
  TensorImplPtr impl_;

  void enforce_invariants();

 public:
  Tensor() : impl_() {}

  Tensor(const Tensor& t) : impl_(t.impl_) {}
  Tensor& operator=(const Tensor& t) {
    impl_ = t.impl_;
    return *this;
  }

  Tensor(Tensor&&) = default;
  Tensor& operator=(Tensor&&) = default;

  operator bool() const {
    return impl_.defined();
  }

  TensorImpl* unsafeGetTensorImpl() const {
    return impl_.get();
  }

  Tensor UnsafeSharedInstance() const {
    return Tensor(*this, IDoWantAliasing);
  }

  /**
   * @brief Creates a tensor of the given device type.
   *
   * Note that the actual data allocation is not going to be carried out until
   * you resize the tensor and then call mutable_data().
   */
  explicit Tensor(at::Device device)
      : impl_(c10::make_intrusive<TensorImpl, UndefinedTensorImpl>(
            Storage::create_legacy(device),
            c10::computeDispatchKey(c10::nullopt, at::kStrided, device),
            TypeMeta())) {}

  /**
   * @brief Creates a tensor of the given dimension.
   *
   * Note that the actual data allocation is not going to be carried out until
   * the first time mutable_data() is called.
   */
  explicit Tensor(at::IntArrayRef dims, DeviceType type) : Tensor(type) {
    // TODO: here, we create a Storage
    // and immediately discard it in Resize() since
    // reset_tensor will be true and FreeMemory will be called,
    // we might want to avoid creating Storage twice?
    Resize(dims);
  }

  // we want to preserve index information
  explicit Tensor(at::IntArrayRef dims, at::Device device) : Tensor(device) {
    Resize(dims);
  }

  // TODO: remove?
  explicit Tensor(const vector<int>& dims, DeviceType type) : Tensor(type) {
    Resize(dims);
  }

  /**
   * @brief: Create a Tensor of at::DeviceType `type` and initialize it with
   * src Tensor
   */
  Tensor(const Tensor& src, DeviceType type) : Tensor(type) {
    CopyFrom(src);
  }

  /**
   * @brief Mutual conversion with at::Tensor
   *
   * The tensor will share the same instance (data, strides, sizes, etc) but
   * a different subset of APIs would be available
   */
#if defined(EXPOSE_C2_OPS) || \
    !defined(CAFFE2_IS_XPLAT_BUILD) && !defined(C10_MOBILE)
  explicit Tensor(at::Tensor tensor);

  explicit operator at::Tensor() const&;

  explicit operator at::Tensor() &&;
#endif

  bool is_same(const Tensor& other) const noexcept {
    return impl_ == other.impl_;
  }

  Tensor Clone() const {
    Tensor x(GetDevice());
    x.CopyFrom(*this);
    return x;
  }

  /**
   * Clone self as a Tensor that share the same Storage,
   * that is, both Tensors are views on the same Storage.
   * If we change the sizes or strides of one Tensor, it
   * does not affect the other Tensor that it shares Storage
   * with.
   * A similar yet different usage is `Tensor x = y;`, this
   * will make x and y pointing to the same Tensor and resizing
   * one of them will resize the other as well.
   *
   * TODO: Deduplicate this with THTensor_(newWithTensor)
   * (exposed in ATen as at::alias but not otherwise available)
   */
  Tensor Alias() const {
    Tensor x(sizes(), GetDevice());
    if (!dtype_initialized()) {
      C10_LOG_EVERY_MS(WARNING, 1000)
          << "Cloning a tensor that don't have a data type (did you call mutable_data<T> on the tensor?)";
    }
    AT_ASSERTM(
        storage_initialized(),
        "Cloning a tensor that has no content and has size > 0");
    // set_storage already sets data_type_ of TensorImpl
    x.impl_->set_storage_and_dtype(storage(), impl_->dtype());
    x.impl_->set_storage_offset(impl_->storage_offset());
    x.impl_->set_sizes_and_strides(sizes(), strides());
    return x;
  }

  DeviceType GetDeviceType() const {
    return impl_->device_type();
  }

  at::Device GetDevice() const {
    return impl_.get()->device();
  }

  /**
   * @brief Copies the data from a source tensor, with a context provided to
   * carry out the underlying memcpy operation.  This method respects
   * caffe2_keep_on_shrink.
   *
   * After CopyFrom, this function guarantees that the destination tensor will
   * have the same initialization state and dtype as src.  This function
   * preserves the DeviceType of the source tensor (so, e.g., if you allocate
   * a tensor on CPU and then CopyFrom a CUDA tensor, that will to a
   * CUDA-to-CPU transfer).
   *
   * 'async' parameter triggers async copy for CUDA tensors
   */
  void CopyFrom(const Tensor& src, bool async = false);

  /**
   * @brief Extend the outer-most dimension of this tensor
   *        to dimension of `num`.
   */
  void ExtendTo(int64_t num, float growthPct) const {
    CAFFE_ENFORCE_GE_WITH_CALLER(impl_->dim(), 1);
    CAFFE_ENFORCE_GE_WITH_CALLER(growthPct, 0);
    Extend(num - impl_->size(0), growthPct);
  }

  void Extend(int64_t num, float growthPct) const {
    impl_.get()->Extend(num, growthPct);
  }

  /**
   * @brief Shrinks the outer-most dimension to given size, keeping the data.
   *
   * This method guarantees that no re-allocations are carried out, which means
   * that the extra capacity after the end of the shrunk tensor is maintained.
   * Notably, this function does NOT respect caffe2_keep_on_shrink.
   */
  void ShrinkTo(int64_t outer_dim) const {
    CAFFE_ENFORCE_WITH_CALLER(
        impl_->is_contiguous(),
        "Right now ShrinkTo is only supported on contiguous Tensor.");
    CAFFE_ENFORCE_WITH_CALLER(impl_->dim() >= 1, "Tensor must be at least 1D");
    CAFFE_ENFORCE_WITH_CALLER(
        outer_dim <= impl_->size(0),
        "New outer dimension must be smaller than current.");
    CAFFE_ENFORCE(
        impl_->storage().unique(),
        "Can't call ShrinkTo on shared storage, please call Resize instead.");
    impl_.get()->set_size(0, outer_dim);
  }

  template <class T>
  void ReserveSpace(const T& outer_dim) const {
    impl_.get()->ReserveSpace(outer_dim);
  }

  template <typename... Ts>
  void Resize(Ts... dim_source) const {
    impl_.get()->Resize(dim_source...);
  }

  /**
   * Resize the tensor like the source tensor. Note that this is just a
   * sugar wrapper that essentially calls Resize(src_tensor.dims()).
   * This method respects caffe2_keep_on_shrink.
   */
  inline void ResizeLike(const Tensor& src_tensor) const {
    CAFFE_ENFORCE_WITH_CALLER(
        src_tensor.is_contiguous(),
        "Right now ResizeLike is only supported for contiguous Tensor.");
    if (impl_ != src_tensor.impl_) {
      impl_.get()->Resize(src_tensor.sizes());
    }
  }

  inline void Reshape(const vector<int64_t>& dims) const {
    impl_.get()->Reshape(dims);
  }

  inline void Reshape(const vector<int>& dims) const {
    impl_.get()->Reshape(ToVectorint64_t(dims));
  }

  inline void FreeMemory() const {
    impl_.get()->FreeMemory();
  }

  /**
   * A utility function to print the debug string for the tensor. Note that this
   * is very slow since it involves quite some string operations, so do not use
   * it in your performance-critical code.
   */
  string DebugString() const {
    std::stringstream ss;
    ss << "A Tensor of item size " << impl_->dtype().itemsize() << " and type "
       << impl_->dtype().name() << " and dimension (";
    for (int d : impl_->sizes()) {
      ss << d << ",";
    }
    ss << ").";
    return ss.str();
  }

  // To be deprecated
  void ShareData(const Tensor& src) const {
    impl_.get()->ShareData(*src.impl_.get());
  }

  /**
   * @brief Shares the data with an externally managed pointer.
   *
   * This is similar to ShareData() but the source is a pointer with an advanced
   * deleter option. In default, no deletion takes place, and one needs to make
   * sure that the external memory is deallocated only after the tensor finishes
   * using it. If a Deleter object is passed in, when this tensor is reallocated
   * or freed, the deleter function is going to be called.
   */
  template <typename T>
  void ShareExternalPointer(
      T* src,
      size_t nbytes = 0,
      MemoryDeleter d = nullptr) const {
    ShareExternalPointer((void*)src, caffe2::TypeMeta::Make<T>(), nbytes, d);
  }

  template <typename T>
  void ShareExternalPointer(at::DataPtr&& data_ptr, size_t nbytes = 0) const {
    ShareExternalPointer(
        std::move(data_ptr), caffe2::TypeMeta::Make<T>(), nbytes);
  }

  void ShareExternalPointer(
      void* src,
      const TypeMeta data_type,
      size_t nbytes = 0,
      MemoryDeleter d = nullptr) const {
    CAFFE_ENFORCE_WITH_CALLER(
        impl_->is_contiguous(),
        "Right now ShareExternalPointer is only supported for contiguous Tensor.");
    CAFFE_ENFORCE_WITH_CALLER(
        data_type != ScalarType::Undefined,
        "To share with a raw external pointer you need to pass in an "
        "initialized data_type(TypeMeta).");
    impl_.get()->ShareExternalPointer(
        at::DataPtr(src, src, d, impl_->device_type()), data_type, nbytes);
  }

  void ShareExternalPointer(
      at::DataPtr&& data_ptr,
      const TypeMeta data_type,
      size_t nbytes) {
    impl_.get()->ShareExternalPointer(std::move(data_ptr), data_type, nbytes);
  }

  const c10::intrusive_ptr<TensorImpl, UndefinedTensorImpl>& getIntrusivePtr()
      const {
    return impl_;
  }

  bool defined() const {
    return impl_;
  }

  /**
   * Returns a raw void* pointer of the underlying storage. mutable_data()
   * or raw_mutable_data() must have been called prior to this function call.
   */
  inline void* raw_data() const {
    return impl_->data();
  }

  template <typename T>
  inline T* data() const {
    return impl_.get()->data<T>();
  }

  inline void* raw_mutable_data(const TypeMeta meta) const {