/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

/*!
 * \file tvm/te/tensor.h
 * \brief Dataflow tensor object
 */
#ifndef TVM_TE_TENSOR_H_
#define TVM_TE_TENSOR_H_

#include <tvm/arith/bound.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>

#include <string>
#include <type_traits>
#include <utility>
#include <vector>

namespace tvm {
namespace te {

using arith::IntSet;
using namespace tvm::tirx;

// internal node container for Operation
class OperationNode;
class Tensor;

/*! \brief Operation that produces tensors */
class Operation : public ffi::ObjectRef {
 public:
  /*! \brief default constructor  */
  Operation() {}
  explicit Operation(ffi::ObjectPtr<ffi::Object> n) : ffi::ObjectRef(n) {}
  explicit Operation(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
  /*!
   * \brief access the internal node container
   * \return the pointer to the internal node container
   */
  inline const OperationNode* operator->() const;
  /*!
   * \brief get the i-th output of the operation.
   * \param i the output index.
   * \return The i-th output.
   */
  TVM_DLL Tensor output(size_t i) const;
  /*! \brief specify container node */
  using ContainerType = OperationNode;
};

/*! \brief Node to represent a tensor */
class TensorNode : public DataProducerNode {
 public:
  /*! \brief The shape of the tensor */
  ffi::Array<PrimExpr> shape;
  /*! \brief dtype in the content of the tensor */
  PrimType dtype = PrimType::Void();
  /*! \brief the source operation, can be None */
  Operation op;
  /*! \brief the output index from source operation */
  int value_index{0};

  static void RegisterReflection();

  ffi::Array<PrimExpr> GetShape() const final { return shape; }

  PrimType GetDataType() const final { return dtype; }

  TVM_DLL PrimExpr ToPrimExpr() const final;

  TVM_DLL ffi::String GetNameHint() const final;

  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode;

  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.Tensor", TensorNode, DataProducerNode);
};

/*!
 * \brief Tensor structure representing a possible input,
 *  or intermediate computation result.
 */
class Tensor : public DataProducer {
 private:
  /*!
   * \brief Helper for indexing operations into tensors
   * \param indices The indices
   * \param support_negative_indices Whether to normalize indices in the case of negative indices.
   * \return the result expression representing tensor read.
   */
  inline PrimExpr IndexTensor(ffi::Array<PrimExpr> indices, bool support_negative_indices) const;

 public:
  TVM_DLL Tensor(ffi::Array<PrimExpr> shape, PrimType dtype, Operation op, int value_index);
  /*!
   * \brief check if two tensors equals each other.
   * \param other tensor to be checked.
   * \return whether the two tensors equals each other.
   */
  inline bool operator==(const Tensor& other) const;
  /*!
   * \brief check if two tensors are different.
   * \param other tensor to be checked.
   * \return whether the two tensors are different.
   */
  inline bool operator!=(const Tensor& other) const;
  /*! \return The dimension of the tensor */
  inline size_t ndim() const;
  /*!
   * \brief Take elements from the tensor
   * \param args The indices
   * \return the result expression representing tensor read.
   */
  template <typename... Args>
  inline PrimExpr operator()(Args&&... args) const {
    ffi::Array<PrimExpr> indices{std::forward<Args>(args)...};
    return operator()(indices);
  }
  /*!
   * \brief Take elements from the tensor
   * \param indices the indices.
   * \return the result expression representing tensor read.
   */
  TVM_DLL PrimExpr operator()(ffi::Array<PrimExpr> indices) const;
  /*!
   * \brief Take elements from the tensor
   * \param indices the indices.
   * \return the result expression representing tensor read.
   */
  TVM_DLL PrimExpr operator()(ffi::Array<PrimVar> indices) const;
  /*!
   * \brief Take elements from the tensor with support for negative indices.
   * \param args The indices
   * \return the result expression representing tensor read.
   */
  template <typename... Args>
  TVM_DLL PrimExpr IndexWithNegativeIndices(Args&&... args) const {
    ffi::Array<PrimExpr> indices{std::forward<Args>(args)...};
    return IndexWithNegativeIndices(indices);
  }
  /*!
   * \brief Take elements from the tensor with support for negative indices.
   * \param indices the indices.
   * \return the result expression representing tensor read.
   */
  TVM_DLL PrimExpr IndexWithNegativeIndices(ffi::Array<PrimExpr> indices) const;
  /*!
   * \brief Take elements from the tensor with support for negative indices.
   * \param indices the indices.
   * \return the result expression representing tensor read.
   */
  TVM_DLL PrimExpr IndexWithNegativeIndices(ffi::Array<PrimVar> indices) const;

  /*!
   * \brief data structure to represent a slice that fixes first k coordinates.
   *  This is used to enable syntax sugar of Tensor[x][y][z] to get the element.
   */
  class Slice {
   public:
    // construct via tensor and indices
    Slice(const Tensor& tensor, std::vector<PrimExpr> indices)
        : tensor_(tensor), indices_(indices) {}
    /*!
     * \brief get i-th slice from the current slice.
     * \param i the index of the coordinate
     * \return the subsequent slice.
     */
    inline Slice operator[](PrimExpr i) {
      std::vector<PrimExpr> other = indices_;
      other.emplace_back(i);
      return Slice(tensor_, other);
    }
    /*!
     * \brief Convert slice to expression.
     *  This is only valid when all the coordinates are fully specified.
     * \return the corresponding expression of this slice.
     */
    inline operator PrimExpr() const { return tensor_(indices_); }

   private:
    const Tensor& tensor_;
    std::vector<PrimExpr> indices_;
  };
  /*!
   * \brief get i-th slice from the current Tensor.
   * \param i the index of the coordinate
   * \return the subsequent slice.
   */
  inline Slice operator[](PrimExpr i) const { return Slice(*this, {i}); }

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Tensor, DataProducer, TensorNode);
};

// Implementations of inline functions
inline size_t Tensor::ndim() const { return (*this)->shape.size(); }

inline bool Tensor::operator==(const Tensor& other) const {
  if (get() == other.get()) return true;
  if (get() == nullptr || other.get() == nullptr) return false;
  if ((*this)->op.defined() || other->op.defined()) {
    return (*this)->op == other->op && (*this)->value_index == other->value_index;
  } else {
    return false;
  }
}

inline bool Tensor::operator!=(const Tensor& other) const { return !(*this == other); }

// macro to turn every operation of slice to expression
#define DEFINE_OVERLOAD_SLICE_UNARY_OP(Op) \
  inline PrimExpr operator Op(const Tensor::Slice& a) { return Op a.operator PrimExpr(); }

#define DEFINE_OVERLOAD_SLICE_BINARY_OP(Op)                                     \
  template <typename T>                                                         \
  inline PrimExpr operator Op(const Tensor::Slice& a, const T& b) {             \
    return a.operator PrimExpr() Op b;                                          \
  }                                                                             \
  template <typename T>                                                         \
  inline PrimExpr operator Op(const T& a, const Tensor::Slice& b) {             \
    return a Op b.operator PrimExpr();                                          \
  }                                                                             \
  inline PrimExpr operator Op(const Tensor::Slice& a, const Tensor::Slice& b) { \
    return a.operator PrimExpr() Op b.operator PrimExpr();                      \
  }

DEFINE_OVERLOAD_SLICE_UNARY_OP(!);
DEFINE_OVERLOAD_SLICE_UNARY_OP(-);
DEFINE_OVERLOAD_SLICE_BINARY_OP(+);
DEFINE_OVERLOAD_SLICE_BINARY_OP(-);
DEFINE_OVERLOAD_SLICE_BINARY_OP(*);
DEFINE_OVERLOAD_SLICE_BINARY_OP(==);
DEFINE_OVERLOAD_SLICE_BINARY_OP(<=);
DEFINE_OVERLOAD_SLICE_BINARY_OP(>=);
DEFINE_OVERLOAD_SLICE_BINARY_OP(!=);
DEFINE_OVERLOAD_SLICE_BINARY_OP(&&);
DEFINE_OVERLOAD_SLICE_BINARY_OP(||);
DEFINE_OVERLOAD_SLICE_BINARY_OP(>>);
DEFINE_OVERLOAD_SLICE_BINARY_OP(<<);
DEFINE_OVERLOAD_SLICE_BINARY_OP(>);  // NOLINT(*)
DEFINE_OVERLOAD_SLICE_BINARY_OP(<);  // NOLINT(*)

}  // namespace te
}  // namespace tvm

namespace std {
template <>
struct hash<::tvm::te::Operation> : public ::tvm::ffi::ObjectPtrHash {};

template <>
struct hash<::tvm::te::Tensor> {
  std::size_t operator()(const ::tvm::te::Tensor& k) const {
    ::tvm::ffi::ObjectPtrHash hasher;
    if (k.defined() && k->op.defined()) {
      return hasher(k->op);
    } else {
      return hasher(k);
    }
  }
};
}  // namespace std
#endif  // TVM_TE_TENSOR_H_