############################################################################################
#                                                                                          #
# ISMR 2025 Workshop on Open-Source Software for Intelligent Image-Guided Medical Robotics #
# May 14, 2025 – Atlanta GA                                                                #
# Session 3 – Part 1: MRI-Guided Robot-Assisted Prostate Biopsy with SmartTemplate         #
#                                                                                          #
# mcostabernardesmatias@bwh.harvard.edu; tokuda@bwh.harvard.edu                            #
#                                                                                          #
############################################################################################



############################################
#                                          #
# STEP 3: Register SmartTemplate to world  #
#                                          #
############################################

# Get robot node
rosNode = slicer.util.getModuleLogic('ROS2').GetDefaultROS2Node()
robotNode = rosNode.GetRobotNodeByName('smart_template')

# Get robot_description
robot_description = slicer.util.getNode(robotNode.GetNodeReferenceID('parameter')).GetParameterAsString('robot_description')

# Get RobotToZFrame (from URDF robot_description)
import xml.etree.ElementTree as ET
root = ET.fromstring(robot_description)
zframe_orientation = root.find('./custom_parameters/zframe_orientation').get('value').strip()
zframe_position = root.find('./custom_parameters/zframe_position').get('value').strip()
rotation_values = list(map(float, zframe_orientation.strip().split()))
translation_values = list(map(float, zframe_position.strip().split()))
mat_ZFrameToRobot = vtk.vtkMatrix4x4()
for i in range(3): # Set rotation part (3x3)
    for j in range(3):
        mat_ZFrameToRobot.SetElement(i, j, rotation_values[3 * i + j])
for i in range(3): # Set translation part (last column)
    mat_ZFrameToRobot.SetElement(i, 3, 1000*translation_values[i])
mat_ZFrameToRobot.SetElement(3, 3, 1.0)
t_RobotToZFrame = vtk.vtkTransform()
t_RobotToZFrame.SetMatrix(mat_ZFrameToRobot)
t_RobotToZFrame.Inverse()

# Get ZFrameToScanner (from ZFrame registration)
mat_ZFrameToScanner = vtk.vtkMatrix4x4()
ZFrameToScannerNode = slicer.util.getFirstNodeByName('COR_TSE_T2_COVER_ZFRAME-label-ZFrameTransform', className='vtkMRMLLinearTransformNode')
ZFrameToScannerNode.GetMatrixTransformToWorld(mat_ZFrameToScanner)
t_ZFrameToScanner = vtk.vtkTransform()
t_ZFrameToScanner.SetMatrix(mat_ZFrameToScanner)

# Calculate RobotToScanner transform (robot world pose)
t_RobotToScanner = vtk.vtkTransform()
t_RobotToScanner.Concatenate(t_ZFrameToScanner)
t_RobotToScanner.Concatenate(t_RobotToZFrame)
mat_RobotToScanner = t_RobotToScanner.GetMatrix()

# Create publisher and send RobotToScanner to \world_listener
pubWorld = rosNode.CreateAndAddPublisherNode('TransformStamped', '/world_pose')
world_msg = pubWorld.GetBlankMessage()
world_msg.SetTransform(mat_RobotToScanner)
pubWorld.Publish(world_msg)

############################################
#                                          #
# STEP 5: Targeted insertion w/ publisher  #
#                                          #
############################################

# Create publisher for /desired_position
pubGoal = rosNode.CreateAndAddPublisherNode('PoseStamped', '/desired_position')

# Get target in scanner coordinates
targetMarkupsNode = slicer.util.getFirstNodeByName('F', className='vtkMRMLMarkupsFiducialNode')
target_scanner = targetMarkupsNode.GetNthControlPointPosition(0)

# Calculate initial alignment (insertion depth [Y coordinate] set to 0.0)
# OBS: Another approach would be to use a tf_lookup from 'world' to 'base_link' to get mat_RobotToScanner 
t_ScannerToRobot = vtk.vtkTransform()
t_ScannerToRobot.SetMatrix(mat_RobotToScanner)
t_ScannerToRobot.Inverse()
target = t_ScannerToRobot.TransformPoint(target_scanner)

mat_desired_position = vtk.vtkMatrix4x4()
mat_desired_position.Identity()
mat_desired_position.SetElement(0,3,target[0])
mat_desired_position.SetElement(2,3,target[2])

#### PHASE 1: ALIGNMENT ####
# Send alignment position message
goal_msg = pubGoal.GetBlankMessage()
h = goal_msg.GetHeader()
h.SetFrameId('base_link')
goal_msg.SetPose(mat_desired_position)
pubGoal.Publish(goal_msg)
print('Desired position = [%.4f, %.4f, %.4f]mm - base_link frame' % (target[0], 0.0, target[2]))

#__________________________________________#
#  WAIT for the robot to finish alignment  #
#__________________________________________#


#### PHASE 2: INSERTION ####
# Send target position message
mat_desired_position.SetElement(1,3,target[1])
goal_msg.SetPose(mat_desired_position)
pubGoal.Publish(goal_msg)
print('Desired position = [%.4f, %.4f, %.4f]mm - base_link frame' % (target[0], target[1], target[2]))

############################################
#                                          #
#  STEP 6: Get needle pose w/ subscriber   #
#                                          #
############################################

# Create subscriber for /end_effector_pose and print last message
subEEPose = rosNode.CreateAndAddSubscriberNode('PoseStamped', '/end_effector_pose')
print(subEEPose.GetLastMessage())


############################################
#                                          #
#    STEP 7: Command robot w/ publisher    #
#                                          #
############################################

# Create publisher for /desired_command
pubCommand = rosNode.CreateAndAddPublisherNode('String', '/desired_command')

# Send command message
pubCommand.Publish('RETRACT')

#__________________________________________#
# WAIT for the robot to finish retraction  #
#__________________________________________#

# Send command message
pubCommand.Publish('HOME')