#!/usr/bin/env python
# coding: utf-8

# # Pick up test
# This example will show a simple pickup test in addition to the TF Broadcaster to visualize the result in rviz. 
# This should be usable inside a binderhub using webviz
# 
# 

# ## Bulletworld

# In[ ]:


from rospy import get_param
from pycram.bullet_world import BulletWorld, Object
from pycram.pose import Pose

from pycram.process_module import simulated_robot, with_simulated_robot
from pycram.language import macros, par
from pycram.designators.location_designator import *
from pycram.designators.action_designator import *
from pycram.enums import Arms
from pycram.designators.object_designator import *
from pycram.designators.object_designator import BelieveObject
from pycram.plan_failures import IKError
import math

try:
    robot = get_param('/nbparam_robot')
    environment = get_param('/nbparam_environment')
except Exception as e:
    robot = 'pr2'
    environment = 'kitchen'
    
print(f"Robot: {robot}")
print(f"Environment: {environment}")

# With Bulletworld
world = BulletWorld()

# Without Bulletworld (use rviz instead)
world = BulletWorld('DIRECT')


# ## Robot

# In[ ]:


robot = Object("pr2", "robot", robot + ".urdf")
robot_desig = ObjectDesignatorDescription(names=["pr2"]).resolve()


# ## Create Objects

# In[ ]:


milk = Object("milk", "milk", "milk.stl", pose=Pose([2, 0, 1]))
milk_BO = BelieveObject(names=["milk"])


# ## TF Broadcaster

# In[ ]:


from pycram.ros.tf_broadcaster import TFBroadcaster
from pycram.ros.joint_state_publisher import JointStatePublisher
from pycram.ros.viz_marker_publisher import VizMarkerPublisher

broadcaster = TFBroadcaster()
joint_publisher = JointStatePublisher("joint_states", 0.1)
#v = VizMarkerPublisher(topic_name='viz_marker')


# ## Navigate to Object

# In[ ]:


with simulated_robot:
    ParkArmsAction([Arms.BOTH]).resolve().perform()

    MoveTorsoAction([0.33]).resolve().perform()

    pickup_pose_knife = CostmapLocation(target=milk_BO.resolve(), reachable_for=robot_desig).resolve()
    pickup_arm = pickup_pose_knife.reachable_arms[0]

    NavigateAction(target_locations=[pickup_pose_knife.pose]).resolve().perform()


# ## Perform the Pickup action

# In[ ]:


with simulated_robot:
    PickUpAction(object_designator_description=milk_BO,
                 arms=["left"],
                 grasps=["left", "right"]).resolve().perform()