#include "chainerx/routines/indexing.h"

#include <algorithm>
#include <cstdint>
#include <iterator>
#include <string>
#include <vector>

#include "nonstd/optional.hpp"

#include "chainerx/array.h"
#include "chainerx/array_index.h"
#include "chainerx/axes.h"
#include "chainerx/backprop_mode.h"
#include "chainerx/backward_builder.h"
#include "chainerx/backward_context.h"
#include "chainerx/constant.h"
#include "chainerx/dtype.h"
#include "chainerx/graph.h"
#include "chainerx/macro.h"
#include "chainerx/routines/creation.h"
#include "chainerx/shape.h"
#include "chainerx/slice.h"
#include "chainerx/strides.h"

namespace chainerx {
namespace internal {
namespace {

// Returns an array where elements at indices are added by the addends `b`.
//
// It is not in-place  operation: the input arrays are not altered.
// It is differentiable with respect to `a` and `b`.
Array AddAt(const Array& a, const std::vector<ArrayIndex>& indices, const Array& b) {
    // TODO(sonots): dtype conversion
    CheckEqual(a.dtype(), b.dtype());

    Array out = a.AsGradStopped(CopyKind::kCopy);
    Array out_view = out.At(indices);

    // TODO(sonots): broadcasting
    CheckEqual(out_view.shape(), b.shape());

    {
        NoBackpropModeScope scope{};
        a.device().Add(b, out_view, out_view);
    }

    {
        BackwardBuilder bb{"add_at", {a, b}, out};
        if (BackwardBuilder::Target bt = bb.CreateTarget(0)) {
            bt.Define([](BackwardContext& bctx) { bctx.input_grad() = *bctx.output_grad(); });
        }
        if (BackwardBuilder::Target bt = bb.CreateTarget(1)) {
            bt.Define([indices](BackwardContext& bctx) { bctx.input_grad() = bctx.output_grad()->At(indices); });
        }
        bb.Finalize();
    }

    return out;
}

}  // namespace

Array At(const Array& a, const std::vector<ArrayIndex>& indices) {
    Shape out_shape{};
    Strides out_strides{};
    int64_t out_offset = a.offset();
    int64_t i_in = 0;
    for (const ArrayIndex& index : indices) {
        switch (index.tag()) {
            case ArrayIndexTag::kSingleElement: {
                int64_t dim = a.shape()[i_in];
                if (index.index() < -dim || dim <= index.index()) {
                    throw DimensionError{"Index ", index.index(), " is out of bounds for axis ", i_in, " with size ", dim};
                }
                out_offset += a.strides()[i_in] * ((index.index() + dim) % dim);
                ++i_in;
                break;
            }
            case ArrayIndexTag::kSlice: {
                const Slice& slice = index.slice();
                int64_t slice_length = slice.GetLength(a.shape()[i_in]);
                out_offset += a.strides()[i_in] * slice.GetStart(a.shape()[i_in]);
                out_shape.emplace_back(slice_length);
                out_strides.emplace_back(a.strides()[i_in] * slice.step());
                ++i_in;
                break;
            }
            case ArrayIndexTag::kNewAxis:
                out_shape.emplace_back(1);
                out_strides.emplace_back(0);
                break;
            default:
                CHAINERX_NEVER_REACH();
        }
    }
    for (int64_t i = i_in; i < a.ndim(); ++i) {
        out_shape.emplace_back(a.shape()[i]);
        out_strides.emplace_back(a.strides()[i]);
    }

    Array out = MakeArray(out_shape, out_strides, a.dtype(), a.device(), a.data(), out_offset);

    BackwardBuilder bb{"get_item", a, out};
    if (BackwardBuilder::Target bt = bb.CreateTarget(0)) {
        bt.Define([indices, a_shape = a.shape(), a_dtype = a.dtype()](BackwardContext& bctx) {
            const Array& gout = *bctx.output_grad();
            Array gin = Zeros(a_shape, a_dtype, gout.device());
            bctx.input_grad() = AddAt(gin, indices, gout);
        });
    }
    bb.Finalize();

    return out;
}

}  // namespace internal

namespace {

// Adds elements of `b` indexed by `indices` into `a` and returns the result.
// Used in backward pass of Take()
//
// It is not in-place operation: the input arrays are not altered.
// It is differentiable with respect to `a` and `b`.
Array AddAt(const Array& a, const Array& indices, int8_t axis, const Array& b) {
    CHAINERX_ASSERT(0 <= axis && axis < a.ndim());
    CHAINERX_ASSERT(b.ndim() == indices.ndim() + a.ndim() - 1);
    CheckEqual(a.dtype(), b.dtype());

    CHAINERX_ASSERT(internal::GetArrayBody(indices)->nodes().empty());

    Array out = EmptyLike(a, a.device());

    {
        NoBackpropModeScope scope{};
        a.device().AddAt(a, indices, axis, b, out);
    }

    {
        BackwardBuilder bb{"add_at", {a, b}, out};
        if (BackwardBuilder::Target bt = bb.CreateTarget(0)) {
            bt.Define([](BackwardContext& bctx) { bctx.input_grad() = *bctx.output_grad(); });
        }
        if (BackwardBuilder::Target bt = bb.CreateTarget(1)) {
            CHAINERX_ASSERT(internal::GetArrayBody(indices)->nodes().empty());
            bt.Define([indices, axis](BackwardContext& bctx) { bctx.input_grad() = Take(*bctx.output_grad(), indices, axis); });
        }
        bb.Finalize();
    }

    return out;
}

}  // namespace

Array Take(const Array& a, const Array& indices, int8_t axis) {
    // TODO(niboshi): Support other dtypes by casting
    if (indices.dtype() != Dtype::kInt64) {
        throw DtypeError(
                std::string{"Only "} + GetDtypeName(Dtype::kInt64) + " is supported as indices, but given " +
                GetDtypeName(indices.dtype()));
    }

    CHAINERX_ASSERT(internal::GetArrayBody(indices)->nodes().empty());

    int8_t axis_norm = internal::NormalizeAxis(axis, a.ndim());

    Shape out_shape{};
    std::copy(a.shape().begin(), a.shape().begin() + axis_norm, std::back_inserter(out_shape));
    std::copy(indices.shape().begin(), indices.shape().end(), std::back_inserter(out_shape));
    std::copy(a.shape().begin() + (axis_norm + 1), a.shape().end(), std::back_inserter(out_shape));
    Array out = Empty(out_shape, a.dtype(), a.device());

    {
        NoBackpropModeScope scope{};
        a.device().Take(a, indices, axis_norm, out);
    }

    BackwardBuilder bb{"take", a, out};
    if (BackwardBuilder::Target bt = bb.CreateTarget(0)) {
        CHAINERX_ASSERT(internal::GetArrayBody(indices)->nodes().empty());
        bt.Define([indices, axis_norm, a_shape = a.shape()](BackwardContext& bctx) {
            const Array& gout = *bctx.output_grad();
            bctx.input_grad() = AddAt(Zeros(a_shape, gout.dtype(), gout.device()), indices, axis_norm, gout);
        });
    }
    bb.Finalize();

    return out;
}

}  // namespace chainerx
