#include "idyntree_core.h"

#include <iDynTree/Axis.h>
#include <iDynTree/Direction.h>
#include <iDynTree/MatrixDynSize.h>
#include <iDynTree/MatrixFixSize.h>
#include <iDynTree/Position.h>
#include <iDynTree/Rotation.h>
#include <iDynTree/RotationalInertia.h>
#include <iDynTree/SpatialInertia.h>
#include <iDynTree/Transform.h>
#include <iDynTree/Twist.h>
#include <iDynTree/VectorDynSize.h>
#include <iDynTree/VectorFixSize.h>
#include <pybind11/operators.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <array>
#include <cstddef>

namespace iDynTree {
namespace bindings {
namespace {
namespace py = ::pybind11;

void vectorDynSizeClassDefinition(py::class_<VectorDynSize>& vector) {
  vector.def(py::init())
      .def(py::init<unsigned>())
      .def("__getitem__", py::overload_cast<const std::size_t>(
                              &VectorDynSize::operator(), py::const_))
      .def("__setitem__",
           [](VectorDynSize& the_vector, std::size_t index, double new_value) {
             the_vector(index) = new_value;
           })
      .def(
          "__iter__",
          [](const VectorDynSize& s) {
            return py::make_iterator(s.begin(), s.end());
          },
          py::keep_alive<
              0, 1>() /* Essential: keep object alive while iterator exists */)
      .def("__len__", &VectorDynSize::size)
      .def("set_zero", &VectorDynSize::zero)
      .def("resize", &VectorDynSize::resize)
      .def("__repr__", &VectorDynSize::toString)
      .def_buffer([](VectorDynSize& v) -> py::buffer_info {
        return py::buffer_info(
            v.data(),                                /* Pointer to buffer */
            sizeof(double),                          /* Size of one scalar */
            py::format_descriptor<double>::format(), /* Python struct-style
                                                        format descriptor */
            1,                                       /* Number of dimensions */
            {v.size()},                              /* Buffer dimensions */
            {sizeof(double)} /* Strides (in bytes) for each index */
        );
      });
}

template <unsigned size>
void createFixSizeVector(pybind11::module& module,
                         const std::string& class_name) {
  py::class_<VectorFixSize<size>>(module, class_name.c_str(),
                                  py::buffer_protocol())
      .def(py::init())
      .def("__getitem__", py::overload_cast<const std::size_t>(
                              &VectorFixSize<size>::operator(), py::const_))
      .def("__setitem__",
           [](VectorFixSize<size>& the_vector, std::size_t index,
              double new_value) { the_vector(index) = new_value; })
      .def("__len__", &VectorFixSize<size>::size)
      .def(
          "__iter__",
          [](const VectorFixSize<size>& s) {
            return py::make_iterator(s.begin(), s.end());
          },
          py::keep_alive<
              0, 1>() /* Essential: keep object alive while iterator exists */)
      .def("set_zero", &VectorFixSize<size>::zero)
      .def("__repr__", &VectorFixSize<size>::toString)
      .def_buffer([](VectorFixSize<size>& v) -> py::buffer_info {
        return py::buffer_info(
            v.data(),                                /* Pointer to buffer */
            sizeof(double),                          /* Size of one scalar */
            py::format_descriptor<double>::format(), /* Python struct-style
                                                        format descriptor */
            1,                                       /* Number of dimensions */
            {v.size()},                              /* Buffer dimensions */
            {sizeof(double)} /* Strides (in bytes) for each index */
        );
      });
}

void matrixDynSizeClassDefinition(py::class_<MatrixDynSize>& matrix) {
  matrix.def(py::init())
      .def(py::init<unsigned, unsigned>())
      .def("__getitem__",
           [](MatrixDynSize& matrix,
              std::pair<std::size_t, std::size_t> indices) {
             return matrix.getVal(indices.first, indices.second);
           })
      .def("__setitem__",
           [](MatrixDynSize& matrix,
              std::pair<std::size_t, std::size_t> indices, double item) {
             return matrix.setVal(indices.first, indices.second, item);
           })
      .def("rows", &MatrixDynSize::rows)
      .def("cols", &MatrixDynSize::cols)
      .def("set_zero", &MatrixDynSize::zero)
      .def("resize", &MatrixDynSize::resize)
      .def("__repr__", &MatrixDynSize::toString)
      .def_buffer([](MatrixDynSize& m) -> py::buffer_info {
        return py::buffer_info(
            m.data(),                                /* Pointer to buffer */
            sizeof(double),                          /* Size of one scalar */
            py::format_descriptor<double>::format(), /* Python struct-style
                                                        format descriptor */
            2,                                       /* Number of dimensions */
            {m.rows(), m.cols()},                    /* Buffer dimensions */
            {sizeof(double) * m.cols(), /* Strides (in bytes) for each index */
             sizeof(double)});
      });
}

template <unsigned rows, unsigned cols>
void createFixSizeMatrix(pybind11::module& module,
                         const std::string& class_name) {
  py::class_<MatrixFixSize<rows, cols>>(module, class_name.c_str(),
                                        py::buffer_protocol())
      .def(py::init())
      .def("__getitem__",
           [](MatrixFixSize<rows, cols>& matrix,
              std::pair<std::size_t, std::size_t> indices) {
             return matrix.getVal(indices.first, indices.second);
           })
      .def("__setitem__",
           [](MatrixFixSize<rows, cols>& matrix,
              std::pair<std::size_t, std::size_t> indices, double item) {
             return matrix.setVal(indices.first, indices.second, item);
           })
      .def("rows", &MatrixFixSize<rows, cols>::rows)
      .def("cols", &MatrixFixSize<rows, cols>::cols)
      .def("set_zero", &MatrixFixSize<rows, cols>::zero)
      .def("__repr__", &MatrixFixSize<rows, cols>::toString)
      .def_buffer([](MatrixFixSize<rows, cols>& m) -> py::buffer_info {
        return py::buffer_info(
            m.data(),                                /* Pointer to buffer */
            sizeof(double),                          /* Size of one scalar */
            py::format_descriptor<double>::format(), /* Python struct-style
                                                        format descriptor */
            2,                                       /* Number of dimensions */
            {m.rows(), m.cols()},                    /* Buffer dimensions */
            {sizeof(double) * m.cols(), /* Strides (in bytes) for each index */
             sizeof(double)});
      });
}

void transformClassDefinition(py::class_<Transform>& transform) {
  transform.def(py::init())
      .def(py::init<const Rotation&, const Position&>())
      .def(py::init<const MatrixFixSize<4, 4>&>())
      .def_static("Identity", &Transform::Identity)
      .def_property_readonly("rotation", &Transform::getRotation)
      .def_property_readonly("position", &Transform::getPosition)
      .def("inverse", &Transform::inverse)
      .def(py::self * py::self)
      .def(
          "__mul__",
          [](const Transform& self, const Position& pos) -> Position {
            return self * pos;
          },
          py::is_operator())
      .def(
          "__mul__",
          [](const Transform& self, const SpatialInertia& inertia)
              -> SpatialInertia { return self * inertia; },
          py::is_operator());
}
}  // namespace

void iDynTreeCoreBindings(pybind11::module& module) {
  // Vectors and matrices.
  py::class_<VectorDynSize> vector_dyn(module, "VectorDynSize",
                                       py::buffer_protocol());
  vectorDynSizeClassDefinition(vector_dyn);
  createFixSizeVector<3>(module, "Vector3");
  createFixSizeVector<4>(module, "Vector4");
  createFixSizeVector<6>(module, "Vector6");

  py::class_<MatrixDynSize> matrix_dyn(module, "MatrixDynSize",
                                       py::buffer_protocol());
  matrixDynSizeClassDefinition(matrix_dyn);
  createFixSizeMatrix<3, 3>(module, "Matrix3x3");
  createFixSizeMatrix<4, 4>(module, "Matrix4x4");
  createFixSizeMatrix<6, 6>(module, "Matrix6x6");

  // Positions, Rotations and Transforms.
  py::class_<Position, VectorFixSize<3>>(module, "Position")
      .def(py::init())
      .def(py::init<double, double, double>())
      .def(py::self + py::self)
      .def(py::self - py::self)
      .def(-py::self)
      .def("__repr__", &Position::toString)
      .def_static("Zero", &Position::Zero);


  py::class_<Rotation, MatrixFixSize<3, 3>>(module, "Rotation")
      .def(py::init())
      .def(py::init<double, double, double,  //
                    double, double, double,  //
                    double, double, double>())
      .def("inverse", &Rotation::inverse)
      .def(py::self * py::self)
      .def(
          "__mul__",
          [](const Rotation& r, const Position& p) -> Position {
            return r * p;
          },
          py::is_operator())
      .def("as_rpy", &Rotation::asRPY)
      .def("as_quaternion", &Rotation::asQuaternion)
      .def("__repr__", &Rotation::toString)
      .def_static("Identity", &Rotation::Identity)
      .def_static("RotX", &Rotation::RotX)
      .def_static("RotY", &Rotation::RotY)
      .def_static("RotZ", &Rotation::RotZ)
      .def_static("RPY", &Rotation::RPY)
      .def_static("Quaternion", &Rotation::RotationFromQuaternion);

  py::class_<Transform> transform(module, "Transform");
  transformClassDefinition(transform);

  // Other classes.
  py::class_<Direction, VectorFixSize<3>>(module, "Direction")
      .def(py::init<double, double, double>());

  py::class_<Axis>(module, "Axis")
      .def(py::init<const Direction&, const Position&>())
      .def_property("direction", &Axis::getDirection, &Axis::setDirection)
      .def_property("origin", &Axis::getOrigin, &Axis::setOrigin)
      .def("__repr__", &Axis::toString);

  py::class_<RotationalInertia, MatrixFixSize<3, 3>>(module,
                                                     "RotationalInertia")
      .def(py::init());

  py::class_<SpatialInertia>(module, "SpatialInertia")
      .def(py::init())
      .def(py::init<double, const Position&, const RotationalInertia&>())
      .def("from_rotational_inertia_wrt_center_of_mass",
           &SpatialInertia::fromRotationalInertiaWrtCenterOfMass)
      .def("get_mass", &SpatialInertia::getMass)
      .def("get_center_of_mass", &SpatialInertia::getCenterOfMass)
      .def("get_rotational_inertia_wrt_frame_origin",
           &SpatialInertia::getRotationalInertiaWrtFrameOrigin)
      .def("get_rotational_inertia_wrt_center_of_mass",
           &SpatialInertia::getRotationalInertiaWrtCenterOfMass)
      .def_static("Zero", &SpatialInertia::Zero)
      .def("as_matrix", &SpatialInertia::asMatrix)
      .def(py::self + py::self)
      .def("__repr__", [](const SpatialInertia& inertia) {
        return inertia.asMatrix().toString();
      });

  // Basic twist interface.
  py::class_<Twist>(module, "Twist")
      .def(py::init([](const std::array<double, 3>& lin_velocity,
                       const std::array<double, 3>& ang_velocity) {
        LinVelocity lin_vel(lin_velocity.data(), 3);
        AngVelocity ang_vel(ang_velocity.data(), 3);
        return Twist(lin_vel, ang_vel);
      }));
}
}  // namespace bindings
}  // namespace iDynTree
