/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

#pragma once

#include <executorch/runtime/core/exec_aten/exec_aten.h>
#include <executorch/runtime/core/result.h>

namespace executorch {
namespace extension {
namespace internal {

/**
 * Base class template storing the underlying data with size and stride helpers.
 * Inherited by TensorAccessor<> which requires specialization on rank.
 */
template <typename T, ssize_t N>
class TensorAccessorBase {
 public:
  /// Returns the size of the underlying tensor at the given dimension.
  executorch::aten::SizesType size(ssize_t i) const {
    ET_CHECK_MSG(
        i < dim_ && i >= 0,
        "Dimension outside of [0, %zd], got %zd",
        dim_ - 1,
        i);
    return sizes_[i];
  }

  /// Returns the stride of the underlying tensor at the given dimension.
  executorch::aten::StridesType stride(ssize_t i) const {
    ET_CHECK_MSG(
        i < dim_ && i >= 0,
        "Dimension outside of [0, %zd], got %zd",
        dim_ - 1,
        i);
    return strides_[i];
  }

 protected:
  TensorAccessorBase(
      T* data,
      const executorch::aten::SizesType* sizes,
      const executorch::aten::StridesType* strides,
      ssize_t dim)
      : data_(data), sizes_(sizes), strides_(strides), dim_(dim) {}

  T* data_;
  const executorch::aten::SizesType* sizes_;
  const executorch::aten::StridesType* strides_;
  ssize_t dim_;
};

} // namespace internal

/**
 * TensorAccessor template with data type and rank as template parameters. No
 * public constructors, can only be created using make_tensor_accessor from a
 * given executorch::aten::Tensor. Use operator[] to index and obtain a lower
 * rank accessor or the underlying scalar value.
 */
template <typename T, ssize_t N>
class TensorAccessor : public internal::TensorAccessorBase<T, N> {
 public:
  /**
   * Index into the the outer most dimension.
   *
   * @param i Index.
   * @return If N > 1, a TensorAccessor with N-1 dimensions. If N == 1, a
   * reference to the underlying scalar. Refer to the TensorAccessor<T, 1>
   * specialization.
   */
  TensorAccessor<T, N - 1> operator[](ssize_t i) {
    return TensorAccessor<T, N - 1>(
        this->data_ + this->strides_[0] * i,
        this->sizes_ + 1,
        this->strides_ + 1,
        N - 1);
  }

  /**
   * Index into the the outer most dimension.
   *
   * @param i Index.
   * @return If N > 1, a constant TensorAccessor with N-1 dimensions. If N == 1,
   * a constant reference to the underlying scalar. Refer to the
   * TensorAccessor<T, 1> specialization.
   */
  const TensorAccessor<T, N - 1> operator[](ssize_t i) const {
    return TensorAccessor<T, N - 1>(
        this->data_ + this->strides_[0] * i,
        this->sizes_ + 1,
        this->strides_ + 1,
        N - 1);
  }

 private:
  TensorAccessor(
      T* data,
      const executorch::aten::SizesType* sizes,
      const executorch::aten::StridesType* strides,
      ssize_t dim)
      : internal::TensorAccessorBase<T, N>(data, sizes, strides, dim) {}

  template <typename T2, ssize_t N2>
  friend class TensorAccessor;

  template <typename T2, ssize_t N2>
  friend executorch::runtime::Result<TensorAccessor<T2, N2>>
  make_tensor_accessor(const executorch::aten::Tensor& t);
};

/**
 * TensorAccessor specialization for N == 1, where operator[] returns a
 * reference to the underlying scalar.
 */
template <typename T>
class TensorAccessor<T, 1> : public internal::TensorAccessorBase<T, 1> {
 public:
  /**
   * Index into the the outer most dimension.
   *
   * @param i Index.
   * @return Reference to the underlying scalar.
   */
  T& operator[](ssize_t i) {
    return this->data_[this->strides_[0] * i];
  }

  /**
   * Index into the the outer most dimension.
   *
   * @param i Index.
   * @return Constant reference to the underlying scalar.
   */
  const T& operator[](ssize_t i) const {
    return this->data_[this->strides_[0] * i];
  }

 private:
  TensorAccessor(
      T* data,
      const executorch::aten::SizesType* sizes,
      const executorch::aten::StridesType* strides,
      ssize_t dim)
      : internal::TensorAccessorBase<T, 1>(data, sizes, strides, dim) {}

  template <typename T2, ssize_t N2>
  friend class TensorAccessor;

  template <typename T2, ssize_t N2>
  friend executorch::runtime::Result<TensorAccessor<T2, N2>>
  make_tensor_accessor(const executorch::aten::Tensor& t);
};

/**
 * Creates a TensorAccessor<T, N> from the given tensor. The number of dimension
 * N and the data type T's size must match those of the input tensor. For
 * Executorch tensors, non-trivial dimension order is not supported.
 *
 * @param tensor Origin tensor. The TensorImpl inside must outlive the returned
 * TensorAccessor.
 * @return TensorAccessor of the input tensor.
 * @retval Error::InvalidArgument Mismatch on data type or number of dimensions.
 * @retval Error::NotSupported Input tensor has non-trivial dimension onrder.
 */
template <typename T, ssize_t N>
executorch::runtime::Result<TensorAccessor<T, N>> make_tensor_accessor(
    const executorch::aten::Tensor& tensor) {
  static_assert(
      N > 0,
      "TensorAccessor is used for indexing tensors, for scalar use *_data_ptr<T>()");

  if (N != tensor.dim()) {
    ET_LOG(
        Error,
        "Expecting %zd dimensions but tensor has %zd.",
        static_cast<ssize_t>(N),
        static_cast<ssize_t>(tensor.dim()));
    return executorch::runtime::Error::InvalidArgument;
  }

  if (sizeof(T) != tensor.element_size()) {
    ET_LOG(
        Error,
        "Size of data type template argument (%zd) not equal to tensor element size (%zd)",
        static_cast<ssize_t>(sizeof(T)),
        static_cast<ssize_t>(tensor.element_size()));
    return executorch::runtime::Error::InvalidArgument;
  }

#ifndef USE_ATEN_LIB
  auto dim_order = tensor.dim_order();
  for (ssize_t i = 0; i < dim_order.size(); i++) {
    if (dim_order[i] != i) {
      ET_LOG(Error, "Non-trival dim_order not supported.");
      return executorch::runtime::Error::NotSupported;
    }
  }
#endif

  T* ptr = nullptr;
  if constexpr (std::is_const_v<T>) {
    ptr = tensor.const_data_ptr<T>();
  } else {
    ptr = tensor.mutable_data_ptr<T>();
  }
  return TensorAccessor<T, N>(
      ptr, tensor.sizes().data(), tensor.strides().data(), N);
}

} // namespace extension
} // namespace executorch