#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/eigen.h>

#include <iostream>
#include <limits>

#include "covalent_wrapper.hpp"

namespace py = pybind11;
using namespace covalentcore;

/*
Summary:
  Register Python bindings for high-level trajectory/sample shaping interfaces.
Args:
  m: Pybind11 module object provided by `PYBIND11_MODULE`.
Returns:
  None.
Side Effects:
  Exposes `RobotState`, `ShapedSample`, `ShapedTrajectory`, and `ShaperInterface`.
Raises:
  None.
Preconditions:
  Wrapper dependencies (Torch, Eigen, pybind11) are available at import time.
*/
PYBIND11_MODULE(covalent, m) {
    m.doc() = "Python bindings for Covalent ShaperInterface";

    py::class_<RobotState>(m, "RobotState")
        .def(py::init<>())
        .def_readwrite("joint_angles", &RobotState::joint_angles)
        .def_readwrite("tcp_position", &RobotState::tcp_position);

    py::class_<ShapedSample>(m, "ShapedSample")
        .def_readwrite("positions", &ShapedSample::positions)
        .def_readwrite("velocities", &ShapedSample::velocities)
        .def_readwrite("accelerations", &ShapedSample::accelerations);

    py::class_<ShapedTrajectory>(m, "ShapedTrajectory")
        .def_readwrite("positions", &ShapedTrajectory::positions)
        .def_readwrite("velocities", &ShapedTrajectory::velocities)
        .def_readwrite("accelerations", &ShapedTrajectory::accelerations)
        .def_readwrite("time", &ShapedTrajectory::time);

    py::class_<ShaperInterface>(m, "ShaperInterface")
        .def(py::init<double, std::string, const std::string&, const std::string&, double, double, int, int, double>(),
             py::arg("sample_time"),
             py::arg("model_directory"),
             py::arg("python_src_root"),
             py::arg("urdf_filepath"),
             py::arg("side_length") = std::numeric_limits<double>::quiet_NaN(),
             py::arg("base_height") = std::numeric_limits<double>::quiet_NaN(),
             py::arg("num_axes") = 0,
             py::arg("num_joints") = 0,
             py::arg("prob_thresh") = defaultProbThresh)
        .def("shape_sample", &ShaperInterface::shapeSample,
             py::arg("command"),
             py::arg("state"),
             py::arg("command_dot") = std::nullopt,
             py::arg("command_ddot") = std::nullopt)
        .def("process_sample", &ShaperInterface::processSample,
             py::arg("command"),
             py::arg("state"),
             py::arg("command_dot") = std::nullopt,
             py::arg("command_ddot") = std::nullopt,
             py::arg("enable_vibration_control") = true)
        .def("shape_trajectory", &ShaperInterface::shapeTrajectory,
             py::arg("command"),
             py::arg("states"),
             py::arg("command_dot") = std::nullopt,
             py::arg("command_ddot") = std::nullopt,
             py::arg("time_vector") = std::nullopt)
        .def("process_trajectory", &ShaperInterface::processTrajectory,
             py::arg("command"),
             py::arg("states"),
             py::arg("command_dot") = std::nullopt,
             py::arg("command_ddot") = std::nullopt,
             py::arg("time_vector") = std::nullopt,
             py::arg("enable_vibration_control") = true)
        .def("finalize", &ShaperInterface::finalize)
        .def("reset", &ShaperInterface::reset)
        .def_property_readonly("sample_time", &ShaperInterface::sampleTime)
        .def_property_readonly("axes", &ShaperInterface::axes)
        .def_property_readonly("joints", &ShaperInterface::joints);
}
