#include "chainerx/python/array.h"

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include <absl/types/optional.h>
#include <pybind11/operators.h>

#include "chainerx/array.h"
#include "chainerx/array_body.h"
#include "chainerx/array_index.h"
#include "chainerx/axes.h"
#include "chainerx/backend_util.h"
#include "chainerx/backward.h"
#include "chainerx/constant.h"
#include "chainerx/context.h"
#include "chainerx/device.h"
#include "chainerx/dtype.h"
#include "chainerx/error.h"
#include "chainerx/graph.h"
#include "chainerx/indexable_array.h"
#include "chainerx/indexer.h"
#include "chainerx/native/data_type.h"
#include "chainerx/native/native_backend.h"
#include "chainerx/routines/arithmetic.h"
#include "chainerx/routines/creation.h"
#include "chainerx/routines/indexing.h"
#include "chainerx/routines/manipulation.h"
#include "chainerx/routines/misc.h"
#include "chainerx/routines/sorting.h"
#include "chainerx/shape.h"
#include "chainerx/slice.h"
#include "chainerx/strides.h"

#include "chainerx/python/array_index.h"
#include "chainerx/python/axes.h"
#include "chainerx/python/common.h"
#include "chainerx/python/device.h"
#include "chainerx/python/dtype.h"
#include "chainerx/python/py_cached_objects.h"
#include "chainerx/python/shape.h"
#include "chainerx/python/strides.h"

namespace chainerx {
namespace python {
namespace python_internal {
namespace {

using internal::GetArrayBody;
using internal::MoveArrayBody;

}  // namespace

namespace py = pybind11;
using py::literals::operator""_a;

py::tuple ToTuple(const std::vector<Array>& ary) {
    py::tuple ret{ary.size()};
    for (size_t i = 0; i < ary.size(); i++) {
        ret[i] = GetArrayBody(ary[i]);
    }
    return ret;
}

std::vector<ArrayBodyPtr> ToArrayBodyPtr(const std::vector<Array>& ary) {
    std::vector<ArrayBodyPtr> ret{ary.size()};
    for (size_t i = 0; i < ary.size(); i++) {
        ArrayBodyPtr array_body = GetArrayBody(ary[i]);
        ret[i] = std::move(array_body);
    }
    return ret;
}

ArrayBodyPtr MakeArrayFromNumpyArray(py::array array, Device& device) {
    Shape shape{array.shape(), array.shape() + array.ndim()};
    Dtype dtype = GetDtypeFromNumpyDtype(array.dtype());
    Strides strides{array.strides(), array.strides() + array.ndim()};

    int64_t first{};
    int64_t last{};
    std::tie(first, last) = GetDataRange(shape, strides, array.itemsize());
    py::buffer_info info = array.request();

    // Some backends may perform zero copy, so increment refcount of numpy ndarray not to be released in user codes.
    // Note that inc_ref() / dec_ref() is performed by the lambda capture.
    std::shared_ptr<void> data{static_cast<char*>(info.ptr) + first, [array](void*) {}};

    return MoveArrayBody(internal::FromHostData(shape, dtype, data, strides, -first, device));
}

namespace {

py::array MakeNumpyArrayFromArray(const py::module& m, const ArrayBodyPtr& self, bool copy) {
    Array array = Array{self}.ToNative();

    py::object dtype = GetNumpyDtypeFromModule(m, array.dtype());
    const Shape& shape = array.shape();
    const Strides& strides = array.strides();
    const void* ptr = internal::GetRawOffsetData(array);

    if (copy) {
        return py::array{dtype, shape, strides, ptr};
    }
    return py::array{dtype, shape, strides, ptr, py::cast(internal::MoveArrayBody(std::move(array)))};
}

// TODO(okapies): this is a workaround for improving performance
py::object MakeCupyArrayFromArray(const py::module& m, py::handle self) {
    Array array{py::cast<ArrayBodyPtr>(self)};
    Device& device = array.device();
    // TODO(okapies): rejects if array's device is not compatible with cupy

    py::object dtype = GetNumpyDtypeFromModule(m, array.dtype());
    const Shape& shape = array.shape();
    const Strides& strides = array.strides();

    // NOLINTNEXTLINE(cppcoreguidelines-pro-type-reinterpret-cast)
    const intptr_t ptr = reinterpret_cast<intptr_t>(array.raw_data());
    const auto range = GetDataRange(shape, strides, array.GetItemSize());
    const auto data_size = std::get<1>(range) - std::get<0>(range);
    const auto device_index = device.index();

    // Convert object to CuPy array using cupy.ndarray()
    auto memory_pointer = GetCachedCupyMemoryPointer();
    auto unowned_memory = GetCachedCupyUnownedMemory();
    py::object memptr = memory_pointer(unowned_memory(ptr, data_size, self, device_index), array.offset());

    auto ndarray = GetCachedCupyNdarray();
    return ndarray(ToTuple(shape), dtype, memptr, ToTuple(strides));
}

}  // namespace

ArrayBodyPtr MakeArray(py::handle object, py::handle dtype, bool copy, py::handle device) {
    absl::optional<Dtype> dtype_ = dtype.is_none() ? absl::nullopt : absl::optional<Dtype>(GetDtype(dtype));
    Device& dev = GetDevice(device);

    return MakeArray(object, dtype_, copy, dev);
}

ArrayBodyPtr MakeArray(py::handle object, const absl::optional<Dtype>& dtype, bool copy, Device& device) {
    // object is chainerx.ndarray
    if (py::isinstance<ArrayBody>(object)) {
        Array a = Array{py::cast<ArrayBodyPtr>(object)};
        Dtype dtype_ = dtype.has_value() ? *dtype : a.dtype();

        if (!copy && a.dtype() == dtype_ && &a.device() == &device) {
            return MoveArrayBody(std::move(a));
        }
        // Note that the graph is connected.
        if (&a.device() != &device) {
            return MoveArrayBody(a.ToDevice(device).AsType(dtype_, false));
        }
        if (a.dtype() != dtype_) {
            return MoveArrayBody(a.AsType(dtype_, true));
        }
        return MoveArrayBody(a.Copy());
    }

    // Convert object to NumPy array using numpy.array()
    // TODO(sonots): Remove dependency on numpy
    auto array_func = GetCachedNumpyArray();
    py::object dtype_name = py::none();
    if (dtype.has_value()) {
        dtype_name = py::str{GetDtypeName(*dtype)};
    }
    py::array np_array = array_func(object, "copy"_a = copy, "dtype"_a = dtype_name);

    // Convert NumPy array to ChainerX array
    return MakeArrayFromNumpyArray(np_array, device);
}

void InitChainerxArrayConversion(pybind11::module& m, py::class_<ArrayBody, ArrayBodyPtr>& c) {
    // TODO(hvy): Support all arguments in the constructor of numpy.ndarray.
    c.def(py::init([](py::handle shape, py::handle dtype, py::handle device) {
              return MoveArrayBody(Empty(ToShape(shape), GetDtype(dtype), GetDevice(device)));
          }),
          "shape"_a,
          "dtype"_a,
          "device"_a = nullptr);
    c.def_property_readonly("__array_priority__", [](const ArrayBodyPtr & /*self*/) -> double { return 100.; });
    m.def("to_numpy",
          [m](const ArrayBodyPtr& array, bool copy) { return MakeNumpyArrayFromArray(m, array, copy); },
          "array"_a,
          "copy"_a = true);
    m.def("_to_cupy", [m](py::handle array) { return MakeCupyArrayFromArray(m, array); }, "array"_a);
    // This is currently for internal use (from Chainer) to support CuPy.
    // TODO(niboshi): Remove this once it will be possible to import cupy.ndarray using chx.array / chx.asarray.
    m.def("_fromrawpointer",
          [](intptr_t ptr, py::handle shape, py::handle dtype, const py::tuple& strides, py::handle device, int64_t offset, py::object base)
                  -> ArrayBodyPtr {
              // TODO(niboshi): Expose `base` as `ndarray.base` attribute.
              void* c_ptr = reinterpret_cast<void*>(ptr);  // NOLINT(cppcoreguidelines-pro-type-reinterpret-cast)
              // Note that inc_ref() / dec_ref() is performed by the lambda capture.
              std::shared_ptr<void> data{c_ptr, [base](void*) {}};
              return MoveArrayBody(FromData(ToShape(shape), GetDtype(dtype), data, ToStrides(strides), offset, GetDevice(device)));
          });
    c.def(py::pickle(
            [m](const ArrayBodyPtr& self) -> py::tuple {
                return py::make_tuple(MakeNumpyArrayFromArray(m, self, true), py::cast(self->device(), py::return_value_policy::reference));
            },
            [](py::tuple state) -> ArrayBodyPtr {
                py::array numpy_array = state[0];
                Device& device = py::cast<Device&>(state[1]);
                return MakeArrayFromNumpyArray(numpy_array, device);
            }));
    // TODO(niboshi): Support arguments
    c.def("item", [](const ArrayBodyPtr& a) -> py::object {
        Scalar s = AsScalar(Array{a});
        switch (s.kind()) {
            case DtypeKind::kBool:
                return py::bool_{static_cast<bool>(s)};
            case DtypeKind::kInt:
                return py::int_{static_cast<int64_t>(s)};
            case DtypeKind::kFloat:
                return py::float_{static_cast<double>(s)};
            default:
                CHAINERX_NEVER_REACH();
        }
    });
    c.def("view", [](const ArrayBodyPtr& self) { return MoveArrayBody(Array{self}.MakeView()); });
    c.def("astype",
          [](const ArrayBodyPtr& self, py::handle dtype, bool copy) { return MoveArrayBody(Array{self}.AsType(GetDtype(dtype), copy)); },
          "dtype"_a,
          "copy"_a = true);
    c.def("copy", [](const ArrayBodyPtr& self) { return MoveArrayBody(Array{self}.Copy()); });
    c.def("fill",
          [](const ArrayBodyPtr& self, Scalar value) {
              Array{self}.Fill(value);
              return;
          },
          "value"_a);
}

void InitChainerxArrayManipulation(py::class_<ArrayBody, ArrayBodyPtr>& c) {
    c.def("take",
          [](const ArrayBodyPtr& self, py::handle indices, const absl::optional<int8_t>& axis) {
              if (!axis.has_value()) {
                  throw NotImplementedError{"axis=None is not yet supported for chainerx.ndarray.take."};
              }
              if (py::isinstance<ArrayBody>(indices)) {
                  return MoveArrayBody(Array{self}.Take(Array{py::cast<ArrayBodyPtr>(indices)}, axis.value()));
              }
              if (py::isinstance<py::sequence>(indices)) {
                  absl::optional<Dtype> dtype = Dtype::kInt64;
                  return MoveArrayBody(Array{self}.Take(Array{MakeArray(indices, dtype, false, self->device())}, axis.value()));
              }
              if (py::isinstance<py::array>(indices)) {
                  return MoveArrayBody(
                          Array{self}.Take(Array{MakeArrayFromNumpyArray(py::cast<py::array>(indices), self->device())}, axis.value()));
              }
              throw py::type_error{"only integers, slices (`:`), sequence, numpy.ndarray and chainerx.newaxis (`None`) are valid indices"};
          },
          "indices"_a,
          "axis"_a = nullptr);
    c.def("transpose",
          [](const ArrayBodyPtr& self, const absl::optional<std::vector<int8_t>>& axes) {
              return MoveArrayBody(Array{self}.Transpose(ToAxes(axes)));
          },
          "axes"_a = nullptr);
    c.def("ravel", [](const ArrayBodyPtr& self) { return MoveArrayBody(Array{self}.Ravel()); });
    c.def("transpose", [](const ArrayBodyPtr& self, py::args args) { return MoveArrayBody(Array{self}.Transpose(ToAxes(args))); });
    c.def("reshape", [](const ArrayBodyPtr& self, py::handle shape) { return MoveArrayBody(Array{self}.Reshape(ToShape(shape))); });
    c.def("reshape", [](const ArrayBodyPtr& self, const std::vector<int64_t>& shape) {
        return MoveArrayBody(Array{self}.Reshape({shape.begin(), shape.end()}));
    });
    c.def("reshape", [](const ArrayBodyPtr& self, py::args args) {
        if (args.size() == 0) {
            throw ChainerxError{"Reshape takes exactly 1 argument (0 given)."};
        }
        return MoveArrayBody(Array{self}.Reshape(ToShape(args)));
    });
    c.def("squeeze",
          [](const ArrayBodyPtr& self, const absl::optional<std::vector<int8_t>>& axis) {
              return MoveArrayBody(Array{self}.Squeeze(ToAxes(axis)));
          },
          "axis"_a = nullptr);
    c.def("squeeze", [](const ArrayBodyPtr& self, int8_t axis) { return MoveArrayBody(Array{self}.Squeeze(Axes{axis})); }, "axis"_a);
    c.def("swapaxes",
          [](const ArrayBodyPtr& self, int8_t axis1, int8_t axis2) { return MoveArrayBody(Array{self}.Swapaxes(axis1, axis2)); },
          "axis1"_a,
          "axis2"_a);
    c.def("repeat",
          [](const ArrayBodyPtr& self, int64_t repeats, absl::optional<int8_t> axis) {
              return MoveArrayBody(Repeat(Array{self}, repeats, axis));
          },
          "repeats"_a,
          "axis"_a = nullptr);
    c.def("repeat",
          [](const ArrayBodyPtr& self, const std::vector<int64_t>& repeats, absl::optional<int8_t> axis) {
              return MoveArrayBody(Repeat(Array{self}, repeats, axis));
          },
          "repeats"_a,
          "axis"_a = nullptr);
    c.def("dot", [](const ArrayBodyPtr& self, const ArrayBodyPtr& b) { return MoveArrayBody(Array{self}.Dot(Array{b})); }, "b"_a);
    c.def("flatten", [](const ArrayBodyPtr& self) { return MoveArrayBody(Array{self}.Flatten()); });
}

void InitChainerxArrayComparison(py::class_<ArrayBody, ArrayBodyPtr>& c) {
    c.def("__eq__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} == Array{rhs}); },
          py::is_operator());
    c.def("__eq__",
          [](const ArrayBodyPtr& self, Scalar rhs) {
              // TODO(niboshi): More efficient implementation
              Array self_array{self};
              return MoveArrayBody(self_array == FullLike(self_array, rhs, self->device()));
          },
          py::is_operator());
    c.def("__ne__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} != Array{rhs}); },
          py::is_operator());
    c.def("__ne__",
          [](const ArrayBodyPtr& self, Scalar rhs) {
              // TODO(niboshi): More efficient implementation
              Array self_array{self};
              return MoveArrayBody(self_array != FullLike(self_array, rhs, self->device()));
          },
          py::is_operator());
    c.def("__gt__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} > Array{rhs}); },
          py::is_operator());
    c.def("__gt__",
          [](const ArrayBodyPtr& self, Scalar rhs) {
              // TODO(niboshi): More efficient implementation
              Array self_array{self};
              return MoveArrayBody(self_array > FullLike(self_array, rhs, self->device()));
          },
          py::is_operator());
    c.def("__ge__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} >= Array{rhs}); },
          py::is_operator());
    c.def("__ge__",
          [](const ArrayBodyPtr& self, Scalar rhs) {
              // TODO(niboshi): More efficient implementation
              Array self_array{self};
              return MoveArrayBody(self_array >= FullLike(self_array, rhs, self->device()));
          },
          py::is_operator());
    c.def("__lt__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} < Array{rhs}); },
          py::is_operator());
    c.def("__lt__",
          [](const ArrayBodyPtr& self, Scalar rhs) {
              // TODO(niboshi): More efficient implementation
              Array self_array{self};
              return MoveArrayBody(self_array < FullLike(self_array, rhs, self->device()));
          },
          py::is_operator());
    c.def("__le__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} <= Array{rhs}); },
          py::is_operator());
    c.def("__le__",
          [](const ArrayBodyPtr& self, Scalar rhs) {
              // TODO(niboshi): More efficient implementation
              Array self_array{self};
              return MoveArrayBody(self_array <= FullLike(self_array, rhs, self->device()));
          },
          py::is_operator());
}

void InitChainerxArrayUnary(py::class_<ArrayBody, ArrayBodyPtr>& c) {
    c.def("__neg__", [](const ArrayBodyPtr& self) { return MoveArrayBody(-Array{self}); });
    c.def("__abs__", [](const ArrayBodyPtr& self) { return MoveArrayBody(Absolute(Array{self})); }, py::is_operator());
}

void InitChainerxArrayInPlace(py::class_<ArrayBody, ArrayBodyPtr>& c) {
    c.def("__iadd__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(std::move(Array{self} += Array{rhs})); },
          py::is_operator());
    c.def("__iadd__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(std::move(Array{self} += rhs)); }, py::is_operator());
    c.def("__isub__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(std::move(Array{self} -= Array{rhs})); },
          py::is_operator());
    c.def("__isub__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(std::move(Array{self} -= rhs)); }, py::is_operator());
    c.def("__imul__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(std::move(Array{self} *= Array{rhs})); },
          py::is_operator());
    c.def("__imul__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(std::move(Array{self} *= rhs)); }, py::is_operator());
    c.def("__ifloordiv__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) {
              internal::IFloorDivide(Array{self}, Array{rhs});
              return self;
          },
          py::is_operator());
    c.def("__ifloordiv__", [](const ArrayBodyPtr& self, Scalar rhs) {
        internal::IFloorDivide(Array{self}, rhs);
        return self;
    });
    c.def("__itruediv__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(std::move(Array{self} /= Array{rhs})); },
          py::is_operator());
    c.def("__itruediv__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(std::move(Array{self} /= rhs)); });
    c.def("__imod__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(std::move(Array{self} %= Array{rhs})); },
          py::is_operator());
    c.def("__imod__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(std::move(Array{self} %= rhs)); });
    c.def("__iremainder__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(std::move(Array{self} %= Array{rhs})); },
          py::is_operator());
    c.def("__iremainder__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(std::move(Array{self} %= rhs)); });
    c.def("__iand__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(std::move(Array{self} &= Array{rhs})); },
          py::is_operator());
    c.def("__iand__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(std::move(Array{self} &= rhs)); }, py::is_operator());
    c.def("__ior__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(std::move(Array{self} |= Array{rhs})); },
          py::is_operator());
    c.def("__ior__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(std::move(Array{self} |= rhs)); }, py::is_operator());
    c.def("__ixor__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(std::move(Array{self} ^= Array{rhs})); },
          py::is_operator());
    c.def("__ixor__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(std::move(Array{self} ^= rhs)); }, py::is_operator());
    c.def("__ilshift__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(std::move(Array{self} <<= Array{rhs})); },
          py::is_operator());
    c.def("__ilshift__",
          [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(std::move(Array{self} <<= rhs)); },
          py::is_operator());
    c.def("__irshift__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(std::move(Array{self} >>= Array{rhs})); },
          py::is_operator());
    c.def("__irshift__",
          [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(std::move(Array{self} >>= rhs)); },
          py::is_operator());
}

void InitChainerxArrayArithmetic(py::class_<ArrayBody, ArrayBodyPtr>& c) {
    c.def("__add__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} + Array{rhs}); },
          py::is_operator());
    c.def("__add__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(Array{self} + rhs); }, py::is_operator());
    c.def("__radd__", [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(lhs + Array{self}); }, py::is_operator());
    c.def("__sub__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} - Array{rhs}); },
          py::is_operator());
    c.def("__sub__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(Array{self} - rhs); }, py::is_operator());
    c.def("__rsub__", [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(lhs - Array{self}); }, py::is_operator());
    c.def("__mul__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} * Array{rhs}); },
          py::is_operator());
    c.def("__mul__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(Array{self} * rhs); }, py::is_operator());
    c.def("__rmul__", [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(lhs * Array{self}); }, py::is_operator());
    c.def("__floordiv__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(FloorDivide(Array{self}, Array{rhs})); },
          py::is_operator());
    c.def("__floordiv__",
          [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(FloorDivide(Array{self}, rhs)); },
          py::is_operator());
    c.def("__rfloordiv__",
          [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(FloorDivide(lhs, Array{self})); },
          py::is_operator());
    c.def("__truediv__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} / Array{rhs}); },
          py::is_operator());
    c.def("__truediv__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(Array{self} / rhs); }, py::is_operator());
    c.def("__pow__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Power(Array{self}, Array{rhs})); },
          py::is_operator());
    c.def("__pow__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(Power(Array{self}, rhs)); }, py::is_operator());
    c.def("__rpow__", [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(Power(lhs, Array{self})); }, py::is_operator());
    c.def("__mod__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Mod(Array{self}, Array{rhs})); },
          py::is_operator());
    c.def("__mod__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(Mod(Array{self}, rhs)); }, py::is_operator());
    c.def("__rmod__", [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(Mod(lhs, Array{self})); }, py::is_operator());
    c.def("__remainder__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Mod(Array{self}, Array{rhs})); },
          py::is_operator());
    c.def("__remainder__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(Mod(Array{self}, rhs)); }, py::is_operator());
    c.def("__rremainder__", [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(Mod(lhs, Array{self})); }, py::is_operator());

    c.def("__rtruediv__", [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(lhs / Array{self}); }, py::is_operator());
    c.def("__and__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} & Array{rhs}); },
          py::is_operator());
    c.def("__and__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(Array{self} & rhs); }, py::is_operator());
    c.def("__rand__", [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(Array{self} & lhs); }, py::is_operator());
    c.def("__or__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} | Array{rhs}); },
          py::is_operator());
    c.def("__or__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(Array{self} | rhs); }, py::is_operator());
    c.def("__ror__", [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(Array{self} | lhs); }, py::is_operator());
    c.def("__xor__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} ^ Array{rhs}); },
          py::is_operator());
    c.def("__xor__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(Array{self} ^ rhs); }, py::is_operator());
    c.def("__rxor__", [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(Array{self} ^ lhs); }, py::is_operator());
    c.def("__lshift__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} << Array{rhs}); },
          py::is_operator());
    c.def("__lshift__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(Array{self} << rhs); }, py::is_operator());
    c.def("__rlshift__", [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(lhs << Array{self}); }, py::is_operator());
    c.def("__rshift__",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& rhs) { return MoveArrayBody(Array{self} >> Array{rhs}); },
          py::is_operator());
    c.def("__rshift__", [](const ArrayBodyPtr& self, Scalar rhs) { return MoveArrayBody(Array{self} >> rhs); }, py::is_operator());
    c.def("__rrshift__", [](const ArrayBodyPtr& self, Scalar lhs) { return MoveArrayBody(lhs >> Array{self}); }, py::is_operator());
}

void InitChainerxArrayCalculation(py::class_<ArrayBody, ArrayBodyPtr>& c) {
    c.def("sum",
          [](const ArrayBodyPtr& self, int8_t axis, bool keepdims) { return MoveArrayBody(Array{self}.Sum(Axes{axis}, keepdims)); },
          "axis"_a,
          "keepdims"_a = false);
    c.def("sum",
          [](const ArrayBodyPtr& self, const absl::optional<std::vector<int8_t>>& axis, bool keepdims) {
              return MoveArrayBody(Array{self}.Sum(ToAxes(axis), keepdims));
          },
          "axis"_a = nullptr,
          "keepdims"_a = false);
    c.def("max",
          [](const ArrayBodyPtr& self, int8_t axis, bool keepdims) { return MoveArrayBody(Array{self}.Max(Axes{axis}, keepdims)); },
          "axis"_a,
          "keepdims"_a = false);
    c.def("max",
          [](const ArrayBodyPtr& self, const absl::optional<std::vector<int8_t>>& axis, bool keepdims) {
              return MoveArrayBody(Array{self}.Max(ToAxes(axis), keepdims));
          },
          "axis"_a = nullptr,
          "keepdims"_a = false);
    c.def("min",
          [](const ArrayBodyPtr& self, int8_t axis, bool keepdims) { return MoveArrayBody(Array{self}.Min(Axes{axis}, keepdims)); },
          "axis"_a,
          "keepdims"_a = false);
    c.def("min",
          [](const ArrayBodyPtr& self, const absl::optional<std::vector<int8_t>>& axis, bool keepdims) {
              return MoveArrayBody(Array{self}.Min(ToAxes(axis), keepdims));
          },
          "axis"_a = nullptr,
          "keepdims"_a = false);
    c.def("mean",
          [](const ArrayBodyPtr& self, int8_t axis, bool keepdims) { return MoveArrayBody(Array{self}.Mean(Axes{axis}, keepdims)); },
          "axis"_a,
          "keepdims"_a = false);
    c.def("mean",
          [](const ArrayBodyPtr& self, const absl::optional<std::vector<int8_t>>& axis, bool keepdims) {
              return MoveArrayBody(Array{self}.Mean(ToAxes(axis), keepdims));
          },
          "axis"_a = nullptr,
          "keepdims"_a = false);
    c.def("var",
          [](const ArrayBodyPtr& self, int8_t axis, bool keepdims) { return MoveArrayBody(Array{self}.Var(Axes{axis}, keepdims)); },
          "axis"_a,
          "keepdims"_a = false);
    c.def("var",
          [](const ArrayBodyPtr& self, const absl::optional<std::vector<int8_t>>& axis, bool keepdims) {
              return MoveArrayBody(Array{self}.Var(ToAxes(axis), keepdims));
          },
          "axis"_a = nullptr,
          "keepdims"_a = false);
    c.def("all",
          [](const ArrayBodyPtr& self, int8_t axis, bool keepdims) { return MoveArrayBody(Array{self}.All(Axes{axis}, keepdims)); },
          "axis"_a,
          "keepdims"_a = false);
    c.def("all",
          [](const ArrayBodyPtr& self, const absl::optional<std::vector<int8_t>>& axis, bool keepdims) {
              return MoveArrayBody(Array{self}.All(ToAxes(axis), keepdims));
          },
          "axis"_a = nullptr,
          "keepdims"_a = false);
    c.def("any",
          [](const ArrayBodyPtr& self, int8_t axis, bool keepdims) { return MoveArrayBody(Array{self}.Any(Axes{axis}, keepdims)); },
          "axis"_a,
          "keepdims"_a = false);
    c.def("any",
          [](const ArrayBodyPtr& self, const absl::optional<std::vector<int8_t>>& axis, bool keepdims) {
              return MoveArrayBody(Array{self}.Any(ToAxes(axis), keepdims));
          },
          "axis"_a = nullptr,
          "keepdims"_a = false);
    c.def("argmax",
          [](const ArrayBodyPtr& self, const absl::optional<int8_t>& axis) { return MoveArrayBody(ArgMax(Array{self}, ToAxes(axis))); },
          "axis"_a = nullptr);
    c.def("argmin",
          [](const ArrayBodyPtr& self, const absl::optional<int8_t>& axis) { return MoveArrayBody(ArgMin(Array{self}, ToAxes(axis))); },
          "axis"_a = nullptr);
}

void InitChainerxArraySpecial(pybind11::module& m, py::class_<ArrayBody, ArrayBodyPtr>& c) {
    c.def("__len__", [](const ArrayBodyPtr& self) -> size_t {
        // TODO(hvy): Do bounds cheking. For reference, Chainer throws an AttributeError.
        if (self->ndim() == 0) {
            throw pybind11::type_error{"len() of unsized object"};
        }
        return self->shape().front();
    });
    c.def("__bool__", [](const ArrayBodyPtr& self) -> bool { return static_cast<bool>(AsScalar(Array{self})); });
    c.def("__int__", [](const ArrayBodyPtr& self) -> int64_t { return static_cast<int64_t>(AsScalar(Array{self})); });
    c.def("__float__", [](const ArrayBodyPtr& self) -> double { return static_cast<double>(AsScalar(Array{self})); });
    c.def("__repr__", [](const ArrayBodyPtr& self) { return Array{self}.ToString(); });
    c.def("__getitem__", [](const ArrayBodyPtr& self, py::handle key) { return MoveArrayBody(Array{self}.At(MakeArrayIndices(key))); });
    c.def("require_grad",
          [](const ArrayBodyPtr& self, const absl::optional<BackpropId>& backprop_id) {
              return MoveArrayBody(std::move(Array{self}.RequireGrad(backprop_id)));
          },
          "backprop_id"_a = nullptr);
    c.def("is_grad_required",
          [](const ArrayBodyPtr& self, const absl::optional<BackpropId>& backprop_id) { return Array{self}.IsGradRequired(backprop_id); },
          "backprop_id"_a = nullptr);
    c.def("is_backprop_required",
          [](const ArrayBodyPtr& self, const absl::optional<BackpropId>& backprop_id) {
              return Array{self}.IsBackpropRequired(backprop_id);
          },
          "backprop_id"_a = nullptr);
    c.def("is_backprop_required",
          [](const ArrayBodyPtr& self, AnyGraph any_graph) { return Array{self}.IsBackpropRequired(any_graph); },
          "backprop_id"_a);
    c.def("get_grad",
          [](const ArrayBodyPtr& self, const absl::optional<BackpropId>& backprop_id) -> ConstArrayBodyPtr {
              const absl::optional<Array>& grad = Array{self}.GetGrad(backprop_id);
              if (!grad.has_value()) {
                  return nullptr;
              }
              return internal::GetArrayBody(*grad);
          },
          "backprop_id"_a = nullptr);
    c.def("set_grad",
          [](const ArrayBodyPtr& self, const ArrayBodyPtr& grad, const absl::optional<BackpropId>& backprop_id) {
              Array array{self};
              if (grad) {
                  array.SetGrad(Array{grad}, backprop_id);
              } else {
                  array.ClearGrad(backprop_id);
              }
          },
          "grad"_a,
          "backprop_id"_a = nullptr);
    c.def("backward",
          [](const ArrayBodyPtr& self,
             const absl::optional<BackpropId>& backprop_id,
             bool enable_double_backprop,
             const absl::optional<float>& loss_scale) {
              auto double_backprop = enable_double_backprop ? DoubleBackpropOption::kEnable : DoubleBackpropOption::kDisable;
              Backward(Array{self}, backprop_id, double_backprop, loss_scale);
          },
          "backprop_id"_a = nullptr,
          "enable_double_backprop"_a = false,
          "loss_scale"_a = nullptr);
    c.def("_debug_dump_computational_graph",
          [](const ArrayBodyPtr& self, const absl::optional<BackpropId>& backprop_id) {
              DebugDumpComputationalGraph(std::cout, Array{self}, backprop_id);
          },
          "backprop_id"_a = nullptr);
    c.def_property(
            "grad",
            [](const ArrayBodyPtr& self) -> ConstArrayBodyPtr {
                const absl::optional<Array>& grad = Array{self}.GetGrad(absl::nullopt);
                if (!grad.has_value()) {
                    return nullptr;
                }
                return internal::GetArrayBody(*grad);
            },
            [](const ArrayBodyPtr& self, const ArrayBodyPtr& grad) {
                Array array{self};
                if (grad) {
                    array.SetGrad(Array{grad}, absl::nullopt);
                } else {
                    array.ClearGrad(absl::nullopt);
                }
            });
    c.def("cleargrad",
          [](const ArrayBodyPtr& self, const absl::optional<BackpropId>& backprop_id) { Array{self}.ClearGrad(backprop_id); },
          "backprop_id"_a = nullptr);
    c.def_property_readonly(
            "device", [](const ArrayBodyPtr& self) -> Device& { return self->device(); }, py::return_value_policy::reference);
    c.def_property_readonly("dtype", [m](const ArrayBodyPtr& self) { return GetNumpyDtypeFromModule(m, self->dtype()); });
    c.def_property_readonly("itemsize", [](const ArrayBodyPtr& self) { return self->GetItemSize(); });
    c.def_property_readonly("is_contiguous", [](const ArrayBodyPtr& self) { return self->IsContiguous(); });
    c.def_property_readonly("ndim", [](const ArrayBodyPtr& self) { return self->ndim(); });
    c.def_property_readonly("offset", [](const ArrayBodyPtr& self) { return self->offset(); });
    c.def_property_readonly("shape", [](const ArrayBodyPtr& self) { return ToTuple(self->shape()); });
    c.def_property_readonly("strides", [](const ArrayBodyPtr& self) { return ToTuple(self->strides()); });
    c.def_property_readonly("nbytes", [](const ArrayBodyPtr& self) { return self->GetNBytes(); });
    c.def_property_readonly("size", [](const ArrayBodyPtr& self) { return self->GetTotalSize(); });
    c.def_property_readonly("T", [](const ArrayBodyPtr& self) { return MoveArrayBody(Array{self}.Transpose()); });
    // Returns the data address, before adding offset.
    // TODO(niboshi): Consider what to do with the backends in which the "pointer" is not available from host.
    c.def_property_readonly("data_ptr", [](const ArrayBodyPtr& self) -> intptr_t {
        return reinterpret_cast<intptr_t>(self->data().get());  // NOLINT(cppcoreguidelines-pro-type-reinterpret-cast)
    });
    c.def_property_readonly("data_size", [](const ArrayBodyPtr& self) -> int64_t {
        auto range = GetDataRange(self->shape(), self->strides(), self->GetItemSize());
        return std::get<1>(range) - std::get<0>(range);
    });
    // TODO(niboshi): Remove this in favor of data_ptr.
    c.def_property_readonly(
            "_debug_data_memory_address",  // These methods starting with `_debug_` are stubs for testing
            [](const ArrayBodyPtr& self) -> intptr_t {
                const void* ptr = self->data().get();
                return reinterpret_cast<intptr_t>(ptr);  // NOLINT(cppcoreguidelines-pro-type-reinterpret-cast)
            });
    c.def_property_readonly("_debug_flat_data", [](const ArrayBodyPtr& self) {
        py::list list;
        Array array = Array{self}.ToNative();

        // Copy data into the list
        VisitDtype(array.dtype(), [&array, &list](auto pt) {
            using T = typename decltype(pt)::type;
            IndexableArray<const T> iarray{array};
            Indexer<> indexer{array.shape()};

            for (auto it = indexer.It(0); it; ++it) {
                T value = native::StorageToDataType<const T>(iarray[it]);
                if (std::is_same<T, chainerx::Float16>::value) {
                    list.append(static_cast<double>(value));
                } else {
                    list.append(value);
                }
            }
        });

        return list;
    });
    // TODO(hvy): Rename `_is_chained` to a less ambiguous function name.
    c.def("_is_chained",
          [](const ArrayBodyPtr& self, const absl::optional<BackpropId>& backprop_id) {
              BackpropId actual_backprop_id = internal::GetArrayBackpropId(Array{self}, backprop_id);
              actual_backprop_id.CheckValid();
              if (!self->HasArrayNode(actual_backprop_id)) {
                  throw ChainerxError{"Array is constant with respect to the computation for backprop ID: '", actual_backprop_id, "'."};
              }
              return self->GetArrayNode(actual_backprop_id)->creator_op_node() != nullptr;
          },
          "backprop_id"_a = nullptr);
}

void InitChainerxArray(pybind11::module& m) {
    py::class_<ArrayBody, ArrayBodyPtr> c{m, "ndarray", py::buffer_protocol()};

    c.def("to_device", [](const ArrayBodyPtr& self, py::handle device) { return MoveArrayBody(Array{self}.ToDevice(GetDevice(device))); });
    c.def("to_device", [](const ArrayBodyPtr& self, const std::string& backend_name, int index) {
        Device& device = GetDefaultContext().GetDevice({backend_name, index});
        return MoveArrayBody(Array{self}.ToDevice(device));
    });
    c.def("as_grad_stopped",
          [](const ArrayBodyPtr& self, bool copy) {
              return MoveArrayBody(Array{self}.AsGradStopped(copy ? CopyKind::kCopy : CopyKind::kView));
          },
          "copy"_a = false);
    c.def("as_grad_stopped",
          [](const ArrayBodyPtr& self, const std::vector<BackpropId>& backprop_ids, bool copy) {
              return MoveArrayBody(Array{self}.AsGradStopped(backprop_ids, copy ? CopyKind::kCopy : CopyKind::kView));
          },
          py::arg().noconvert(),
          "copy"_a = false);

    InitChainerxArrayConversion(m, c);
    InitChainerxArrayComparison(c);
    InitChainerxArrayManipulation(c);
    InitChainerxArrayUnary(c);
    InitChainerxArrayInPlace(c);
    InitChainerxArrayArithmetic(c);
    InitChainerxArrayCalculation(c);
    InitChainerxArraySpecial(m, c);
}

}  // namespace python_internal
}  // namespace python
}  // namespace chainerx
