/*
    tests/eigen_tensor.cpp -- automatic conversion of Eigen Tensor

    All rights reserved. Use of this source code is governed by a
    BSD-style license that can be found in the LICENSE file.
*/

#include <pybind11/eigen/tensor.h>

PYBIND11_NAMESPACE_BEGIN(eigen_tensor_test)

namespace py = pybind11;

PYBIND11_WARNING_DISABLE_MSVC(4127)

template <typename M>
void reset_tensor(M &x) {
    for (int i = 0; i < x.dimension(0); i++) {
        for (int j = 0; j < x.dimension(1); j++) {
            for (int k = 0; k < x.dimension(2); k++) {
                x(i, j, k) = i * (5 * 2) + j * 2 + k;
            }
        }
    }
}

template <typename M>
bool check_tensor(M &x) {
    for (int i = 0; i < x.dimension(0); i++) {
        for (int j = 0; j < x.dimension(1); j++) {
            for (int k = 0; k < x.dimension(2); k++) {
                if (x(i, j, k) != (i * (5 * 2) + j * 2 + k)) {
                    return false;
                }
            }
        }
    }
    return true;
}

template <int Options>
Eigen::Tensor<double, 3, Options> &get_tensor() {
    static Eigen::Tensor<double, 3, Options> *x;

    if (!x) {
        x = new Eigen::Tensor<double, 3, Options>(3, 5, 2);
        reset_tensor(*x);
    }

    return *x;
}

template <int Options>
Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> &get_tensor_map() {
    static Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> *x;

    if (!x) {
        x = new Eigen::TensorMap<Eigen::Tensor<double, 3, Options>>(get_tensor<Options>());
    }

    return *x;
}

template <int Options>
Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options> &get_fixed_tensor() {
    static Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options> *x;

    if (!x) {
        Eigen::aligned_allocator<Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options>>
            allocator;
        x = new (allocator.allocate(1))
            Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options>();
        reset_tensor(*x);
    }

    return *x;
}

template <int Options>
const Eigen::Tensor<double, 3, Options> &get_const_tensor() {
    return get_tensor<Options>();
}

template <int Options>
struct CustomExample {
    CustomExample() : member(get_tensor<Options>()), view_member(member) {}

    Eigen::Tensor<double, 3, Options> member;
    Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> view_member;
};

template <int Options>
void init_tensor_module(pybind11::module &m) {
    const char *needed_options = "";
    if (Options == Eigen::ColMajor) {
        needed_options = "F";
    } else {
        needed_options = "C";
    }
    m.attr("needed_options") = needed_options;

    m.def("setup", []() {
        reset_tensor(get_tensor<Options>());
        reset_tensor(get_fixed_tensor<Options>());
    });

    m.def("is_ok", []() {
        return check_tensor(get_tensor<Options>()) && check_tensor(get_fixed_tensor<Options>());
    });

    py::class_<CustomExample<Options>>(m, "CustomExample", py::module_local())
        .def(py::init<>())
        .def_readonly(
            "member", &CustomExample<Options>::member, py::return_value_policy::reference_internal)
        .def_readonly("member_view",
                      &CustomExample<Options>::view_member,
                      py::return_value_policy::reference_internal);

    m.def(
        "copy_fixed_tensor",
        []() { return &get_fixed_tensor<Options>(); },
        py::return_value_policy::copy);

    m.def(
        "copy_tensor", []() { return &get_tensor<Options>(); }, py::return_value_policy::copy);

    m.def(
        "copy_const_tensor",
        []() { return &get_const_tensor<Options>(); },
        py::return_value_policy::copy);

    m.def(
        "move_fixed_tensor_copy",
        []() -> Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options> {
            return get_fixed_tensor<Options>();
        },
        py::return_value_policy::move);

    m.def(
        "move_tensor_copy",
        []() -> Eigen::Tensor<double, 3, Options> { return get_tensor<Options>(); },
        py::return_value_policy::move);

    m.def(
        "move_const_tensor",
        []() -> const Eigen::Tensor<double, 3, Options> & { return get_const_tensor<Options>(); },
        py::return_value_policy::move);

    m.def(
        "take_fixed_tensor",

        []() {
            Eigen::aligned_allocator<
                Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options>>
                allocator;
            return new (allocator.allocate(1))
                Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options>(
                    get_fixed_tensor<Options>());
        },
        py::return_value_policy::take_ownership);

    m.def(
        "take_tensor",
        []() { return new Eigen::Tensor<double, 3, Options>(get_tensor<Options>()); },
        py::return_value_policy::take_ownership);

    m.def(
        "take_const_tensor",
        []() -> const Eigen::Tensor<double, 3, Options> * {
            return new Eigen::Tensor<double, 3, Options>(get_tensor<Options>());
        },
        py::return_value_policy::take_ownership);

    m.def(
        "take_view_tensor",
        []() -> const Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> * {
            return new Eigen::TensorMap<Eigen::Tensor<double, 3, Options>>(get_tensor<Options>());
        },
        py::return_value_policy::take_ownership);

    m.def(
        "reference_tensor",
        []() { return &get_tensor<Options>(); },
        py::return_value_policy::reference);

    m.def(
        "reference_tensor_v2",
        []() -> Eigen::Tensor<double, 3, Options> & { return get_tensor<Options>(); },
        py::return_value_policy::reference);

    m.def(
        "reference_tensor_internal",
        []() { return &get_tensor<Options>(); },
        py::return_value_policy::reference_internal);

    m.def(
        "reference_fixed_tensor",
        []() { return &get_tensor<Options>(); },
        py::return_value_policy::reference);

    m.def(
        "reference_const_tensor",
        []() { return &get_const_tensor<Options>(); },
        py::return_value_policy::reference);

    m.def(
        "reference_const_tensor_v2",
        []() -> const Eigen::Tensor<double, 3, Options> & { return get_const_tensor<Options>(); },
        py::return_value_policy::reference);

    m.def(
        "reference_view_of_tensor",
        []() -> Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> {
            return get_tensor_map<Options>();
        },
        py::return_value_policy::reference);

    m.def(
        "reference_view_of_tensor_v2",
        // NOLINTNEXTLINE(readability-const-return-type)
        []() -> const Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> {
            return get_tensor_map<Options>(); // NOLINT(readability-const-return-type)
        },                                    // NOLINT(readability-const-return-type)
        py::return_value_policy::reference);

    m.def(
        "reference_view_of_tensor_v3",
        []() -> Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> * {
            return &get_tensor_map<Options>();
        },
        py::return_value_policy::reference);

    m.def(
        "reference_view_of_tensor_v4",
        []() -> const Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> * {
            return &get_tensor_map<Options>();
        },
        py::return_value_policy::reference);

    m.def(
        "reference_view_of_tensor_v5",
        []() -> Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> & {
            return get_tensor_map<Options>();
        },
        py::return_value_policy::reference);

    m.def(
        "reference_view_of_tensor_v6",
        []() -> const Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> & {
            return get_tensor_map<Options>();
        },
        py::return_value_policy::reference);

    m.def(
        "reference_view_of_fixed_tensor",
        []() {
            return Eigen::TensorMap<
                Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options>>(
                get_fixed_tensor<Options>());
        },
        py::return_value_policy::reference);

    m.def("round_trip_tensor",
          [](const Eigen::Tensor<double, 3, Options> &tensor) { return tensor; });

    m.def(
        "round_trip_tensor_noconvert",
        [](const Eigen::Tensor<double, 3, Options> &tensor) { return tensor; },
        py::arg("tensor").noconvert());

    m.def("round_trip_tensor2",
          [](const Eigen::Tensor<int32_t, 3, Options> &tensor) { return tensor; });

    m.def("round_trip_fixed_tensor",
          [](const Eigen::TensorFixedSize<double, Eigen::Sizes<3, 5, 2>, Options> &tensor) {
              return tensor;
          });

    m.def(
        "round_trip_view_tensor",
        [](Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> view) { return view; },
        py::return_value_policy::reference);

    m.def(
        "round_trip_view_tensor_ref",
        [](Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> &view) { return view; },
        py::return_value_policy::reference);

    m.def(
        "round_trip_view_tensor_ptr",
        [](Eigen::TensorMap<Eigen::Tensor<double, 3, Options>> *view) { return view; },
        py::return_value_policy::reference);

    m.def(
        "round_trip_aligned_view_tensor",
        [](Eigen::TensorMap<Eigen::Tensor<double, 3, Options>, Eigen::Aligned> view) {
            return view;
        },
        py::return_value_policy::reference);

    m.def(
        "round_trip_const_view_tensor",
        [](Eigen::TensorMap<const Eigen::Tensor<double, 3, Options>> view) {
            return Eigen::Tensor<double, 3, Options>(view);
        },
        py::return_value_policy::move);

    m.def(
        "round_trip_rank_0",
        [](const Eigen::Tensor<double, 0, Options> &tensor) { return tensor; },
        py::return_value_policy::move);

    m.def(
        "round_trip_rank_0_noconvert",
        [](const Eigen::Tensor<double, 0, Options> &tensor) { return tensor; },
        py::arg("tensor").noconvert(),
        py::return_value_policy::move);

    m.def(
        "round_trip_rank_0_view",
        [](Eigen::TensorMap<Eigen::Tensor<double, 0, Options>> &tensor) { return tensor; },
        py::return_value_policy::reference);
}

void test_module(py::module_ &m) {
    auto f_style = m.def_submodule("f_style");
    auto c_style = m.def_submodule("c_style");

    init_tensor_module<Eigen::ColMajor>(f_style);
    init_tensor_module<Eigen::RowMajor>(c_style);
}

PYBIND11_NAMESPACE_END(eigen_tensor_test)