/*
 * 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/ffi/container/shape.h
 * \brief Container to store shape of an Tensor.
 */
#ifndef TVM_FFI_CONTAINER_SHAPE_H_
#define TVM_FFI_CONTAINER_SHAPE_H_

#include <tvm/ffi/container/array.h>
#include <tvm/ffi/error.h>
#include <tvm/ffi/type_traits.h>

#include <algorithm>
#include <ostream>
#include <utility>
#include <vector>

namespace tvm {
namespace ffi {

/*!
 * \brief Lightweight view non-owning class for shape.
 */
class ShapeView {
 public:
  /*! \brief Default constructor. */
  ShapeView() : cell_{nullptr, 0} {}
  /*! \brief Copy constructor. */
  ShapeView(const ShapeView& other) = default;
  /*! \brief Copy assignment operator. */
  ShapeView& operator=(const ShapeView& other) = default;
  /*! \brief Move constructor. */
  ShapeView(ShapeView&& other) = default;
  /*! \brief Move assignment operator. */
  ShapeView& operator=(ShapeView&& other) = default;
  /*! \brief Constructor from data and size. */
  ShapeView(const int64_t* data, size_t size) : cell_{data, size} {}
  /*! \brief Constructor from initializer list. */
  ShapeView(const std::initializer_list<int64_t>& other) : cell_{other.begin(), other.size()} {}
  /*! \brief Get the data pointer. */
  const int64_t* data() const { return cell_.data; }
  /*! \brief Get the size of the shape. */
  size_t size() const { return cell_.size; }

  /*! \brief Get the product of the shape. */
  int64_t Product() const {
    int64_t product = 1;
    for (size_t i = 0; i < cell_.size; ++i) {
      product *= cell_.data[i];
    }
    return product;
  }

  /*! \brief Get the i-th element of the shape. */
  int64_t operator[](size_t idx) const { return cell_.data[idx]; }

  /*! \return begin iterator */
  const int64_t* begin() const { return cell_.data; }

  /*! \return end iterator */
  const int64_t* end() const { return cell_.data + cell_.size; }

  /*! \return Whether shape tuple is empty */
  bool empty() const { return size() == 0; }

  /*! \return The first element of the shape tuple */
  int64_t front() const { return this->at(0); }

  /*! \return The last element of the shape tuple */
  int64_t back() const { return this->at(this->size() - 1); }

  /*! \brief Get the i-th element of the shape. */
  int64_t at(size_t idx) const {
    if (idx >= this->size()) {
      TVM_FFI_THROW(IndexError) << "indexing " << idx << " on a Shape of size " << this->size();
    }
    return cell_.data[idx];
  }

 private:
  TVMFFIShapeCell cell_;
};

/*! \brief An object representing a shape tuple. */
class ShapeObj : public Object, public TVMFFIShapeCell {
 public:
  /*! \brief The type of shape index element. */
  using index_type = int64_t;

  /*! \brief Get "numel", meaning the number of elements of an array if the array has this shape */
  int64_t Product() const {
    int64_t product = 1;
    for (size_t i = 0; i < this->size; ++i) {
      product *= this->data[i];
    }
    return product;
  }

  /// \cond Doxygen_Suppress
  static constexpr const uint32_t _type_index = TypeIndex::kTVMFFIShape;
  TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIShape, ShapeObj, Object);
  /// \endcond
};

namespace details {

class ShapeObjStdImpl : public ShapeObj {
 public:
  explicit ShapeObjStdImpl(std::vector<int64_t> other) : data_{std::move(other)} {
    this->data = data_.data();
    this->size = static_cast<size_t>(data_.size());
  }

 private:
  std::vector<int64_t> data_;
};

TVM_FFI_INLINE ObjectPtr<ShapeObj> MakeEmptyShape(size_t length, int64_t** mutable_data) {
  ObjectPtr<ShapeObj> p = make_inplace_array_object<ShapeObj, int64_t>(length);
  static_assert(alignof(ShapeObj) % alignof(int64_t) == 0);
  static_assert(sizeof(ShapeObj) % alignof(int64_t) == 0);
  int64_t* data = reinterpret_cast<int64_t*>(reinterpret_cast<char*>(p.get()) + sizeof(ShapeObj));
  if (mutable_data) {
    *mutable_data = data;
  }
  p->data = data;
  p->size = length;
  return p;
}

// inplace shape allocation
template <typename IterType>
TVM_FFI_INLINE ObjectPtr<ShapeObj> MakeInplaceShape(IterType begin, IterType end) {
  size_t length = std::distance(begin, end);
  int64_t* mutable_data;
  ObjectPtr<ShapeObj> p = MakeEmptyShape(length, &mutable_data);
  std::copy(begin, end, mutable_data);
  return p;
}

/*!
 * \brief Get the product of a shape.
 * \param shape The input shape.
 * \param out_strides The output strides.
 * \return The product of the shape.
 */
TVM_FFI_INLINE void FillStridesFromShape(ShapeView shape, int64_t* out_strides) {
  int64_t stride = 1;
  for (int64_t i = static_cast<int64_t>(shape.size()) - 1; i >= 0; --i) {
    out_strides[i] = stride;
    stride *= shape[i];
  }
}

/*!
 * \brief Make a strides shape from a shape view.
 * \param shape The input shape.
 * \return The shape.
 */
TVM_FFI_INLINE ObjectPtr<ShapeObj> MakeStridesFromShape(ShapeView shape) {
  int64_t* strides_data;
  ObjectPtr<ShapeObj> strides = details::MakeEmptyShape(shape.size(), &strides_data);
  FillStridesFromShape(shape, strides_data);
  return strides;
}

}  // namespace details

/*!
 * \brief Managed reference to shape object.
 *
 * When possible, use ShapeView instead of Shape to reduce memory allocation.
 * Use Shape when you need to have a managed reference to on-heap allocated shape.
 *
 * \sa ShapeView
 */
class Shape : public ObjectRef {
 public:
  /*! \brief The type of shape index element. */
  using index_type = ShapeObj::index_type;

  /*! \brief Default constructor */
  Shape() : ObjectRef(details::MakeEmptyShape(0, nullptr)) {}

  /*!
   * \brief Constructor from iterator
   * \param begin begin of iterator
   * \param end end of iterator
   * \tparam IterType The type of iterator
   */
  template <typename IterType>
  Shape(IterType begin, IterType end) : Shape(details::MakeInplaceShape(begin, end)) {}

  /**
   * \brief Constructor from Array<int64_t>
   * \param shape The Array<int64_t>
   *
   * \note This constructor will copy the data content.
   */
  Shape(Array<int64_t> shape)  // NOLINT(*)
      : Shape(shape.begin(), shape.end()) {}

  /*!
   * \brief constructor from initializer list
   * \param shape The initializer list
   */
  Shape(std::initializer_list<int64_t> shape) : Shape(shape.begin(), shape.end()) {}

  /*!
   * \brief constructor from int64_t [N]
   *
   * \param other a int64_t array.
   */
  Shape(std::vector<int64_t> other)  // NOLINT(*)
      : ObjectRef(make_object<details::ShapeObjStdImpl>(std::move(other))) {}

  /*!
   * \brief constructor from shape view.
   * \param other The shape view.
   */
  Shape(ShapeView other) : Shape(other.begin(), other.end()) {}  // NOLINT(*)

  /*!
   * \brief Create a strides from a shape.
   * \param shape The input shape.
   * \return The strides.
   */
  static Shape StridesFromShape(ShapeView shape) {
    return Shape(details::MakeStridesFromShape(shape));
  }

  /*!
   * \brief Convert to shape view.
   * \return The shape view.
   */
  operator ShapeView() const { return ShapeView(data(), size()); }  // NOLINT(*)

  /*!
   * \brief Return the data pointer
   *
   * \return const index_type* data pointer
   */
  const int64_t* data() const { return get()->data; }

  /*!
   * \brief Return the size of the shape tuple
   *
   * \return size_t shape tuple size
   */
  size_t size() const { return get()->size; }

  /*!
   * \brief Immutably read i-th element from the shape tuple.
   * \param idx The index
   * \return the i-th element.
   */
  int64_t operator[](size_t idx) const { return this->data()[idx]; }

  /*!
   * \brief Immutably read i-th element from the shape tuple.
   * \param idx The index
   * \return the i-th element.
   */
  int64_t at(size_t idx) const {
    if (idx >= this->size()) {
      TVM_FFI_THROW(IndexError) << "indexing " << idx << " on a Shape of size " << this->size();
    }
    return this->operator[](idx);
  }

  /*! \return Whether shape tuple is empty */
  bool empty() const { return size() == 0; }

  /*! \return The first element of the shape tuple */
  int64_t front() const { return this->at(0); }

  /*! \return The last element of the shape tuple */
  int64_t back() const { return this->at(this->size() - 1); }

  /*! \return begin iterator */
  const int64_t* begin() const { return get()->data; }

  /*! \return end iterator */
  const int64_t* end() const { return (get()->data + size()); }

  /*! \return The product of the shape tuple */
  int64_t Product() const { return get()->Product(); }

  /// \cond Doxygen_Suppress
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Shape, ObjectRef, ShapeObj);
  /// \endcond

 private:
  explicit Shape(ObjectPtr<ShapeObj> ptr) : ObjectRef(std::move(ptr)) {}
};

inline std::ostream& operator<<(std::ostream& os, const Shape& shape) {
  os << '[';
  for (size_t i = 0; i < shape.size(); ++i) {
    if (i != 0) {
      os << ", ";
    }
    os << shape[i];
  }
  os << ']';
  return os;
}

// Shape
template <>
inline constexpr bool use_default_type_traits_v<Shape> = false;

// Allow auto conversion from Array<int64_t> to Shape, but not from Shape to Array<int64_t>
template <>
struct TypeTraits<Shape> : public ObjectRefWithFallbackTraitsBase<Shape, Array<int64_t>> {
  static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIShape;
  TVM_FFI_INLINE static Shape ConvertFallbackValue(Array<int64_t> src) {
    return Shape(std::move(src));
  }
};

}  // namespace ffi
}  // namespace tvm

#endif  // TVM_FFI_CONTAINER_SHAPE_H_
