#!/usr/bin/env python

from __future__ import print_function

import argparse
import os

from pybullet_tools.pr2_problems import create_floor, create_table
from pybullet_tools.pr2_utils import get_top_grasps
from pybullet_tools.utils import (
    BLUE,
    GREEN,
    RED,
    TAN,
    WSG_50_URDF,
    HideOutput,
    LockRenderer,
    Point,
    add_line,
    connect,
    create_box,
    create_cylinder,
    disconnect,
    draw_pose,
    get_max_limit,
    get_model_path,
    get_movable_joints,
    get_pose,
    get_relative_pose,
    invert,
    link_from_name,
    load_pybullet,
    multiply,
    set_camera_pose,
    set_joint_position,
    set_point,
    set_pose,
    stable_z,
    unit_pose,
    wait_for_user,
)

# from pddlstream.utils import get_file_path

# https://www.generationrobots.com/en/403318-fe-gripper-for-panda-robotic-arm.html
# /usr/local/lib/python2.7/site-packages/pybullet_data/gripper/
# drake/manipulation/models/ycb/sdf/005_tomato_soup_can.sdf


def close_gripper(robot):
    for joint in get_movable_joints(robot):
        set_joint_position(robot, joint, get_max_limit(robot, joint))


def open_gripper(robot):
    for joint in get_movable_joints(robot):
        set_joint_position(robot, joint, get_max_limit(robot, joint))


BLOCK_SIDE = 0.07
TABLE_WIDTH = 1.0
EPSILON = 1e-3

DRAKE_PATH = "/Users/caelan/Programs/external/drake"
DRAKE_YCB = "manipulation/models/ycb/sdf/003_cracker_box.sdf"

YCB_PATH = "/Users/caelan/Programs/external/ycb_benchmarks/16k_laser_scan"
YCB_TEMPLATE = "{}/google_16k/textured.obj"


def main():
    parser = argparse.ArgumentParser()  # Automatically includes help
    parser.add_argument("-viewer", action="store_true", help="enable viewer.")
    args = parser.parse_args()

    connect(use_gui=True)
    # ycb_path = os.path.join(DRAKE_PATH, DRAKE_YCB)
    # ycb_path = os.path.join(YCB_PATH, YCB_TEMPLATE.format('003_cracker_box'))
    # print(ycb_path)
    # load_pybullet(ycb_path)

    with LockRenderer():
        draw_pose(unit_pose(), length=1, width=1)
        floor = create_floor()
        set_point(floor, Point(z=-EPSILON))
        table = create_table(
            width=TABLE_WIDTH,
            length=TABLE_WIDTH / 2,
            height=TABLE_WIDTH / 2,
            top_color=TAN,
            cylinder=False,
        )
        # set_euler(table, Euler(yaw=np.pi/2))
        with HideOutput(False):
            # data_path = add_data_path()
            # robot_path = os.path.join(data_path, WSG_GRIPPER)
            robot_path = get_model_path(WSG_50_URDF)  # WSG_50_URDF | PANDA_HAND_URDF
            # robot_path = get_file_path(__file__, 'mit_arch_suction_gripper/mit_arch_suction_gripper.urdf')
            robot = load_pybullet(robot_path, fixed_base=True)
            # dump_body(robot)
            # robot = create_cylinder(radius=0.5*BLOCK_SIDE, height=4*BLOCK_SIDE) # vacuum gripper

        block1 = create_box(w=BLOCK_SIDE, l=BLOCK_SIDE, h=BLOCK_SIDE, color=RED)
        block_z = stable_z(block1, table)
        set_point(block1, Point(z=block_z))

        block2 = create_box(w=BLOCK_SIDE, l=BLOCK_SIDE, h=BLOCK_SIDE, color=GREEN)
        set_point(block2, Point(x=+0.25, z=block_z))

        block3 = create_box(w=BLOCK_SIDE, l=BLOCK_SIDE, h=BLOCK_SIDE, color=BLUE)
        set_point(block3, Point(x=-0.15, z=block_z))

        blocks = [block1, block2, block3]

        add_line(
            Point(x=-TABLE_WIDTH / 2, z=block_z - BLOCK_SIDE / 2 + EPSILON),
            Point(x=+TABLE_WIDTH / 2, z=block_z - BLOCK_SIDE / 2 + EPSILON),
            color=RED,
        )
        set_camera_pose(
            camera_point=Point(y=-1, z=block_z + 1), target_point=Point(z=block_z)
        )

    wait_for_user()
    block_pose = get_pose(block1)
    open_gripper(robot)
    tool_link = link_from_name(robot, "tool_link")
    base_from_tool = get_relative_pose(robot, tool_link)
    # draw_pose(unit_pose(), parent=robot, parent_link=tool_link)

    y_grasp, x_grasp = get_top_grasps(
        block1, tool_pose=unit_pose(), grasp_length=0.03, under=False
    )
    grasp = y_grasp  # fingers won't collide
    gripper_pose = multiply(block_pose, invert(grasp))
    set_pose(robot, multiply(gripper_pose, invert(base_from_tool)))
    wait_for_user("Finish?")
    disconnect()


if __name__ == "__main__":
    main()
