#! /usr/bin/env python
'''
IKFast Plugin Generator for arm_kinematics_constraint_aware

Creates a kinematics plugin using the output of IKFast from OpenRAVE.
This plugin and the constraint aware kinematics node can be used as a general 
kinematics service, from within the arm_navigation planning warehouse, or in 
your own ROS node.

To use, you should have created your own arm_navigation stack using the planning
wizard e.g. <your_robot_name>_arm_navigation
Use IKFast to generate a kinematics solution for your robot.
Copy the output files to your arm_navigation stack.
Then run:
$ create_ikfast_plugin.py <your_robot_name>

A full tutorial is available at:
http://www.ros.org/wiki/Industrial/Tutorials/Create_a_Fast_IK_Solution


Tested on ROS Fuerte
with code generated by IKFast54 from OpenRAVE 0.6.0
and IKFast 56/61 from OpenRAVE 0.8.2
using a 6 DOF manipulator

Note: you should apply the patch for arm_kinematics_constraint_aware 
mentioned here: https://code.ros.org/trac/ros-pkg/ticket/5586
This fixes a bug in the FK routine, and also allows you to switch from using
TF by default for forward kinematics, to the FK routine from this plugin.


Author: David Butterworth, KAIST
        Based on code by an unknown author
Date: November 2012

'''
'''
Copyright (c) 2012, David Butterworth, KAIST
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of the Willow Garage, Inc. nor the names of its
      contributors may be used to endorse or promote products derived from
      this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
IABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
POSSIBILITY OF SUCH DAMAGE.
'''

import sys
import yaml
import roslib
import re
import os
from lxml import etree

plugin_generator_package = 'arm_kinematics_tools'

if __name__ == '__main__':
   # Expect robot name as argument
   if(len(sys.argv) != 2):
       print("\nUsage: create_ikfast_plugin.py <your_robot_name> \n")
       sys.exit(-1)
   robot_name = sys.argv[1]
  
   print '\nIKFast Plugin Generator'

   # Create plugin inside <your_robot_name>_arm_navigation stack
   try: 
      directory_name = roslib.packages.get_pkg_dir(robot_name+'_arm_navigation')
   except:
      print '\nERROR: can\'t find package '+robot_name+'_arm_navigation \n'
      sys.exit(-1)
   print '\nCreating plugin in \''+robot_name+'_arm_navigation stack\'... '

   # Check for at least 1 planning group, should be called 'manipulator' by ROS specifications
   try: 
      yaml_file_name = directory_name + '/config/' + robot_name + '_planning_description.yaml'
      data = yaml.safe_load(open(yaml_file_name, 'r')) # Dictionary of yaml defs
      groups = data['groups'] # List of groups
      if(len(groups) < 1) : # No groups
         raise
      if groups[0]['name'] == None: # Group name is blank
         raise
   except:
      print("\nERROR: need at least 1 planning group in robot planning description ")
      print directory_name + '/config/' + robot_name + '_planning_description.yaml \n'
      sys.exit(-1)
   print '\nFound ' + str(len(groups)) + ' planning groups:  '

   # Check for source code generated by IKFast
   for i in groups:
        if 'tip_link' in i:
            group_name = i['name']
            print '  '+group_name
            solver_file_name = directory_name+'/src/'+robot_name+'_'+group_name+'_ikfast_solver.cpp'
            if not os.path.exists(solver_file_name):
               print '\nERROR: can\'t find IKFast source code at '
               print directory_name+'/src/'+robot_name+'_'+group_name+'_ikfast_solver.cpp \n'
               print 'Copy the source file generated by IKFast to this location \n'
               sys.exit(-1)
   print ' '

   # Detect version of IKFast used to generate solver code
   group_name = groups[0]['name']
   ikfast_source = open(directory_name+'/src/'+robot_name+'_'+group_name+'_ikfast_solver.cpp','r')
   lines = ikfast_source.readlines()
   ikfast_source.close()
   solver_version = 0
   for line in lines:
      if line.startswith('/// ikfast version'):
         line_search = re.search('ikfast version (.*) generated', line)
         if line_search:
            solver_version = int(line_search.group(1))
         break
   print 'Found source code generated by IKFast version ' + str(solver_version) + '\n'

   # For IKFast56+, check for header file
   if solver_version >= 56:
      header_file_name = directory_name+'/include/ikfast.h'
      if not os.path.exists(solver_file_name):
         print '\nERROR: can\'t find IKFast header file at '
         print directory_name+'/include/ikfast.h \n'
         print 'Copy the header file generated by IKFast to this location \n'
         sys.exit(-1)

   # Chose template depending on IKFast version
   if solver_version >= 56:
      template_version = 61
   else:
      template_version = 54

   # Check if template exists
   try: 
      plugin_generator_directory = roslib.packages.get_pkg_dir(plugin_generator_package)
   except:
      print '\nERROR: can\'t find package '+plugin_generator_package+' \n'
      sys.exit(-1)
   template_file_name = plugin_generator_directory+'/templates/ikfast'+str(template_version)+'_plugin_template.cxx'
   if not os.path.exists(template_file_name):
      print '\nERROR: can\'t find template file at '
      print template_file_name+' \n'
      sys.exit(-1)

   # Create plugin source from template
   for i in groups:
        if 'tip_link' in i:
            template_file_data = open(template_file_name, 'r')
            group_name = i['name']
            template_text = template_file_data.read()
            template_text = re.sub('_ROBOT_NAME_', robot_name, template_text)
            template_text = re.sub('_GROUP_NAME_', group_name, template_text)
            if not os.path.exists(directory_name+'/src'):
                os.makedirs(directory_name+'/src')
            output_template = open(directory_name+'/src/'+robot_name+"_"+group_name+"_ikfast_plugin.cpp",'w')
            output_template.write(template_text)
            print 'Created '+directory_name+'/src/'+robot_name+'_'+group_name+'_ikfast_plugin.cpp'
                       
   # Create plugin definition
   plugin_file = open(directory_name+"/kinematics_plugins.xml",'w')
   root_one = etree.Element("library", path="lib/lib"+robot_name+"_kinematics_lib")
   for i in groups:
        if 'tip_link' in i:
            group_name = i['name']
            cl = etree.SubElement(root_one, "class")
            cl.set("name", robot_name+'_'+group_name+'_kinematics/IKFastKinematicsPlugin')
            cl.set("type", robot_name+'_'+group_name+'_kinematics::IKFastKinematicsPlugin')
            cl.set("base_class_type", "kinematics::KinematicsBase")
            desc = etree.SubElement(cl, "description")
            desc.text = 'IKFast'+str(template_version)+' plugin for closed-form kinematics'
   etree.ElementTree(root_one).write(plugin_file, xml_declaration=True, pretty_print=True)
   print 'Created '+directory_name+'/kinematics_plugins.xml'

   # Modify CMakelists
   cmake_file = open(directory_name+"/CMakeLists.txt", 'a+')
   cmake_file_text = cmake_file.read()
   if re.search(robot_name+'_kinematics_lib', cmake_file_text) == None:
        extra_text = '\n# Library containing IKFast plugin \nrosbuild_add_library('+robot_name+'_kinematics_lib\n'
        for i in groups:
            if 'tip_link' in i:
                group_name = i['name']
                extra_text += '  src/'+robot_name+"_"+group_name+'_ikfast_plugin.cpp\n'
        extra_text += "  )"     
        cmake_file.write(extra_text)
        print 'Modified '+directory_name+'/CMakeLists.txt'
   cmake_file.close()

   parser = etree.XMLParser(remove_blank_text=True)

   # Add plugin export to package manifest
   manifest_xml = etree.parse(directory_name+"/manifest.xml", parser)

   # Check that kinematics_base is in the depends list
   found = False
   for depend_entry in manifest_xml.getroot().findall("depend"):
        if depend_entry.get("package") == "kinematics_base":
            found = True
            break
   if not found:
        for idx, entry in enumerate(manifest_xml.getroot().getchildren()):
            if entry.tag == "depend":
                manifest_xml.getroot().insert(idx, etree.Element("depend", package="kinematics_base"))
                break

   export_element = manifest_xml.getroot().find("export")
   if export_element == None:
        export_element = etree.SubElement(manifest_xml.getroot(), "export")

   kinematics_element = export_element.find("kinematics_base")
   if kinematics_element == None:
        kinematics_element = etree.SubElement(export_element, "kinematics_base", plugin="${prefix}/kinematics_plugins.xml")

   manifest_xml_file = open(directory_name+"/manifest.xml", "w")
   manifest_xml.write(manifest_xml_file, xml_declaration=True, pretty_print=True)
   print 'Modified '+directory_name+'/manifest.xml'

   # Modify constraint_aware kinematics launch file
   constraint_aware_xml = etree.parse(directory_name+"/launch/constraint_aware_kinematics.launch", parser)
   for node_entry in constraint_aware_xml.getroot().findall("node"):
        for param_entry in node_entry.getchildren():
            if param_entry.get("name") == "group":
                group_name = param_entry.get("value")
        # modify solver
        for param_entry in node_entry.getchildren():
            if param_entry.get("name") == "kinematics_solver":
                old_solver = param_entry.get("value")
                param_entry.set("value", robot_name+"_"+group_name+"_kinematics/IKFastKinematicsPlugin")
        # disable old solver
        param = etree.SubElement(node_entry, "param")
        param.set('name', 'OLDkinematics_solver')
        param.set('type', 'string')
        param.set('value', old_solver)
   constraint_aware_xml_file = open(directory_name+"/launch/constraint_aware_kinematics.launch", "w")
   constraint_aware_xml.write(constraint_aware_xml_file, xml_declaration=True, pretty_print=True)
   print 'Modified '+directory_name+'/launch/constraint_aware_kinematics.launch \n'

