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

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

To use, you should have created your own moveit_config stack using the setup_assistant e.g. <your_robot_name>_moveit_config
Use IKFast to generate a kinematics solution for your robot.
Copy the output files to a separate <your_robot>_moveit_plugins package.
Then run:
$ create_ikfast_moveit_plugin.py <your_robot_name> <moveit_plugin_pkg>

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

Tested on ROS Groovy
with code generated by IKFast61 from OpenRAVE 0.8.2
using a 6 DOF manipulator

Author: Jeremy Zoss, SwRI
        Based heavily on the arm_navigation plugin generator by David Butterworth, KAIST
Date: January 2013

'''
'''
Copyright (c) 2013, Jeremy Zoss, SwRI
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 glob
import sys
import roslib
import re
import os
import yaml
from lxml import etree

plugin_gen_pkg = 'arm_kinematics_tools'  # package containing this file

if __name__ == '__main__':
   # Check input arguments
   try:
      robot_name = sys.argv[1]
      moveit_plugin_pkg = sys.argv[2]
      assert( len(sys.argv) < 4 )   # invalid num-arguments
   except:
       print("\nUsage: create_ikfast_plugin.py <your_robot_name> <moveit_plugin_pkg>\n")
       sys.exit(-1)
   print '\nIKFast Plugin Generator'

   # Setup key package directories
   try: 
      plan_pkg = robot_name + '_moveit_config'
      plan_pkg_dir = roslib.packages.get_pkg_dir(plan_pkg)
      print 'Loading robot from \''+plan_pkg+'\' package ... '
   except:
      print '\nERROR: can\'t find package '+plan_pkg+'\n'
      sys.exit(-1)
   try:
      plugin_pkg = moveit_plugin_pkg
      plugin_pkg_dir = roslib.packages.get_pkg_dir(plugin_pkg)
      print 'Creating plugin in \''+plugin_pkg+'\' package ... '
   except:
      print '\nERROR: can\'t find package '+plugin_pkg+'\n'
      sys.exit(-1)

   # Check for at least 1 planning group, should be called 'manipulator' by ROS specifications
   try:
      srdf_files = glob.glob(plan_pkg_dir+'/config/*.srdf')
      if (len(srdf_files) == 1):
         srdf_file_name = srdf_files[0]
      else:
         srdf_file_name = plan_pkg_dir + '/config/' + robot_name + '.srdf'
      srdf = etree.parse(srdf_file_name).getroot()
      if (robot_name != srdf.get('name')):
          print '\nERROR: non-matching robot name found in ' \
                + srdf_file_name + '.' \
                + '  Expected \'' + robot_name + '\',' \
                + ' found \''+srdf.get('name')+'\''
          raise
      
      groups = srdf.findall('group')
      if(len(groups) < 1) : # No groups
         raise
      if groups[0].get('name') == None: # Group name is blank
         raise
   except:
      print("\nERROR: need at least 1 planning group in robot planning description ")
      print srdf_file_name + '\n'
      sys.exit(-1)
   print '  found ' + str(len(groups)) + ' planning groups: ' \
         + ", ".join([g.get('name') for g in groups])

   # Select manipulator arm group
   arm_group = None
   for g in groups:
      foundName  = (g.get('name').lower() == 'manipulator2')
      foundChain = (g.find('chain') != None)
      found6DOF  = ( len(g.findall('joint')) == 6 )

      if (foundName or foundChain or found6DOF):
         arm_group = g
         group_name = g.get('name')
         break
   if arm_group is None:
      print '\nERROR: could not identify manipulator group in SRDF.\n'
      sys.exit(-1)
   print '  found group \'' + arm_group.get('name') + '\''

   # Check for source code generated by IKFast
   solver_file_name = plugin_pkg_dir+'/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 \'' + \
            solver_file_name + '\'\n'
      print 'Copy the source file generated by IKFast to this location \n'
      sys.exit(-1)

   # Detect version of IKFast used to generate solver code
   solver_version = 0
   with open(solver_file_name,'r') as src:
      for line in src:
         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)

   # For IKFast56+, check for header file
   if solver_version >= 56:
      header_file_name = plugin_pkg_dir+'/include/ikfast.h'
      if not os.path.exists(header_file_name):
         print '\nERROR: can\'t find IKFast header file at \'' + \
               header_file_name + '\'\n'
         print 'Copy the IKFast header file 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_gen_dir = roslib.packages.get_pkg_dir(plugin_gen_pkg)
   except:
      print '\nERROR: can\'t find package '+plugin_gen_pkg+' \n'
      sys.exit(-1)
   template_file_name = plugin_gen_dir + '/templates/ikfast' + \
                        str(template_version) + \
                        '_moveit_plugin_template.cxx'
   if not os.path.exists(template_file_name):
      print '\nERROR: can\'t find template file at \'' + template_file_name + '\'\n'
      sys.exit(-1)

   # Create plugin source from template
   template_file_data = open(template_file_name, 'r')
   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)
   plugin_file_base = robot_name + '_' + group_name + \
                      '_ikfast_moveit_plugin.cpp'
   plugin_file_name = plugin_pkg_dir + '/src/' + plugin_file_base
   with open(plugin_file_name,'w') as f:
      f.write(template_text)
   print '\nCreated plugin file at \'' + plugin_file_name + '\''

   # Create plugin definition
   lib_file_base = robot_name + "_moveit_arm_kinematics"
   plugin_name = robot_name + '_' + group_name + \
                 '_kinematics/IKFastKinematicsPlugin'
   plugin_def = etree.Element("library", path="lib/lib"+lib_file_base)
   cl = etree.SubElement(plugin_def, "class")
   cl.set("name", plugin_name)
   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'

   # Write plugin definition to file
   def_file_base = robot_name + "_moveit_kinematics_plugin.xml"
   def_file_name = plugin_pkg_dir + "/" + def_file_base
   with open(def_file_name,'w') as f:
      etree.ElementTree(plugin_def).write(f, xml_declaration=True, pretty_print=True)
   print '\nCreated plugin definition at: \''+def_file_name+'\''

   # Modify CMakelists
   with open(plugin_pkg_dir+"/CMakeLists.txt", 'a+') as cmake_file:
      cmake_file_text = cmake_file.read()
      if re.search(lib_file_base, cmake_file_text) is None:
        extra_text = '\n' + \
                     '# Library containing IKFast plugin\n' + \
                     'rosbuild_add_library(' + lib_file_base + '\n' + \
                     '  src/' + plugin_file_base + '\n' + \
                     '  )'
        cmake_file.write(extra_text)
        print 'Modified '+plugin_pkg_dir+'/CMakeLists.txt'

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

   # Check that moveit_core is in the depends list
   found = False
   for depend_entry in manifest_xml.getroot().findall("depend"):
        if depend_entry.get("package") == "moveit_core":
            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="moveit_core"))
                break

   # Check that plugin definition file is in the export list
   new_export = etree.Element("moveit_core", \
                              plugin="${prefix}/"+def_file_base)
   export_element = manifest_xml.getroot().find("export")
   if export_element == None:
        export_element = etree.SubElement(manifest_xml.getroot(), "export")
   found = False
   for el in export_element.findall("moveit_core"):
      found = (etree.tostring(new_export) == etree.tostring(el))
      if found: break

   if not found:
      export_element.append(new_export)
      manifest_file_name = plugin_pkg_dir+"/manifest.xml"
      with open(manifest_file_name,"w") as f:
         manifest_xml.write(f, xml_declaration=True, pretty_print=True)
      print '\nModified manifest.xml at \''+manifest_file_name+'\''

   # Modify kinematics.yaml file
   kin_yaml_file_name = plan_pkg_dir+"/config/kinematics.yaml"
   with open(kin_yaml_file_name, 'r') as f:
      kin_yaml_data = yaml.safe_load(f)
   kin_yaml_data["manipulator"]["kinematics_solver"] = plugin_name
   with open(kin_yaml_file_name, 'w') as f:
      yaml.dump(kin_yaml_data, f,default_flow_style=False)
   print '\nModified kinematics.yaml at ' + kin_yaml_file_name

