#pragma once

#include <ATen/Tensor.h>
#include <c10/core/SymIntArrayRef.h>
#include <c10/core/TensorImpl.h>

#include <torch/csrc/lazy/core/tensor.h>

namespace torch {
namespace lazy {

// Tensor implementation class used to be fed to the at::Tensor.
// Its scope is just to handle an LazyTensor.
class TORCH_API LTCTensorImpl final : public c10::TensorImpl {
 public:
  explicit LTCTensorImpl(const LazyTensorPtr& tensor);
  explicit LTCTensorImpl(const LazyTensor& tensor);
  explicit LTCTensorImpl(LazyTensor&& tensor);

  LazyTensorPtr tensor() { return tensor_; }

  void set_tensor(const LazyTensorPtr& lazy_tensor);

  void force_refresh_sizes() { generation_ = 0; }

  c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach(
      const c10::VariableVersion& version_counter,
      bool allow_tensor_metadata_change) const override;

  c10::intrusive_ptr<TensorImpl> shallow_copy_and_detach(
      c10::VariableVersion&& version_counter,
      bool allow_tensor_metadata_change) const override;

  void shallow_copy_from(const c10::intrusive_ptr<TensorImpl>& impl) override;

  at::IntArrayRef sizes_custom() const override;
  at::IntArrayRef strides_custom() const override;
  int64_t dim_custom() const override;
  int64_t numel_custom() const override;
  bool is_contiguous_custom(at::MemoryFormat memory_format) const override;

  virtual c10::SymIntArrayRef sym_sizes_custom() const override;

#ifndef C10_DISABLE_TENSORIMPL_EXTENSIBILITY
  const at::Storage& storage() const override { return tensor_->Storage(); }
  bool has_storage() const override { return tensor_->Storage(); }
#endif  // C10_DISABLE_TENSORIMPL_EXTENSIBILITY

 private:
  void setup_size_properties();

  LazyTensorPtr tensor_;
  std::vector<c10::SymInt> sym_sizes_;
  size_t generation_ {0};
};

}  // namespace lazy
}  // namespace torch
