#include "ros/ros.h"

#include "tf2_geometry_msgs/tf2_geometry_msgs.h"
#include "tf2/LinearMath/Quaternion.h"
#include "tf2/LinearMath/Vector3.h"

#include <moveit/planning_scene_interface/planning_scene_interface.h>
#include <moveit/move_group_interface/move_group_interface.h>

#include <gtest/gtest.h>
#include <grasp_util.h>

/**
 * Testing grasping from side.
 */
class Grasping : public ::testing::Test {

protected:

    GraspUtil grasp_util;

    geometry_msgs::Pose pick_table_pose;
    geometry_msgs::Pose place_table_pose;
    geometry_msgs::Pose obj_pose;
    geometry_msgs::Pose place_pose;

    tf2::Quaternion gripper_orientation;
    tf2::Quaternion place_orientation;

    geometry_msgs::Vector3 dimensions;

    void SetUp() override {

        ros::NodeHandle n("panda_grasping");

        gripper_orientation.setRPY(0, 0, -M_PI_4);
        place_pose.position.z = 0.4 + 0.05; // table height + object height / 2 + buffer
        place_orientation.setRPY(0, 0, M_PI / 2);
        dimensions.x = 0.05;
        dimensions.y = 0.05;
        dimensions.z = 0.1;

        CleanupScene();
    }

    static void CleanupScene() {
        moveit::planning_interface::PlanningSceneInterface planning_scene_interface;

        for (auto p : planning_scene_interface.getObjects()) {
            p.second.operation = moveit_msgs::CollisionObject::REMOVE;
            planning_scene_interface.applyCollisionObject(p.second);
        }
        ASSERT_EQ(planning_scene_interface.getObjects().size(), 0);
        ASSERT_EQ(planning_scene_interface.getAttachedObjects().size(), 0);
    }

    void doTest(geometry_msgs::Pose &pick_pose, geometry_msgs::Pose &pick_table_pose, geometry_msgs::Pose &place_table_pose,
            geometry_msgs::Pose &obj_pose, geometry_msgs::Pose &place_pose, geometry_msgs::Vector3 &approach_direction){

        moveit::planning_interface::PlanningSceneInterface planning_scene_interface;
        moveit::planning_interface::MoveGroupInterface group("panda_arm");

        std::vector<moveit_msgs::CollisionObject> collision_objects;
        collision_objects.resize(3);

        collision_objects[0].id = "table1";
        collision_objects[0].header.frame_id = "panda_link0";

        collision_objects[0].primitives.resize(1);
        collision_objects[0].primitives[0].type = collision_objects[0].primitives[0].BOX;
        collision_objects[0].primitives[0].dimensions.resize(3);
        collision_objects[0].primitives[0].dimensions[0] = 0.2;
        collision_objects[0].primitives[0].dimensions[1] = 0.2;
        collision_objects[0].primitives[0].dimensions[2] = 0.4;
        collision_objects[0].primitive_poses.resize(1);
        collision_objects[0].primitive_poses[0] = pick_table_pose;
        collision_objects[0].operation = collision_objects[0].ADD;

        collision_objects[1].id = "table2";
        collision_objects[1].header.frame_id = "panda_link0";

        collision_objects[1].primitives.resize(1);
        collision_objects[1].primitives[0].type = collision_objects[1].primitives[0].BOX;
        collision_objects[1].primitives[0].dimensions.resize(3);
        collision_objects[1].primitives[0].dimensions[0] = 0.2;
        collision_objects[1].primitives[0].dimensions[1] = 0.2;
        collision_objects[1].primitives[0].dimensions[2] = 0.4;
        collision_objects[1].primitive_poses.resize(1);
        collision_objects[1].primitive_poses[0] = place_table_pose;
        collision_objects[1].operation = collision_objects[1].ADD;

        collision_objects[2].header.frame_id = "panda_link0";
        collision_objects[2].id = "object";

        collision_objects[2].primitives.resize(1);
        collision_objects[2].primitives[0].type = collision_objects[1].primitives[0].BOX;
        collision_objects[2].primitives[0].dimensions.resize(3);
        collision_objects[2].primitives[0].dimensions[0] = dimensions.x;
        collision_objects[2].primitives[0].dimensions[1] = dimensions.y;
        collision_objects[2].primitives[0].dimensions[2] = dimensions.z;
        collision_objects[2].primitive_poses.resize(1);
        collision_objects[2].primitive_poses[0] = obj_pose;
        collision_objects[2].operation = collision_objects[2].ADD;

        planning_scene_interface.applyCollisionObjects(collision_objects);

        pickAndPlace(pick_pose, planning_scene_interface, group, place_pose, approach_direction);
    }

    void pickAndPlace(geometry_msgs::Pose &pick_pose, moveit::planning_interface::PlanningSceneInterface &planning_scene_interface,
            moveit::planning_interface::MoveGroupInterface &group, geometry_msgs::Pose &place_pose, geometry_msgs::Vector3 &approach_direction) {

        shape_msgs::SolidPrimitive object;
        object.type = shape_msgs::SolidPrimitive::CYLINDER;
        object.dimensions.resize(2);
        object.dimensions[0] = dimensions.x;
        object.dimensions[1] = dimensions.y;

        GraspUtil grasp_util;


        ros::WallDuration(1.0).sleep();

        ASSERT_TRUE(grasp_util.pickFromSide(group, pick_pose, object, approach_direction,0.06, "table1", "object"));

        ros::WallDuration(1.0).sleep();

        ASSERT_TRUE(grasp_util.placeFromSide(group, place_pose, 0.06, "table2", "object", false));
    }
};

/**
 * test from -x grasp
*/
TEST_F(Grasping, SIDE_GRASPING_1) {

    pick_table_pose.position.x = 0.5;
    pick_table_pose.position.y = -0.5;
    pick_table_pose.position.z = 0.2;
    pick_table_pose.orientation.w = 1.0;

    place_table_pose.position.x = 0;
    place_table_pose.position.y = 0.5;
    place_table_pose.position.z = 0.2;
    place_table_pose.orientation.w = 1.0;

    obj_pose.position.x = 0.5;
    obj_pose.position.y = -0.5;
    obj_pose.position.z = 0.45;
    obj_pose.orientation.w = 1.0;


    place_pose.orientation = tf2::toMsg(place_orientation);
    place_pose.position.x = place_table_pose.position.x;
    place_pose.position.y = place_table_pose.position.y;

    geometry_msgs::Vector3 approach_direction;
    approach_direction.x = 1;

    geometry_msgs::Pose pick_pose = grasp_util.getPickFromBack(obj_pose, dimensions, tf2::toMsg(gripper_orientation));
    doTest(pick_pose, pick_table_pose, place_table_pose, obj_pose, place_pose, approach_direction);
}

/**
 * test from +x grasp
*/
TEST_F(Grasping, SIDE_GRASPING_2) {

    pick_table_pose.position.x = -0.5;
    pick_table_pose.position.y = -0.5;
    pick_table_pose.position.z = 0.2;
    pick_table_pose.orientation.w = 1.0;

    place_table_pose.position.x = 0;
    place_table_pose.position.y = 0.5;
    place_table_pose.position.z = 0.2;
    place_table_pose.orientation.w = 1.0;

    obj_pose.position.x = -0.5;
    obj_pose.position.y = -0.5;
    obj_pose.position.z = 0.45;
    obj_pose.orientation.w = 1.0;

    place_pose.orientation = tf2::toMsg(place_orientation);
    place_pose.position.x = place_table_pose.position.x;
    place_pose.position.y = place_table_pose.position.y;

    geometry_msgs::Vector3 approach_direction;
    approach_direction.x = -1;

    geometry_msgs::Pose pick_pose = grasp_util.getPickFromFront(obj_pose, dimensions, tf2::toMsg(gripper_orientation));
    doTest(pick_pose, pick_table_pose, place_table_pose, obj_pose, place_pose, approach_direction);
}

/**
 * test from -y grasp
*/
TEST_F(Grasping, SIDE_GRASPING_3) {

    pick_table_pose.position.x = 0.5;
    pick_table_pose.position.y = 0.5;
    pick_table_pose.position.z = 0.2;
    pick_table_pose.orientation.w = 1.0;

    place_table_pose.position.x = -0.5;
    place_table_pose.position.y = 0;
    place_table_pose.position.z = 0.2;
    place_table_pose.orientation.w = 1.0;

    obj_pose.position.x = 0.5;
    obj_pose.position.y = 0.5;
    obj_pose.position.z = 0.45;
    obj_pose.orientation.w = 1.0;

    place_pose.orientation = tf2::toMsg(place_orientation);
    place_pose.position.x = place_table_pose.position.x;
    place_pose.position.y = place_table_pose.position.y;

    geometry_msgs::Vector3 approach_direction;
    approach_direction.y = +1;

    geometry_msgs::Pose pick_pose = grasp_util.getPickFromLeft(obj_pose, dimensions, tf2::toMsg(gripper_orientation));
    doTest(pick_pose, pick_table_pose, place_table_pose, obj_pose, place_pose, approach_direction);
}

/**
 * test from +y grasp
 */
TEST_F(Grasping, SIDE_GRASPING_4) {

    pick_table_pose.position.x = 0.5;
    pick_table_pose.position.y = -0.5;
    pick_table_pose.position.z = 0.2;
    pick_table_pose.orientation.w = 1.0;

    place_table_pose.position.x = -0.5;
    place_table_pose.position.y = 0;
    place_table_pose.position.z = 0.2;
    place_table_pose.orientation.w = 1.0;

    obj_pose.position.x = 0.5;
    obj_pose.position.y = -0.5;
    obj_pose.position.z = 0.45;
    obj_pose.orientation.w = 1.0;

    place_pose.orientation = tf2::toMsg(place_orientation);
    place_pose.position.x = place_table_pose.position.x;
    place_pose.position.y = place_table_pose.position.y;

    geometry_msgs::Vector3 approach_direction;
    approach_direction.y = -1;

    geometry_msgs::Pose pick_pose = grasp_util.getPickFromRight(obj_pose, dimensions, tf2::toMsg(gripper_orientation));
    doTest(pick_pose, pick_table_pose, place_table_pose, obj_pose, place_pose, approach_direction);
}

int main(int argc, char **argv) {
    testing::InitGoogleTest(&argc, argv);
    ros::init(argc, argv, "tests_side_grasping");
    ros::NodeHandle n("panda_grasping");

    ros::AsyncSpinner spinner(4);
    spinner.start();

    ROS_INFO_NAMED("constraint_planner", ">>>>>>>>>>>>>>>> WAITING FOR ROBOT INIT <<<<<<<<<<<<<<<");
    ros::Duration(3.0).sleep();
    ROS_INFO_NAMED("constraint_planner", ">>>>>>>>>>>>>>>>> STARTING TEST AFTER INIT <<<<<<<<<<<<<<<<");


    return RUN_ALL_TESTS();
}