#include "common/curve/include/s_curve.h"
#include "planner/s_curve_planner/include/s_curve_position_planner.h"

#include <algorithm>
#include <stdexcept>
#include <vector>

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

namespace py = pybind11;

namespace {

py::dict sample_trajectory(const cmvr::SCurve& curve, const cmvr::SCurveProfile& profile, double dt)
{
  if (dt <= 0.0) {
    throw std::invalid_argument("dt must be > 0");
  }

  std::vector<double> times;
  std::vector<double> positions;
  std::vector<double> velocities;
  std::vector<double> accelerations;
  std::vector<double> jerks;

  if (profile.total_time <= 0.0) {
    times.push_back(0.0);
    positions.push_back(profile.p0);
    velocities.push_back(profile.v0);
    accelerations.push_back(profile.a0);
    jerks.push_back(0.0);
  } else {
    for (double t = 0.0; t <= profile.total_time + 1e-12; t += dt) {
      const double sample_time = std::min(t, profile.total_time);
      times.push_back(sample_time);
      positions.push_back(curve.getPositionAtTime(profile, sample_time));
      velocities.push_back(curve.getVelocityAtTime(profile, sample_time));
      accelerations.push_back(curve.getAccelerationAtTime(profile, sample_time));
      jerks.push_back(curve.getJerkAtTime(profile, sample_time));
    }

    if (std::abs(times.back() - profile.total_time) > 1e-9) {
      times.push_back(profile.total_time);
      positions.push_back(curve.getPositionAtTime(profile, profile.total_time));
      velocities.push_back(curve.getVelocityAtTime(profile, profile.total_time));
      accelerations.push_back(curve.getAccelerationAtTime(profile, profile.total_time));
      jerks.push_back(curve.getJerkAtTime(profile, profile.total_time));
    }
  }

  py::dict result;
  result["times"] = times;
  result["positions"] = positions;
  result["velocities"] = velocities;
  result["accelerations"] = accelerations;
  result["jerks"] = jerks;
  return result;
}

}  // namespace

PYBIND11_MODULE(_s_curve_native, m)
{
  m.doc() = "pybind11 bindings for cmvr::SCurve";

  py::class_<cmvr::SCurveProfile>(m, "SCurveProfile")
    .def(py::init<>())
    .def_readwrite("t1", &cmvr::SCurveProfile::t1)
    .def_readwrite("t2", &cmvr::SCurveProfile::t2)
    .def_readwrite("t3", &cmvr::SCurveProfile::t3)
    .def_readwrite("t4", &cmvr::SCurveProfile::t4)
    .def_readwrite("t5", &cmvr::SCurveProfile::t5)
    .def_readwrite("t6", &cmvr::SCurveProfile::t6)
    .def_readwrite("t7", &cmvr::SCurveProfile::t7)
    .def_readwrite("total_time", &cmvr::SCurveProfile::total_time)
    .def_readwrite("j_max", &cmvr::SCurveProfile::j_max)
    .def_readwrite("a_max", &cmvr::SCurveProfile::a_max)
    .def_readwrite("v_max", &cmvr::SCurveProfile::v_max)
    .def_readwrite("distance", &cmvr::SCurveProfile::distance)
    .def_readwrite("direction", &cmvr::SCurveProfile::direction)
    .def_readwrite("v_cruise", &cmvr::SCurveProfile::v_cruise)
    .def_readwrite("a_limit", &cmvr::SCurveProfile::a_limit)
    .def_readwrite("p0", &cmvr::SCurveProfile::p0)
    .def_readwrite("v0", &cmvr::SCurveProfile::v0)
    .def_readwrite("a0", &cmvr::SCurveProfile::a0)
    .def_readwrite("vf", &cmvr::SCurveProfile::vf);

  py::class_<cmvr::SCurve>(m, "SCurve")
    .def(py::init<double, double, double>(),
      py::arg("max_velocity") = 3.0,
      py::arg("max_acceleration") = 10.0,
      py::arg("max_jerk") = 50.0)
    .def("setConstraints", &cmvr::SCurve::setConstraints,
      py::arg("max_velocity"),
      py::arg("max_acceleration"),
      py::arg("max_jerk"))
    .def("getConstraints", [](const cmvr::SCurve& self) {
      double max_velocity = 0.0;
      double max_acceleration = 0.0;
      double max_jerk = 0.0;
      self.getConstraints(max_velocity, max_acceleration, max_jerk);
      return py::make_tuple(max_velocity, max_acceleration, max_jerk);
    })
    .def("calculateProfile", &cmvr::SCurve::calculateProfile,
      py::arg("start_position"),
      py::arg("end_position"),
      py::arg("start_velocity") = 0.0,
      py::arg("end_velocity") = 0.0)
    .def("getPositionAtTime", &cmvr::SCurve::getPositionAtTime,
      py::arg("profile"),
      py::arg("t"))
    .def("getVelocityAtTime", &cmvr::SCurve::getVelocityAtTime,
      py::arg("profile"),
      py::arg("t"))
    .def("getAccelerationAtTime", &cmvr::SCurve::getAccelerationAtTime,
      py::arg("profile"),
      py::arg("t"))
    .def("getJerkAtTime", &cmvr::SCurve::getJerkAtTime,
      py::arg("profile"),
      py::arg("t"))
    .def("sampleTrajectory", &sample_trajectory,
      py::arg("profile"),
      py::arg("dt"));

  py::class_<cmvr::SCurvePositionPlannerState>(m, "SCurvePositionPlannerState")
    .def(py::init<>())
    .def_readwrite("position", &cmvr::SCurvePositionPlannerState::position)
    .def_readwrite("velocity", &cmvr::SCurvePositionPlannerState::velocity)
    .def_readwrite("acceleration", &cmvr::SCurvePositionPlannerState::acceleration)
    .def_readwrite("jerk", &cmvr::SCurvePositionPlannerState::jerk)
    .def_readwrite("target_position", &cmvr::SCurvePositionPlannerState::target_position)
    .def_readwrite("is_moving", &cmvr::SCurvePositionPlannerState::is_moving);

  py::class_<cmvr::SCurvePositionPlanner1D>(m, "SCurvePositionPlanner1D")
    .def(py::init<double, double, double>(),
      py::arg("max_velocity") = 3.0,
      py::arg("max_acceleration") = 10.0,
      py::arg("max_jerk") = 50.0)
    .def("setConstraints", &cmvr::SCurvePositionPlanner1D::setConstraints,
      py::arg("max_velocity"),
      py::arg("max_acceleration"),
      py::arg("max_jerk"))
    .def("getConstraints", [](const cmvr::SCurvePositionPlanner1D& self) {
      double max_velocity = 0.0;
      double max_acceleration = 0.0;
      double max_jerk = 0.0;
      self.getConstraints(max_velocity, max_acceleration, max_jerk);
      return py::make_tuple(max_velocity, max_acceleration, max_jerk);
    })
    .def("setPositionGain", &cmvr::SCurvePositionPlanner1D::setPositionGain,
      py::arg("position_gain"))
    .def("getPositionGain", &cmvr::SCurvePositionPlanner1D::getPositionGain)
    .def("initialize", &cmvr::SCurvePositionPlanner1D::initialize,
      py::arg("position"),
      py::arg("velocity") = 0.0,
      py::arg("acceleration") = 0.0)
    .def("reset", &cmvr::SCurvePositionPlanner1D::reset)
    .def("setTarget", &cmvr::SCurvePositionPlanner1D::setTarget,
      py::arg("target_position"))
    .def("update", &cmvr::SCurvePositionPlanner1D::update,
      py::arg("dt"))
    .def("getState", &cmvr::SCurvePositionPlanner1D::getState,
      py::return_value_policy::reference_internal)
    .def("getPosition", &cmvr::SCurvePositionPlanner1D::getPosition)
    .def("getVelocity", &cmvr::SCurvePositionPlanner1D::getVelocity)
    .def("getAcceleration", &cmvr::SCurvePositionPlanner1D::getAcceleration)
    .def("getJerk", &cmvr::SCurvePositionPlanner1D::getJerk)
    .def("getTargetPosition", &cmvr::SCurvePositionPlanner1D::getTargetPosition)
    .def("isMoving", &cmvr::SCurvePositionPlanner1D::isMoving);
}
