/*!
 * Copyright 2019-2021 by Contributors
 * \file array_interface.h
 * \brief View of __array_interface__
 */
#ifndef XGBOOST_DATA_ARRAY_INTERFACE_H_
#define XGBOOST_DATA_ARRAY_INTERFACE_H_

#include <algorithm>
#include <cinttypes>
#include <map>
#include <string>
#include <utility>
#include <vector>

#include "../common/bitfield.h"
#include "../common/common.h"
#include "xgboost/base.h"
#include "xgboost/data.h"
#include "xgboost/json.h"
#include "xgboost/linalg.h"
#include "xgboost/logging.h"
#include "xgboost/span.h"

namespace xgboost {
// Common errors in parsing columnar format.
struct ArrayInterfaceErrors {
  static char const *Contiguous() { return "Memory should be contiguous."; }
  static char const *TypestrFormat() {
    return "`typestr' should be of format <endian><type><size of type in bytes>.";
  }
  static char const *Dimension(int32_t d) {
    static std::string str;
    str.clear();
    str += "Only ";
    str += std::to_string(d);
    str += " dimensional array is valid.";
    return str.c_str();
  }
  static char const *Version() {
    return "Only version <= 3 of `__cuda_array_interface__' and `__array_interface__' are "
           "supported.";
  }
  static char const *OfType(std::string const &type) {
    static std::string str;
    str.clear();
    str += " should be of ";
    str += type;
    str += " type.";
    return str.c_str();
  }

  static std::string TypeStr(char c) {
    switch (c) {
      case 't':
        return "Bit field";
      case 'b':
        return "Boolean";
      case 'i':
        return "Integer";
      case 'u':
        return "Unsigned integer";
      case 'f':
        return "Floating point";
      case 'c':
        return "Complex floating point";
      case 'm':
        return "Timedelta";
      case 'M':
        return "Datetime";
      case 'O':
        return "Object";
      case 'S':
        return "String";
      case 'U':
        return "Unicode";
      case 'V':
        return "Other";
      default:
        LOG(FATAL) << "Invalid type code: " << c << " in `typestr' of input array."
                   << "\nPlease verify the `__cuda_array_interface__/__array_interface__' "
                   << "of your input data complies to: "
                   << "https://docs.scipy.org/doc/numpy/reference/arrays.interface.html"
                   << "\nOr open an issue.";
        return "";
    }
  }

  static std::string UnSupportedType(StringView typestr) {
    return TypeStr(typestr[1]) + "-" + typestr[2] + " is not supported.";
  }
};

/**
 * Utilities for consuming array interface.
 */
class ArrayInterfaceHandler {
 public:
  enum Type : std::int8_t { kF4, kF8, kF16, kI1, kI2, kI4, kI8, kU1, kU2, kU4, kU8 };

  template <typename PtrType>
  static PtrType GetPtrFromArrayData(Object::Map const &obj) {
    auto data_it = obj.find("data");
    if (data_it == obj.cend() || IsA<Null>(data_it->second)) {
      LOG(FATAL) << "Empty data passed in.";
    }
    auto p_data = reinterpret_cast<PtrType>(
        static_cast<size_t>(get<Integer const>(get<Array const>(data_it->second).at(0))));
    return p_data;
  }

  static void Validate(Object::Map const &array) {
    auto version_it = array.find("version");
    if (version_it == array.cend() || IsA<Null>(version_it->second)) {
      LOG(FATAL) << "Missing `version' field for array interface";
    }
    if (get<Integer const>(version_it->second) > 3) {
      LOG(FATAL) << ArrayInterfaceErrors::Version();
    }

    auto typestr_it = array.find("typestr");
    if (typestr_it == array.cend() || IsA<Null>(typestr_it->second)) {
      LOG(FATAL) << "Missing `typestr' field for array interface";
    }

    auto typestr = get<String const>(typestr_it->second);
    CHECK(typestr.size() == 3 || typestr.size() == 4) << ArrayInterfaceErrors::TypestrFormat();

    auto shape_it = array.find("shape");
    if (shape_it == array.cend() || IsA<Null>(shape_it->second)) {
      LOG(FATAL) << "Missing `shape' field for array interface";
    }
    auto data_it = array.find("data");
    if (data_it == array.cend() || IsA<Null>(data_it->second)) {
      LOG(FATAL) << "Missing `data' field for array interface";
    }
  }

  // Find null mask (validity mask) field
  // Mask object is also an array interface, but with different requirements.
  static size_t ExtractMask(Object::Map const &column,
                            common::Span<RBitField8::value_type> *p_out) {
    auto &s_mask = *p_out;
    auto const &mask_it = column.find("mask");
    if (mask_it != column.cend() && !IsA<Null>(mask_it->second)) {
      auto const &j_mask = get<Object const>(mask_it->second);
      Validate(j_mask);

      auto p_mask = GetPtrFromArrayData<RBitField8::value_type *>(j_mask);

      auto j_shape = get<Array const>(j_mask.at("shape"));
      CHECK_EQ(j_shape.size(), 1) << ArrayInterfaceErrors::Dimension(1);
      auto typestr = get<String const>(j_mask.at("typestr"));
      // For now this is just 1, we can support different size of interger in mask.
      int64_t const type_length = typestr.at(2) - 48;

      if (typestr.at(1) == 't') {
        CHECK_EQ(type_length, 1) << "mask with bitfield type should be of 1 byte per bitfield.";
      } else if (typestr.at(1) == 'i') {
        CHECK_EQ(type_length, 1) << "mask with integer type should be of 1 byte per integer.";
      } else {
        LOG(FATAL) << "mask must be of integer type or bit field type.";
      }
      /*
       * shape represents how many bits is in the mask. (This is a grey area, don't be
       * suprised if it suddently represents something else when supporting a new
       * implementation).  Quoting from numpy array interface:
       *
       *   The shape of this object should be "broadcastable" to the shape of the original
       *   array.
       *
       * And that's the only requirement.
       */
      size_t const n_bits = static_cast<size_t>(get<Integer>(j_shape.at(0)));
      // The size of span required to cover all bits.  Here with 8 bits bitfield, we
      // assume 1 byte alignment.
      size_t const span_size = RBitField8::ComputeStorageSize(n_bits);

      auto strides_it = j_mask.find("strides");
      if (strides_it != j_mask.cend() && !IsA<Null>(strides_it->second)) {
        auto strides = get<Array const>(strides_it->second);
        CHECK_EQ(strides.size(), 1) << ArrayInterfaceErrors::Dimension(1);
        CHECK_EQ(get<Integer>(strides.at(0)), type_length) << ArrayInterfaceErrors::Contiguous();
      }

      s_mask = {p_mask, span_size};
      return n_bits;
    }
    return 0;
  }
  /**
   * \brief Handle vector inputs.  For higher dimension, we require strictly correct shape.
   */
  template <int32_t D>
  static void HandleRowVector(std::vector<size_t> const &shape, std::vector<size_t> *p_out) {
    auto &out = *p_out;
    if (shape.size() == 2 && D == 1) {
      auto m = shape[0];
      auto n = shape[1];
      CHECK(m == 1 || n == 1);
      if (m == 1) {
        // keep the number of columns
        out[0] = out[1];
        out.resize(1);
      } else if (n == 1) {
        // keep the number of rows.
        out.resize(1);
      }
      // when both m and n are 1, above logic keeps the column.
      // when neither m nor n is 1, caller should throw an error about Dimension.
    }
  }

  template <int32_t D>
  static void ExtractShape(Object::Map const &array, size_t (&out_shape)[D]) {
    auto const &j_shape = get<Array const>(array.at("shape"));
    std::vector<size_t> shape_arr(j_shape.size(), 0);
    std::transform(j_shape.cbegin(), j_shape.cend(), shape_arr.begin(),
                   [](Json in) { return get<Integer const>(in); });
    // handle column vector vs. row vector
    HandleRowVector<D>(shape_arr, &shape_arr);
    // Copy shape.
    size_t i;
    for (i = 0; i < shape_arr.size(); ++i) {
      CHECK_LT(i, D) << ArrayInterfaceErrors::Dimension(D);
      out_shape[i] = shape_arr[i];
    }
    // Fill the remaining dimensions
    std::fill(out_shape + i, out_shape + D, 1);
  }

  /**
   * \brief Extracts the optiona `strides' field and returns whether the array is c-contiguous.
   */
  template <int32_t D>
  static bool ExtractStride(Object::Map const &array, size_t itemsize,
                            size_t (&shape)[D], size_t (&stride)[D]) {
    auto strides_it = array.find("strides");
    // No stride is provided
    if (strides_it == array.cend() || IsA<Null>(strides_it->second)) {
      // No stride is provided, we can calculate it from shape.
      linalg::detail::CalcStride(shape, stride);
      // Quote:
      //
      //   strides: Either None to indicate a C-style contiguous array or a Tuple of
      //            strides which provides the number of bytes
      return true;
    }
    // Get shape, we need to make changes to handle row vector, so some duplicated code
    // from `ExtractShape` for copying out the shape.
    auto const &j_shape = get<Array const>(array.at("shape"));
    std::vector<size_t> shape_arr(j_shape.size(), 0);
    std::transform(j_shape.cbegin(), j_shape.cend(), shape_arr.begin(),
                   [](Json in) { return get<Integer const>(in); });
    // Get stride
    auto const &j_strides = get<Array const>(strides_it->second);
    CHECK_EQ(j_strides.size(), j_shape.size()) << "stride and shape don't match.";
    std::vector<size_t> stride_arr(j_strides.size(), 0);
    std::transform(j_strides.cbegin(), j_strides.cend(), stride_arr.begin(),
                   [](Json in) { return get<Integer const>(in); });

    // Handle column vector vs. row vector
    HandleRowVector<D>(shape_arr, &stride_arr);
    size_t i;
    for (i = 0; i < stride_arr.size(); ++i) {
      // If one of the dim has shape 0 then total size is 0, stride is meaningless, but we
      // set it to 0 here just to be consistent
      CHECK_LT(i, D) << ArrayInterfaceErrors::Dimension(D);
      // We use number of items instead of number of bytes
      stride[i] = stride_arr[i] / itemsize;
    }
    std::fill(stride + i, stride + D, 1);
    // If the stride can be calculated from shape then it's contiguous.
    size_t stride_tmp[D];
    linalg::detail::CalcStride(shape, stride_tmp);
    return std::equal(stride_tmp, stride_tmp + D, stride);
  }

  static void *ExtractData(Object::Map const &array, size_t size) {
    Validate(array);
    void *p_data = ArrayInterfaceHandler::GetPtrFromArrayData<void *>(array);
    if (!p_data) {
      CHECK_EQ(size, 0) << "Empty data with non-zero shape.";
    }
    return p_data;
  }
  /**
   * \brief Whether the ptr is allocated by CUDA.
   */
  static bool IsCudaPtr(void const *ptr);
  /**
   * \brief Sync the CUDA stream.
   */
  static void SyncCudaStream(int64_t stream);
};

/**
 * Dispatch compile time type to runtime type.
 */
template <typename T, typename E = void>
struct ToDType;
// float
template <>
struct ToDType<float> {
  static constexpr ArrayInterfaceHandler::Type kType = ArrayInterfaceHandler::kF4;
};
template <>
struct ToDType<double> {
  static constexpr ArrayInterfaceHandler::Type kType = ArrayInterfaceHandler::kF8;
};
template <typename T>
struct ToDType<T,
               std::enable_if_t<std::is_same<T, long double>::value && sizeof(long double) == 16>> {
  static constexpr ArrayInterfaceHandler::Type kType = ArrayInterfaceHandler::kF16;
};
// uint
template <>
struct ToDType<uint8_t> {
  static constexpr ArrayInterfaceHandler::Type kType = ArrayInterfaceHandler::kU1;
};
template <>
struct ToDType<uint16_t> {
  static constexpr ArrayInterfaceHandler::Type kType = ArrayInterfaceHandler::kU2;
};
template <>
struct ToDType<uint32_t> {
  static constexpr ArrayInterfaceHandler::Type kType = ArrayInterfaceHandler::kU4;
};
template <>
struct ToDType<uint64_t> {
  static constexpr ArrayInterfaceHandler::Type kType = ArrayInterfaceHandler::kU8;
};
// int
template <>
struct ToDType<int8_t> {
  static constexpr ArrayInterfaceHandler::Type kType = ArrayInterfaceHandler::kI1;
};
template <>
struct ToDType<int16_t> {
  static constexpr ArrayInterfaceHandler::Type kType = ArrayInterfaceHandler::kI2;
};
template <>
struct ToDType<int32_t> {
  static constexpr ArrayInterfaceHandler::Type kType = ArrayInterfaceHandler::kI4;
};
template <>
struct ToDType<int64_t> {
  static constexpr ArrayInterfaceHandler::Type kType = ArrayInterfaceHandler::kI8;
};

#if !defined(XGBOOST_USE_CUDA)
inline void ArrayInterfaceHandler::SyncCudaStream(int64_t) { common::AssertGPUSupport(); }
inline bool ArrayInterfaceHandler::IsCudaPtr(void const *) { return false; }
#endif  // !defined(XGBOOST_USE_CUDA)

/**
 * \brief A type erased view over __array_interface__ protocol defined by numpy
 *
 *   <a href="https://numpy.org/doc/stable/reference/arrays.interface.html">numpy</a>.
 *
 * \tparam D The number of maximum dimension.

 *   User input array must have dim <= D for all non-trivial dimensions.  During
 *   construction, the ctor can automatically remove those trivial dimensions.
 *
 * \tparam allow_mask Whether masked array is accepted.
 *
 *   Currently this only supported for 1-dim vector, which is used by cuDF column
 *   (apache arrow format).  For general masked array, as the time of writting, only
 *   numpy has the proper support even though it's in the __cuda_array_interface__
 *   protocol defined by numba.
 */
template <int32_t D, bool allow_mask = (D == 1)>
class ArrayInterface {
  static_assert(D > 0, "Invalid dimension for array interface.");

  /**
   * \brief Initialize the object, by extracting shape, stride and type.
   *
   *   The function also perform some basic validation for input array.  Lastly it will
   *   also remove trivial dimensions like converting a matrix with shape (n_samples, 1)
   *   to a vector of size n_samples.  For for inputs like weights, this should be a 1
   *   dimension column vector even though user might provide a matrix.
   */
  void Initialize(Object::Map const &array) {
    ArrayInterfaceHandler::Validate(array);

    auto typestr = get<String const>(array.at("typestr"));
    this->AssignType(StringView{typestr});
    ArrayInterfaceHandler::ExtractShape(array, shape);
    size_t itemsize = typestr[2] - '0';
    is_contiguous = ArrayInterfaceHandler::ExtractStride(array, itemsize, shape, strides);
    n = linalg::detail::CalcSize(shape);

    data = ArrayInterfaceHandler::ExtractData(array, n);
    static_assert(allow_mask ? D == 1 : D >= 1, "Masked ndarray is not supported.");
    if (allow_mask) {
      common::Span<RBitField8::value_type> s_mask;
      size_t n_bits = ArrayInterfaceHandler::ExtractMask(array, &s_mask);

      valid = RBitField8(s_mask);

      if (s_mask.data()) {
        CHECK_EQ(n_bits, n) << "Shape of bit mask doesn't match data shape. "
                            << "XGBoost doesn't support internal broadcasting.";
      }
    } else {
      auto mask_it = array.find("mask");
      CHECK(mask_it == array.cend() || IsA<Null>(mask_it->second))
          << "Masked array is not yet supported.";
    }

    auto stream_it = array.find("stream");
    if (stream_it != array.cend() && !IsA<Null>(stream_it->second)) {
      int64_t stream = get<Integer const>(stream_it->second);
      ArrayInterfaceHandler::SyncCudaStream(stream);
    }
  }

 public:
  ArrayInterface() = default;
  explicit ArrayInterface(Object::Map const &array) { this->Initialize(array); }

  explicit ArrayInterface(Json const &array) {
    if (IsA<Object>(array)) {
      this->Initialize(get<Object const>(array));
      return;
    }
    if (IsA<Array>(array)) {
      CHECK_EQ(get<Array const>(array).size(), 1)
          << "Column: " << ArrayInterfaceErrors::Dimension(1);
      this->Initialize(get<Object const>(get<Array const>(array)[0]));
      return;
    }
  }

  explicit ArrayInterface(std::string const &str) : ArrayInterface{StringView{str}} {}

  explicit ArrayInterface(StringView str) : ArrayInterface<D>{Json::Load(str)} {}

  void AssignType(StringView typestr) {
    using T = ArrayInterfaceHandler::Type;
    if (typestr.size() == 4 && typestr[1] == 'f' && typestr[2] == '1' && typestr[3] == '6') {
      type = T::kF16;
      CHECK(sizeof(long double) == 16)
          << "128-bit floating point is not supported on current platform.";
    } else if (typestr[1] == 'f' && typestr[2] == '4') {
      type = T::kF4;
    } else if (typestr[1] == 'f' && typestr[2] == '8') {
      type = T::kF8;
    } else if (typestr[1] == 'i' && typestr[2] == '1') {
      type = T::kI1;
    } else if (typestr[1] == 'i' && typestr[2] == '2') {
      type = T::kI2;
    } else if (typestr[1] == 'i' && typestr[2] == '4') {
      type = T::kI4;
    } else if (typestr[1] == 'i' && typestr[2] == '8') {
      type = T::kI8;
    } else if (typestr[1] == 'u' && typestr[2] == '1') {
      type = T::kU1;
    } else if (typestr[1] == 'u' && typestr[2] == '2') {
      type = T::kU2;
    } else if (typestr[1] == 'u' && typestr[2] == '4') {
      type = T::kU4;
    } else if (typestr[1] == 'u' && typestr[2] == '8') {
      type = T::kU8;
    } else {
      LOG(FATAL) << ArrayInterfaceErrors::UnSupportedType(typestr);
      return;
    }
  }

  XGBOOST_DEVICE size_t Shape(size_t i) const { return shape[i]; }
  XGBOOST_DEVICE size_t Stride(size_t i) const { return strides[i]; }

  template <typename Fn>
  XGBOOST_HOST_DEV_INLINE decltype(auto) DispatchCall(Fn func) const {
    using T = ArrayInterfaceHandler::Type;
    switch (type) {
      case T::kF4:
        return func(reinterpret_cast<float const *>(data));
      case T::kF8:
        return func(reinterpret_cast<double const *>(data));
#ifdef __CUDA_ARCH__
      case T::kF16: {
        // CUDA device code doesn't support long double.
        SPAN_CHECK(false);
        return func(reinterpret_cast<double const *>(data));
      }
#else
      case T::kF16:
        return func(reinterpret_cast<long double const *>(data));
#endif
      case T::kI1:
        return func(reinterpret_cast<int8_t const *>(data));
      case T::kI2:
        return func(reinterpret_cast<int16_t const *>(data));
      case T::kI4:
        return func(reinterpret_cast<int32_t const *>(data));
      case T::kI8:
        return func(reinterpret_cast<int64_t const *>(data));
      case T::kU1:
        return func(reinterpret_cast<uint8_t const *>(data));
      case T::kU2:
        return func(reinterpret_cast<uint16_t const *>(data));
      case T::kU4:
        return func(reinterpret_cast<uint32_t const *>(data));
      case T::kU8:
        return func(reinterpret_cast<uint64_t const *>(data));
    }
    SPAN_CHECK(false);
    return func(reinterpret_cast<uint64_t const *>(data));
  }

  XGBOOST_DEVICE size_t ElementSize() {
    return this->DispatchCall(
        [](auto *p_values) { return sizeof(std::remove_pointer_t<decltype(p_values)>); });
  }

  template <typename T = float, typename... Index>
  XGBOOST_DEVICE T operator()(Index &&...index) const {
    static_assert(sizeof...(index) <= D, "Invalid index.");
    return this->DispatchCall([=](auto const *p_values) -> T {
      size_t offset = linalg::detail::Offset<0ul>(strides, 0ul, index...);
      return static_cast<T>(p_values[offset]);
    });
  }

  // Used only by columnar format.
  RBitField8 valid;
  // Array stride
  size_t strides[D]{0};
  // Array shape
  size_t shape[D]{0};
  // Type earsed pointer referencing the data.
  void const *data{nullptr};
  // Total number of items
  size_t n{0};
  // Whether the memory is c-contiguous
  bool is_contiguous{false};
  // RTTI, initialized to the f16 to avoid masking potential bugs in initialization.
  ArrayInterfaceHandler::Type type{ArrayInterfaceHandler::kF16};
};

/**
 * \brief Helper for type casting.
 */
template <typename T, int32_t D>
struct TypedIndex {
  ArrayInterface<D> const &array;
  template <typename... I>
  XGBOOST_DEVICE T operator()(I &&...ind) const {
    static_assert(sizeof...(ind) <= D, "Invalid index.");
    return array.template operator()<T>(ind...);
  }
};

template <int32_t D>
inline void CheckArrayInterface(StringView key, ArrayInterface<D> const &array) {
  CHECK(!array.valid.Data()) << "Meta info " << key << " should be dense, found validity mask";
}
}  // namespace xgboost
#endif  // XGBOOST_DATA_ARRAY_INTERFACE_H_