#include "pybind11/eigen.h"
#include "pybind11/pybind11.h"
#include "pybind11/stl.h"

#include "drake/bindings/pydrake/documentation_pybind.h"
#include "drake/bindings/pydrake/pydrake_pybind.h"
#include "drake/examples/manipulation_station/manipulation_station.h"
#include "drake/examples/manipulation_station/manipulation_station_hardware_interface.h"  // noqa

using std::make_unique;
using std::unique_ptr;
using std::vector;

namespace drake {
namespace pydrake {

PYBIND11_MODULE(manipulation_station, m) {
  // NOLINTNEXTLINE(build/namespaces): Emulate placement in namespace.
  using namespace drake::systems;
  // NOLINTNEXTLINE(build/namespaces): Emulate placement in namespace.
  using namespace drake::examples::manipulation_station;

  m.doc() = "Bindings for the Manipulation Station example.";
  constexpr auto& doc = pydrake_doc.drake.examples.manipulation_station;

  py::module::import("pydrake.systems.framework");

  // ManipulationStation currently only supports double.
  using T = double;

  py::enum_<IiwaCollisionModel>(
      m, "IiwaCollisionModel", doc.IiwaCollisionModel.doc)
      .value("kNoCollision", IiwaCollisionModel::kNoCollision,
          doc.IiwaCollisionModel.kNoCollision.doc)
      .value("kBoxCollision", IiwaCollisionModel::kBoxCollision,
          doc.IiwaCollisionModel.kBoxCollision.doc)
      .export_values();

  py::enum_<SchunkCollisionModel>(
      m, "SchunkCollisionModel", doc.SchunkCollisionModel.doc)
      .value(
          "kBox", SchunkCollisionModel::kBox, doc.SchunkCollisionModel.kBox.doc)
      .value("kBoxPlusFingertipSpheres",
          SchunkCollisionModel::kBoxPlusFingertipSpheres,
          doc.SchunkCollisionModel.kBoxPlusFingertipSpheres.doc)
      .export_values();

  py::class_<ManipulationStation<T>, Diagram<T>>(
      m, "ManipulationStation", doc.ManipulationStation.doc)
      .def(py::init<double>(), py::arg("time_step") = 0.002,
          doc.ManipulationStation.ctor.doc)
      .def("SetupManipulationClassStation",
          &ManipulationStation<T>::SetupManipulationClassStation,
          py::arg("collision_model") = IiwaCollisionModel::kNoCollision,
          py::arg("schunk_model") = SchunkCollisionModel::kBox,
          doc.ManipulationStation.SetupManipulationClassStation.doc)
      .def("SetupClutterClearingStation",
          &ManipulationStation<T>::SetupClutterClearingStation,
          py::arg("X_WCameraBody") = std::nullopt,
          py::arg("collision_model") = IiwaCollisionModel::kNoCollision,
          py::arg("schunk_model") = SchunkCollisionModel::kBox,
          doc.ManipulationStation.SetupClutterClearingStation.doc)
      .def("SetupPlanarIiwaStation",
          &ManipulationStation<T>::SetupPlanarIiwaStation,
          py::arg("schunk_model") = SchunkCollisionModel::kBox,
          doc.ManipulationStation.SetupPlanarIiwaStation.doc)
      .def("AddManipulandFromFile",
          &ManipulationStation<T>::AddManipulandFromFile, py::arg("model_file"),
          py::arg("X_WObject"),
          doc.ManipulationStation.AddManipulandFromFile.doc)
      .def("RegisterIiwaControllerModel",
          &ManipulationStation<T>::RegisterIiwaControllerModel,
          doc.ManipulationStation.RegisterIiwaControllerModel.doc)
      .def("RegisterRgbdSensor", &ManipulationStation<T>::RegisterRgbdSensor,
          doc.ManipulationStation.RegisterRgbdSensor.doc)
      .def("RegisterWsgControllerModel",
          &ManipulationStation<T>::RegisterWsgControllerModel,
          doc.ManipulationStation.RegisterWsgControllerModel.doc)
      .def("Finalize", py::overload_cast<>(&ManipulationStation<T>::Finalize),
          doc.ManipulationStation.Finalize.doc_0args)
      .def("num_iiwa_joints", &ManipulationStation<T>::num_iiwa_joints,
          doc.ManipulationStation.num_iiwa_joints.doc)
      .def("get_multibody_plant", &ManipulationStation<T>::get_multibody_plant,
          py_rvp::reference_internal,
          doc.ManipulationStation.get_multibody_plant.doc)
      .def("get_mutable_multibody_plant",
          &ManipulationStation<T>::get_mutable_multibody_plant,
          py_rvp::reference_internal,
          doc.ManipulationStation.get_mutable_multibody_plant.doc)
      .def("get_scene_graph", &ManipulationStation<T>::get_scene_graph,
          py_rvp::reference_internal,
          doc.ManipulationStation.get_scene_graph.doc)
      .def("get_mutable_scene_graph",
          &ManipulationStation<T>::get_mutable_scene_graph,
          py_rvp::reference_internal,
          doc.ManipulationStation.get_mutable_scene_graph.doc)
      .def("get_controller_plant",
          &ManipulationStation<T>::get_controller_plant,
          py_rvp::reference_internal,
          doc.ManipulationStation.get_controller_plant.doc)
      .def("GetIiwaPosition", &ManipulationStation<T>::GetIiwaPosition,
          doc.ManipulationStation.GetIiwaPosition.doc)
      .def("SetIiwaPosition",
          overload_cast_explicit<void, systems::Context<T>*,
              const Eigen::Ref<const VectorX<T>>&>(
              &ManipulationStation<T>::SetIiwaPosition),
          py::arg("station_context"), py::arg("q"),
          doc.ManipulationStation.SetIiwaPosition.doc_2args)
      .def("GetIiwaVelocity", &ManipulationStation<T>::GetIiwaVelocity,
          doc.ManipulationStation.GetIiwaVelocity.doc)
      .def("SetIiwaVelocity",
          overload_cast_explicit<void, systems::Context<T>*,
              const Eigen::Ref<const VectorX<T>>&>(
              &ManipulationStation<T>::SetIiwaVelocity),
          py::arg("station_context"), py::arg("v"),
          doc.ManipulationStation.SetIiwaVelocity.doc_2args)
      .def("GetWsgPosition", &ManipulationStation<T>::GetWsgPosition,
          doc.ManipulationStation.GetWsgPosition.doc)
      .def("SetWsgPosition",
          overload_cast_explicit<void, systems::Context<T>*, const T&>(
              &ManipulationStation<T>::SetWsgPosition),
          py::arg("station_context"), py::arg("q"),
          doc.ManipulationStation.SetWsgPosition.doc_2args)
      .def("GetWsgVelocity", &ManipulationStation<T>::GetWsgVelocity,
          doc.ManipulationStation.GetWsgVelocity.doc)
      .def("SetWsgVelocity",
          overload_cast_explicit<void, systems::Context<T>*, const T&>(
              &ManipulationStation<T>::SetWsgVelocity),
          py::arg("station_context"), py::arg("v"),
          doc.ManipulationStation.SetWsgVelocity.doc_2args)
      .def("GetStaticCameraPosesInWorld",
          &ManipulationStation<T>::GetStaticCameraPosesInWorld,
          py_rvp::reference_internal,
          doc.ManipulationStation.GetStaticCameraPosesInWorld.doc)
      .def("get_camera_names", &ManipulationStation<T>::get_camera_names,
          doc.ManipulationStation.get_camera_names.doc)
      .def("SetWsgGains", &ManipulationStation<T>::SetWsgGains,
          doc.ManipulationStation.SetWsgGains.doc)
      .def("SetIiwaPositionGains",
          &ManipulationStation<T>::SetIiwaPositionGains,
          doc.ManipulationStation.SetIiwaPositionGains.doc)
      .def("SetIiwaVelocityGains",
          &ManipulationStation<T>::SetIiwaVelocityGains,
          doc.ManipulationStation.SetIiwaVelocityGains.doc)
      .def("SetIiwaIntegralGains",
          &ManipulationStation<T>::SetIiwaIntegralGains,
          doc.ManipulationStation.SetIiwaIntegralGains.doc);

  py::class_<ManipulationStationHardwareInterface, Diagram<double>>(m,
      "ManipulationStationHardwareInterface",
      doc.ManipulationStationHardwareInterface.doc)
      .def(py::init<const std::vector<std::string>>(),
          py::arg("camera_names") = std::vector<std::string>{},
          doc.ManipulationStationHardwareInterface.ctor.doc)
      .def("Connect", &ManipulationStationHardwareInterface::Connect,
          py::arg("wait_for_cameras") = true,
          doc.ManipulationStationHardwareInterface.Connect.doc)
      .def("get_controller_plant",
          &ManipulationStationHardwareInterface::get_controller_plant,
          py_rvp::reference_internal,
          doc.ManipulationStationHardwareInterface.get_controller_plant.doc)
      .def("get_camera_names",
          &ManipulationStationHardwareInterface::get_camera_names,
          py_rvp::reference_internal,
          doc.ManipulationStationHardwareInterface.get_camera_names.doc)
      .def("num_iiwa_joints",
          &ManipulationStationHardwareInterface::num_iiwa_joints,
          doc.ManipulationStationHardwareInterface.num_iiwa_joints.doc);

  ExecuteExtraPythonCode(m);
}

}  // namespace pydrake
}  // namespace drake
