//
// Created by Sebastian Ebert on 31.03.20.
//
#include <moveit/move_group_interface/move_group_interface.h>

#include <moveit_msgs/DisplayTrajectory.h>
#include <moveit_msgs/CollisionObject.h>

#include "panda_mqtt_connector/Trajectory.h"
#include "panda_mqtt_connector/Waypoint.h"

#include <moveit/planning_scene_interface/planning_scene_interface.h>
#include <moveit/trajectory_processing/iterative_time_parameterization.h>
#include <tf2/transform_datatypes.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.h>

#include <boost/optional.hpp>

#include "util/TrajectoryUtil.h"
#include "TimedPlanner.h"

#include <cmath>
#include <utility>


void
constructPlate(std::vector<moveit_msgs::CollisionObject> &collision_objects, double x_dimension, double y_dimension) {

  moveit_msgs::CollisionObject plate;

  plate.header.frame_id = "panda_link0";
  plate.id = "plate";

  plate.primitives.resize(1);
  plate.primitives[0].type = 1u;
  plate.primitives[0].dimensions.resize(3);
  plate.primitives[0].dimensions[0] = x_dimension;
  plate.primitives[0].dimensions[1] = y_dimension;
  plate.primitives[0].dimensions[2] = 0.1;

  plate.primitive_poses.resize(1);
  plate.primitive_poses[0].position.x = 0;
  plate.primitive_poses[0].position.y = 0;
  plate.primitive_poses[0].position.z = -0.1;

  plate.operation = plate.ADD;

  collision_objects.push_back(plate);
}

void TimedPlanner::loadCircularTrajectory(double radius, geometry_msgs::Point origin, int granularity) {

  nextTrajectory = panda_mqtt_connector::Trajectory{};

  double step_size_rad = 2 * M_PI / granularity;

  for (int i = 1; i <= granularity; i++) {

    int pos = i % granularity;

    double next_rad = step_size_rad * (2 * pos > granularity ? granularity - pos : pos);

    tf2::Quaternion orientation; // the RPY constructor is deprecated
    orientation.setRPY(0, M_PI, next_rad);

    panda_mqtt_connector::Waypoint nextWaypoint{};
    nextWaypoint.pose.orientation = tf2::toMsg(orientation);
    nextWaypoint.pose.position.x = origin.x + radius * sin(next_rad);
    nextWaypoint.pose.position.y = origin.y + radius * cos(next_rad);
    nextWaypoint.pose.position.z = origin.z;
    nextWaypoint.mode = TrajectoryUtil::FLUID_PATH;

    nextTrajectory->waypoints.push_back(nextWaypoint);
  }
}

bool TimedPlanner::updateWaypoints() {
  if (nextTrajectory.is_initialized()) {
    currentTrajectory = nextTrajectory.get();
    nextTrajectory.reset();
    return true;
  } else {
    return false;
  }
}

std::vector<moveit_msgs::RobotTrajectory> TimedPlanner::split(moveit::planning_interface::MoveGroupInterface &group,
                                                              const moveit::planning_interface::MoveGroupInterface::Plan &plan) {

  std::vector<moveit_msgs::RobotTrajectory> trajectories;

  auto currentState = group.getCurrentState();
  auto robotModel = currentState->getRobotModel();

  robot_trajectory::RobotTrajectory rt(robotModel, "panda_arm");
  rt.setRobotTrajectoryMsg(*currentState, plan.start_state_, plan.trajectory_);

  for (int wpIndex = 0; wpIndex < rt.getWayPointCount() - 1; ++wpIndex) {
    robot_trajectory::RobotTrajectory wpTrajectory{robotModel, "panda_arm"};
    moveit::core::RobotState waypointState{robotModel};
    wpTrajectory.insertWayPoint(0, rt.getWayPoint(wpIndex), 0);
    wpTrajectory.insertWayPoint(1, rt.getWayPoint(wpIndex + 1), rt.getWayPointDurationFromPrevious(wpIndex + 1));
    moveit_msgs::RobotTrajectory wpTrajectoryMsg{};
    wpTrajectory.getRobotTrajectoryMsg(wpTrajectoryMsg);
    trajectories.push_back(wpTrajectoryMsg);
  }

  return trajectories;

}

void TimedPlanner::doMotion() {
  updateWaypoints();

  do {
    for (auto &waypoint : currentTrajectory.waypoints) {
      moveit::planning_interface::MoveGroupInterface::Plan plan;
      if (TrajectoryUtil::computePathToPose(*group, plan, waypoint.pose, waypoint.mode, 1, 1)) {
        for (auto trajectory : split(*group, plan)) {
          if (nextTrajectory.is_initialized()) {
            ROS_WARN_STREAM("terminating current trajectory");
            return;
          }
          group->getCurrentState()->copyJointGroupPositions("panda_arm", trajectory.joint_trajectory.points[0].positions);
          ROS_INFO_STREAM(
              "Moving to the next waypoint with speed factor " << motionSpeedFactor << " in mode " << waypoint.mode);
          TrajectoryUtil::applyMotionSpeedFactor(trajectory, motionSpeedFactor);
          moveit_msgs::MoveItErrorCodes errorCode = group->execute(trajectory);
          if (errorCode.val != errorCode.SUCCESS) {
            ROS_ERROR_STREAM("Error Code from executing segment: " << errorCode);
          }
        }
      } else {
        ROS_ERROR_STREAM("Unable to find " << waypoint.mode << " path to the next waypoint." << std::endl << waypoint);
      }
    }
  } while (currentTrajectory.loop);

  currentTrajectory.waypoints.clear(); // the trajectory is completed and can be cleared
}

void TimedPlanner::newTrajectoryCallback(const panda_mqtt_connector::Trajectory::ConstPtr &msg) {
  ROS_INFO_STREAM("Received new trajectory with " << msg->waypoints.size() << " waypoints.");
  nextTrajectory = *msg;
}

void TimedPlanner::newMotionSpeedFactorCallback(const std_msgs::Float64ConstPtr &msg) {
  ROS_INFO_STREAM("Received new motionSpeedFactor " << msg->data);

  if (msg->data <= 0 || msg->data > 1.0) {
    ROS_ERROR_STREAM(
        "Invalid motion speed factor " << msg->data << ". Must be in (0,1]. Setting to default " << default_velocity);
    motionSpeedFactor = default_velocity;
  } else {
    motionSpeedFactor = msg->data;
  }

}

TimedPlanner::TimedPlanner(ros::NodeHandle &node_handle, moveit::planning_interface::MoveGroupInterface &group,
                           const double defaultVelocity, std::string defaultPlanningMode) :
    node_handle(&node_handle), group(&group), default_velocity(defaultVelocity),
    default_planning_mode(std::move(defaultPlanningMode)), motionSpeedFactor(defaultVelocity) {

  std::string defaultTrajectory = "<no value provided>";
  node_handle.getParam("default_trajectory", defaultTrajectory);

  if (defaultTrajectory == "circle") {
    ROS_INFO_STREAM("loading default trajectory 'circle'");
    geometry_msgs::Point origin;
    origin.z = .3;
    loadCircularTrajectory(0.6, origin, 20);
  } else if (defaultTrajectory == "square") {
    ROS_INFO_STREAM("loading default trajectory 'square'");
    loadSquareTrajectory();
  } else {
    ROS_WARN_STREAM("Not loading a default trajectory! There is no '" + defaultTrajectory + "'.");
  }

}

void TimedPlanner::loadSquareTrajectory() {
  panda_mqtt_connector::Waypoint waypoint_1;
  waypoint_1.pose.position.x = 0.4;
  waypoint_1.pose.position.y = 0.4;
  waypoint_1.pose.position.z = 0.3;
  waypoint_1.pose.orientation.w = 0;
  waypoint_1.pose.orientation.x = 0.7071067811865476;
  waypoint_1.pose.orientation.y = 0.7071067811865476;
  waypoint_1.pose.orientation.z = 0;
  waypoint_1.mode = TrajectoryUtil::FLUID_PATH;
  panda_mqtt_connector::Waypoint waypoint_2{waypoint_1};
  panda_mqtt_connector::Waypoint waypoint_3{waypoint_1};
  panda_mqtt_connector::Waypoint waypoint_4{waypoint_1};

  currentTrajectory.waypoints.push_back(waypoint_1);

  waypoint_2.pose.position.y *= -1;
  currentTrajectory.waypoints.push_back(waypoint_2);

  waypoint_3.pose.position.x *= -1;
  waypoint_3.pose.position.y *= -1;
  currentTrajectory.waypoints.push_back(waypoint_3);

  waypoint_4.pose.position.x *= -1;
  currentTrajectory.waypoints.push_back(waypoint_4);
}


int main(int argc, char **argv) {

  // setup this ros-node
  ros::init(argc, argv, "timed_planner");
  ros::NodeHandle node_handle("panda_mqtt_connector");
  ros::AsyncSpinner spinner(2);
  spinner.start();

  // wait for robot init of robot_state_initializer
  ros::Duration(3.0).sleep();
  ROS_INFO(">>>>>>>>>>>>>>>>> START UP FINISHED <<<<<<<<<<<<<<<< ");

  // Initialize start state of robot and target trajectory.
  moveit::planning_interface::MoveGroupInterface group("panda_arm");
  moveit::planning_interface::PlanningSceneInterface planning_scene_interface;

  // add the ground plate
  std::vector<moveit_msgs::CollisionObject> collision_objects;
  constructPlate(collision_objects, 5.0, 5.0);
  planning_scene_interface.applyCollisionObjects(collision_objects);

  // everything is set up, we can start the planner now
  TimedPlanner planner(node_handle, group, .5, TrajectoryUtil::FLUID_PATH);
  ros::Subscriber sub = node_handle.subscribe("trajectory_update", 1000, &TimedPlanner::newTrajectoryCallback,
                                              &planner);
  ros::Subscriber velocitySubscriber = node_handle.subscribe("max_velocity", 1000,
                                                             &TimedPlanner::newMotionSpeedFactorCallback,
                                                             &planner);

  // init motionSpeedFactor
  double factor = -1;
  if (node_handle.getParam("robot_speed_factor", factor)) {
    ros::Publisher pub = node_handle.advertise<std_msgs::Float64>("max_velocity", 1000);
    std_msgs::Float64 velocity;
    velocity.data = factor;
    pub.publish(velocity);
  } else {
    ROS_WARN_STREAM("No default robot_speed_factor provided!");
  }

  while (ros::ok()) {

    // execute a trajectory
    planner.doMotion();
    ros::Rate(100).sleep();
    ros::spinOnce();
  }

  return 0;
}