//
// Created by sebastian on 19.05.20.
//
#include "grasp_util.h"
#include <tf2/transform_datatypes.h>
#include <franka_gripper/HomingAction.h>
#include <franka_gripper/MoveAction.h>
#include <franka_gripper/HomingGoal.h>

void GraspUtil::setupOpenGripper(trajectory_msgs::JointTrajectory &posture, double amount) {
    posture.joint_names.resize(2);
    posture.joint_names[0] = "panda_finger_joint1";
    posture.joint_names[1] = "panda_finger_joint2";

    posture.points.resize(1);
    posture.points[0].positions.resize(2);
    posture.points[0].positions[0] = amount / 2;
    posture.points[0].positions[1] = amount / 2;
    posture.points[0].time_from_start = ros::Duration(0.5);
}

void GraspUtil::setupClosedGripper(trajectory_msgs::JointTrajectory &posture, double amount) {
    posture.joint_names.resize(2);
    posture.joint_names[0] = "panda_finger_joint1";
    posture.joint_names[1] = "panda_finger_joint2";

    posture.points.resize(1);
    posture.points[0].positions.resize(2);
    posture.points[0].positions[0] = amount / 2;
    posture.points[0].positions[1] = amount / 2;
    posture.points[0].time_from_start = ros::Duration(0.5);
}

geometry_msgs::Pose
GraspUtil::getPickFromBack(geometry_msgs::Pose &object_pose_msg, geometry_msgs::Vector3 dimensions, geometry_msgs::Quaternion gripper_orientation_msg, double reach){

    double approach_roll, approach_pitch, approach_yaw;
    approach_roll =  -M_PI_2;
    approach_pitch = -M_PI_4;
    approach_yaw =   -M_PI_4;
    tf2::Quaternion approach_orientation;
    approach_orientation.setRPY(approach_roll, approach_pitch, approach_yaw);

    return getPickFromSidePose(object_pose_msg, dimensions, gripper_orientation_msg,
            approach_orientation, 0, -0.1, reach);
}

geometry_msgs::Pose
GraspUtil::getPickFromFront(geometry_msgs::Pose &object_pose_msg, geometry_msgs::Vector3 dimensions, geometry_msgs::Quaternion gripper_orientation_msg, double reach){

    double approach_roll, approach_pitch, approach_yaw;
    approach_roll =  M_PI_2;
    approach_pitch = -M_PI_4 - M_PI_2;
    approach_yaw =   -M_PI_4;
    tf2::Quaternion approach_orientation;
    approach_orientation.setRPY(approach_roll, approach_pitch, approach_yaw);

    return getPickFromSidePose(object_pose_msg, dimensions, gripper_orientation_msg,
                               approach_orientation, 0, 0.1, reach);
}

geometry_msgs::Pose
GraspUtil::getPickFromLeft(geometry_msgs::Pose &object_pose_msg, geometry_msgs::Vector3 dimensions, geometry_msgs::Quaternion gripper_orientation_msg, double reach){

    double approach_roll, approach_pitch, approach_yaw;
    approach_roll =  M_PI_2;
    approach_pitch = -M_PI_4 - M_PI_2;
    approach_yaw =   -M_PI_4 - M_PI_2;
    tf2::Quaternion approach_orientation;
    approach_orientation.setRPY(approach_roll, approach_pitch, approach_yaw);

    return getPickFromSidePose(object_pose_msg, dimensions, gripper_orientation_msg,
                               approach_orientation, -0.1, 0, reach);
}

geometry_msgs::Pose
GraspUtil::getPickFromRight(geometry_msgs::Pose &object_pose_msg, geometry_msgs::Vector3 dimensions, geometry_msgs::Quaternion gripper_orientation_msg, double reach){

    double approach_roll, approach_pitch, approach_yaw;
    approach_roll =  -M_PI_2;
    approach_pitch = -M_PI_4;
    approach_yaw =   -M_PI_4 - M_PI_2;
    tf2::Quaternion approach_orientation;
    approach_orientation.setRPY(approach_roll, approach_pitch, approach_yaw);

    return getPickFromSidePose(object_pose_msg, dimensions, gripper_orientation_msg,
                               approach_orientation, 0.1, 0, reach);
}

geometry_msgs::Pose
GraspUtil::getPickFromSidePose(geometry_msgs::Pose &object_pose_msg, geometry_msgs::Vector3 dimensions, geometry_msgs::Quaternion gripper_orientation_msg,
                    tf2::Quaternion approach_orientation, float y_distance_diff, float x_distance_diff, double reach){

    // object orientation
    tf2::Quaternion object_orientation;
    tf2::fromMsg(object_pose_msg.orientation, object_orientation);

    // gripper orientation
    tf2::Quaternion gripper_orientation;
    tf2::fromMsg(gripper_orientation_msg, gripper_orientation);
    tf2::Quaternion orientation = object_orientation * gripper_orientation * approach_orientation;

    geometry_msgs::Pose target_pose;
    target_pose.orientation = tf2::toMsg(orientation);
    target_pose.position.x = object_pose_msg.position.x + x_distance_diff;
    target_pose.position.y = object_pose_msg.position.y + y_distance_diff;

    if(x_distance_diff > 0){
        target_pose.position.x = object_pose_msg.position.x + (0.5 * dimensions.x) + FRANKA_GRIPPER_DISTANCE - reach;
    }

    if(x_distance_diff < 0){
        target_pose.position.x = object_pose_msg.position.x - (0.5 * dimensions.x) - FRANKA_GRIPPER_DISTANCE + reach;
    }

    if(y_distance_diff > 0){
        target_pose.position.y = object_pose_msg.position.y + (0.5 * dimensions.y) + FRANKA_GRIPPER_DISTANCE - reach;
    }

    if(y_distance_diff < 0){
        target_pose.position.y = object_pose_msg.position.y - (0.5 * dimensions.y) - FRANKA_GRIPPER_DISTANCE + reach;
    }

    target_pose.position.z = object_pose_msg.position.z;

    return target_pose;
}

geometry_msgs::Pose
GraspUtil::getPickFromAbovePose(geometry_msgs::Pose &object_pose_msg, geometry_msgs::Vector3 dimensions,
                                geometry_msgs::Quaternion gripper_orientation_msg, double reach) {

    // object orientation
    tf2::Quaternion object_orientation;
    tf2::fromMsg(object_pose_msg.orientation, object_orientation);

    // gripper orientation
    tf2::Quaternion gripper_orientation;
    tf2::fromMsg(gripper_orientation_msg, gripper_orientation);

    // approach orientation
    double approach_roll, approach_pitch, approach_yaw;
    approach_roll = M_PI;  // turn upside down
    approach_pitch = 0;
    approach_yaw = 0;
    tf2::Quaternion approach_orientation;
    approach_orientation.setRPY(approach_roll, approach_pitch, approach_yaw);

    tf2::Quaternion orientation = object_orientation * gripper_orientation * approach_orientation;

    // Actually, the maximum reach should be FRANKA_GRIPPER_FINGER_LENGTH - FRANKA_GRIPPER_FINGERTIP_SIZE/2, but half a
    // finger length is considered a safety distance.
    if (reach > FRANKA_GRIPPER_FINGER_LENGTH - FRANKA_GRIPPER_FINGERTIP_SIZE) {
        ROS_WARN_STREAM("Reach around object reduced from " << reach << " to " << FRANKA_GRIPPER_FINGER_LENGTH -
                                                                                  FRANKA_GRIPPER_FINGERTIP_SIZE << ".");
        reach = FRANKA_GRIPPER_FINGER_LENGTH - FRANKA_GRIPPER_FINGERTIP_SIZE;
    } else if (reach < -FRANKA_GRIPPER_FINGERTIP_SIZE / 2) {
        ROS_WARN_STREAM("Reach around object increased from " << reach << " to 0." << ". Minimum allowed reach is "
                                                              << (-FRANKA_GRIPPER_FINGERTIP_SIZE / 2) << ".");
        reach = 0;
    }

    // The pose is 10 cm straight above the object to pick. This only works on upright boxes.
    geometry_msgs::Pose target_pose;
    target_pose.orientation = tf2::toMsg(orientation);
    target_pose.position.x = object_pose_msg.position.x;
    target_pose.position.y = object_pose_msg.position.y;
    target_pose.position.z = object_pose_msg.position.z + (0.5 * dimensions.z) + FRANKA_GRIPPER_DISTANCE - reach;

    return target_pose;
}

bool
GraspUtil::pickFromSide(moveit::planning_interface::MoveGroupInterface &move_group, geometry_msgs::Pose &pick_pose,
                        shape_msgs::SolidPrimitive &object_to_pick, geometry_msgs::Vector3 &approach_direction,
                         double open_amount, std::string supporting_surface_id, std::string object_to_pick_id, bool plan_only) {

    std::vector<moveit_msgs::Grasp> grasps;
    grasps.resize(1);
    grasps[0].grasp_pose.header.frame_id = "panda_link0";
    grasps[0].grasp_pose.pose = pick_pose;

    geometry_msgs::Vector3Stamped stamped_approach_dir;
    stamped_approach_dir.vector = approach_direction;

    // Setup pre-grasp approach
    grasps[0].pre_grasp_approach.direction.header.frame_id = "panda_link0";
    grasps[0].pre_grasp_approach.direction = stamped_approach_dir;
    grasps[0].pre_grasp_approach.min_distance = 0.095;
    grasps[0].pre_grasp_approach.desired_distance = 0.115;

    // Setup post-grasp retreat
    grasps[0].post_grasp_retreat.direction.header.frame_id = "panda_link0";
    grasps[0].post_grasp_retreat.direction.vector.z = 1.0;
    grasps[0].post_grasp_retreat.min_distance = 0.1;
    grasps[0].post_grasp_retreat.desired_distance = 0.25;

    setupOpenGripper(grasps[0].pre_grasp_posture, open_amount);

    if(object_to_pick.type == shape_msgs::SolidPrimitive::CYLINDER){

        if(object_to_pick.dimensions[1] >= 0.08){
            ROS_ERROR("Could not grasp: Cylinder object to big.");
            return false;
        }

        setupClosedGripper(grasps[0].grasp_posture, object_to_pick.dimensions[1]);
    }else{
        if (object_to_pick.dimensions[0] <= 0.08) {
            setupClosedGripper(grasps[0].grasp_posture, object_to_pick.dimensions[0]);
        }

        if (object_to_pick.dimensions[0] > 0.08 && object_to_pick.dimensions[1] <= 0.08) {
            setupClosedGripper(grasps[0].grasp_posture, object_to_pick.dimensions[1]);
        }

        if (object_to_pick.dimensions[0] > 0.08 && object_to_pick.dimensions[1] > 0.08) {
            ROS_ERROR("Could not grasp: Object to big (in any dimension).");
            return false;
        }
    }

    move_group.setSupportSurfaceName(supporting_surface_id);

    moveit_msgs::MoveItErrorCodes miec = move_group.pick(object_to_pick_id, grasps, plan_only);

    if (miec.val == miec.SUCCESS) {
        return true;
    }

    ROS_ERROR_STREAM("Error while executing pick task. MoveItErrorCode: " << miec.val);

    return false;
}

bool
GraspUtil::pickFromAbove(moveit::planning_interface::MoveGroupInterface &move_group, geometry_msgs::Pose &pick_pose,
                         geometry_msgs::Vector3 dimensions,
                         double open_amount, std::string supporting_surface_id, std::string object_to_pick_id, bool plan_only) {

    std::vector<moveit_msgs::Grasp> grasps;
    grasps.resize(1);
    grasps[0].grasp_pose.header.frame_id = "panda_link0";
    grasps[0].grasp_pose.pose = pick_pose;

    // Setup pre-grasp approach
    grasps[0].pre_grasp_approach.direction.header.frame_id = "panda_link0";
    grasps[0].pre_grasp_approach.direction.vector.z = -1.0;
    grasps[0].pre_grasp_approach.min_distance = 0.095;
    grasps[0].pre_grasp_approach.desired_distance = 0.115;

    // Setup post-grasp retreat
    grasps[0].post_grasp_retreat.direction.header.frame_id = "panda_link0";
    grasps[0].post_grasp_retreat.direction.vector.z = 1.0;
    grasps[0].post_grasp_retreat.min_distance = 0.1;
    grasps[0].post_grasp_retreat.desired_distance = 0.25;

    // Setup the opening / closing - mechanism of the gripper

    setupOpenGripper(grasps[0].pre_grasp_posture, open_amount);

    if (dimensions.x <= 0.08) {
        setupClosedGripper(grasps[0].grasp_posture, dimensions.x);
    }

    if (dimensions.x > 0.08 && dimensions.y <= 0.08) {
        setupClosedGripper(grasps[0].grasp_posture, dimensions.y);
    }

    if (dimensions.x > 0.08 && dimensions.y > 0.08) {
        ROS_ERROR("Could not grasp: Object to big (in any dimension).");
        return false;
    }

    move_group.setSupportSurfaceName(supporting_surface_id);

    moveit_msgs::MoveItErrorCodes miec = move_group.pick(object_to_pick_id, grasps, plan_only);

    if (miec.val == miec.SUCCESS) {
        return true;
    }

    ROS_ERROR_STREAM("Error while executing pick task. MoveItErrorCode: " << miec.val);
    return false;
}
bool
GraspUtil::placeFromSide(moveit::planning_interface::MoveGroupInterface &move_group, geometry_msgs::Pose &place_pose,
                   double open_amount, std::string supporting_surface_id, std::string object_to_place_id, bool plan_only){

    std::vector<moveit_msgs::PlaceLocation> place_location;
    place_location.resize(1);

    place_location[0].place_pose.header.frame_id = "panda_link0";
    place_location[0].place_pose.pose.orientation.w = 1;
    place_location[0].place_pose.pose.position = place_pose.position;
    place_location[0].place_pose.pose.position.z = place_location[0].place_pose.pose.position.z;

    place_location[0].pre_place_approach.direction.header.frame_id = "panda_link0";
    place_location[0].pre_place_approach.direction.vector.z = -1.0;
    place_location[0].pre_place_approach.min_distance = 0.095;
    place_location[0].pre_place_approach.desired_distance = 0.115;

    place_location[0].post_place_retreat.direction.header.frame_id = "panda_link0";
    place_location[0].post_place_retreat.direction.vector.z = 1.0;
    place_location[0].post_place_retreat.min_distance = 0.1;
    place_location[0].post_place_retreat.desired_distance = 0.25;

    setupOpenGripper(place_location[0].post_place_posture, open_amount);
    move_group.setSupportSurfaceName(supporting_surface_id);

    moveit_msgs::MoveItErrorCodes miec = move_group.place(object_to_place_id, place_location, plan_only);

    if (miec.val == miec.SUCCESS) {
        return true;
    }

    ROS_ERROR_STREAM("Error while executing place task. MoveItErrorCode: " << miec.val);
    return false;
}

bool
GraspUtil::placeFromAbove(moveit::planning_interface::MoveGroupInterface &move_group, geometry_msgs::Pose &place_pose,
                          double open_amount, std::string supporting_surface_id, std::string object_to_place_id, bool plan_only) {

    std::vector<moveit_msgs::PlaceLocation> place_location;
    place_location.resize(1);

    place_location[0].place_pose.header.frame_id = "panda_link0";
    //  tf2::Quaternion orientation;
    //  orientation.setRPY(0, 0, M_PI / 2);
    place_location[0].place_pose.pose.orientation.w = 1;// = tf2::toMsg(orientation);
    place_location[0].place_pose.pose.position = place_pose.position;

    place_location[0].pre_place_approach.direction.header.frame_id = "panda_link0";
    place_location[0].pre_place_approach.direction.vector.z = -1.0;
    place_location[0].pre_place_approach.min_distance = 0.095;
    place_location[0].pre_place_approach.desired_distance = 0.115;

    place_location[0].post_place_retreat.direction.header.frame_id = "panda_link0";
    place_location[0].post_place_retreat.direction.vector.z = 1.0;
    place_location[0].post_place_retreat.min_distance = 0.1;
    place_location[0].post_place_retreat.desired_distance = 0.25;

    setupOpenGripper(place_location[0].post_place_posture, open_amount);
    move_group.setSupportSurfaceName(supporting_surface_id);

    moveit_msgs::MoveItErrorCodes miec = move_group.place(object_to_place_id, place_location, plan_only);

    if (miec.val == miec.SUCCESS) {
        return true;
    }

    ROS_ERROR_STREAM("Error while executing place task. MoveItErrorCode: " << miec.val);
    return false;
}

bool GraspUtil::homeGripper() {
    actionlib::SimpleActionClient<franka_gripper::MoveAction> moveActionClient("franka_gripper/move", true);
    if (moveActionClient.waitForServer(ros::Duration(1.0))) {
        moveActionClient.cancelAllGoals();
    }

    actionlib::SimpleActionClient<franka_gripper::HomingAction> ac("/franka_gripper/homing", true);
    if (!ac.waitForServer(ros::Duration(10.0))) {
        ROS_ERROR("Gripper homing failed (action server not available).");
        return false;
    }

    ac.cancelAllGoals();
    franka_gripper::HomingGoal goal;
    ac.sendGoal(goal);
    if (ac.waitForResult(ros::Duration(15.0))) {
        actionlib::SimpleClientGoalState state = ac.getState();
        ROS_INFO_STREAM("Homing finished: " <<  state.toString());
        return state == actionlib::SimpleClientGoalState::SUCCEEDED;
    } else {
        ROS_ERROR("Gripper homing failed (with timeout).");
        return false;
    }
}