
# \file orchomp.py
# \brief Python interface to orchomp.
# \author Temple Price
# \date 2012

# (C) Copyright 2012 Carnegie Mellon University

import types
import openravepy

def bind(mod):
   mod.viewspheres = types.MethodType(viewspheres,mod)
   mod.computedistancefield = types.MethodType(computedistancefield,mod)
   mod.addfield_fromobsarray = types.MethodType(addfield_fromobsarray,mod)
   mod.create = types.MethodType(create,mod)
   mod.iterate = types.MethodType(iterate,mod)
   mod.gettraj = types.MethodType(gettraj,mod)
   mod.destroy = types.MethodType(destroy,mod)
   mod.execute = types.MethodType(execute,mod)
   mod.playback = types.MethodType(playback,mod)
   mod.viewtsr  = types.MethodType(viewtsr,mod)
   mod.removeconstraint = types.MethodType(removeconstraint,mod)
   mod.runchomp = types.MethodType( runchomp, mod );


def shquot(s):
   return "'" + s.replace("'","'\\''") + "'"

def viewspheres(mod, robot=None, releasegil=False):
   cmd = 'viewspheres'
   if robot is not None:
      if hasattr(robot,'GetName'):
         cmd += ' robot %s' % robot.GetName()
      else:
         cmd += ' robot %s' % robot
   print 'cmd:', cmd
   return mod.SendCommand(cmd, releasegil)

def computedistancefield(mod, kinbody=None, cube_extent=None, aabb_padding=None,
                         cache_filename=None, releasegil=False):
   cmd = 'computedistancefield'
   if kinbody is not None:
      if hasattr(kinbody,'GetName'):
         cmd += ' kinbody ' + kinbody.GetName()
      else:
         cmd += ' kinbody ' + kinbody
   if cube_extent is not None:
      cmd += ' cube_extent %f' % cube_extent
   if aabb_padding is not None:
      cmd += ' aabb_padding %f' % aabb_padding
   if cache_filename is not None:
      cmd += ' cache_filename %s' % cache_filename
   print 'cmd:', cmd
   return mod.SendCommand(cmd, releasegil)

def addfield_fromobsarray(mod, kinbody=None, obsarray=None, sizes=None, lengths=None,
                          pose=None, releasegil=False):
   cmd = 'addfield_fromobsarray'
   if kinbody is not None:
      if hasattr(kinbody,'GetName'):
         cmd += ' kinbody %s' % kinbody.GetName()
      else:
         cmd += ' kinbody %s' % kinbody
   if obsarray is not None:
      cmd += ' obsarray %s' % obsarray
   if sizes is not None:
      cmd += ' sizes %s' % ' '.join([str(v) for v in sizes])
   if lengths is not None:
      cmd += ' lengths %s' % ' '.join([str(v) for v in lengths])
   if pose is not None:
      cmd += ' pose %s' % ' '.join([str(v) for v in pose])
   print 'cmd:', cmd
   return mod.SendCommand(cmd, releasegil)

def create(mod, robot=None, adofgoal=None, lambda_=None,
   starttraj=None, n_points=None, start_tsr=None, start_cost=None,
   everyn_tsr=None,
   use_momentum=None, use_hmc=None, hmc_resample_lambda=None, seed=None,
   epsilon=None, epsilon_self=None, obs_factor=None, obs_factor_self=None,
   no_report_cost=None, dat_filename=None, releasegil=False, **kwargs):
   cmd = 'create'
   if robot is not None:
      if hasattr(robot,'GetName'):
         cmd += ' robot ' + robot.GetName()
      else:
         cmd += ' robot ' + robot
   if adofgoal is not None:
      cmd += ' adofgoal ' + ' '.join([str(v) for v in adofgoal])
   if lambda_ is not None:
      cmd += ' lambda ' + str( lambda_ )
   if starttraj is not None:
      in_traj_data = starttraj.serialize(0) # options
      cmd += ' starttraj %s' % shquot(in_traj_data)
   if n_points is not None:
      cmd += ' n_points %d' % n_points
   if start_cost is not None:
      if isinstance(start_cost, str):
         cmd += ' start_cost \'%s\'' % start_cost
      else:
         cmd += ' start_cost \'%s %s\'' % (start_cost[0], start_cost[1])
   if use_momentum is not None and use_momentum:
      cmd += ' use_momentum'
   if use_hmc is not None and use_hmc:
      cmd += ' use_hmc'
   if hmc_resample_lambda is not None:
      cmd += ' hmc_resample_lambda %f' % hmc_resample_lambda
   if seed is not None:
      cmd += ' seed %d' % seed
   if epsilon is not None:
      cmd += ' epsilon %f' % epsilon
   if epsilon_self is not None:
      cmd += ' epsilon_self %f' % epsilon_self
   if obs_factor is not None:
      cmd += ' obs_factor %f' % obs_factor
   if obs_factor_self is not None:
      cmd += ' obs_factor_self %f' % obs_factor_self
   if no_report_cost is not None and no_report_cost:
      cmd += ' no_report_cost'
   if dat_filename is not None:
      cmd += ' dat_filename %s' % dat_filename
   print 'cmd:', cmd
   retval = mod.SendCommand(cmd, releasegil)

   if start_tsr is not None:
      command = "createtsr "
      cmd += start_tsr.serialize()
      cmd += " endtime 0"
      mod.SendCommand( command, releasegil )
   if everyn_tsr is not None:
      command = "createtsr "
      cmd += everyn_tsr.serialize()
      mod.SendCommand( command, releasegil )
   return retval;

def iterate(mod, run=None, n_iter=None, max_time=None,
            trajs_fileformstr=None, releasegil=False):

   cmd = 'iterate'
   if n_iter is not None:
      cmd += ' n_iter %d' % n_iter
   if max_time is not None:
      cmd += ' max_time %f' % max_time
   if trajs_fileformstr is not None:
      cmd += ' trajs_fileformstr %s' % shquot(trajs_fileformstr)

   return mod.SendCommand(cmd, releasegil)

def gettraj(mod, run=None, no_collision_check=None,
            no_collision_exception=None, no_collision_details=None,
            releasegil=False):

   cmd = 'gettraj'
   if no_collision_check is not None and no_collision_check:
      cmd += ' no_collision_check'
   if no_collision_exception is not None and no_collision_exception:
      cmd += ' no_collision_exception'
   if no_collision_details is not None and no_collision_details:
      cmd += ' no_collision_details'
   out_traj_data = mod.SendCommand(cmd, releasegil)
   
   return openravepy.RaveCreateTrajectory(mod.GetEnv(),'').deserialize(out_traj_data)
   
def destroy(mod, run=None, releasegil=False):
   cmd = 'destroy'
   return mod.SendCommand(cmd, releasegil)

def runchomp(mod, releasegil=False, **kwargs):
   # extract non-create args (run)
   n_iter = None
   max_time = None
   trajs_fileformstr = None
   if 'n_iter' in kwargs:
      n_iter = kwargs['n_iter']
      del kwargs['n_iter']
   if 'max_time' in kwargs:
      max_time = kwargs['max_time']
      del kwargs['max_time']
   if 'trajs_fileformstr' in kwargs:
      trajs_fileformstr = kwargs['trajs_fileformstr']
      del kwargs['trajs_fileformstr']
   # extract non-create args (gettraj)
   no_collision_check = None
   no_collision_exception = None
   no_collision_details = None
   if 'no_collision_check' in kwargs:
      no_collision_check = kwargs['no_collision_check']
      del kwargs['no_collision_check']
   if 'no_collision_exception' in kwargs:
      no_collision_exception = kwargs['no_collision_exception']
      del kwargs['no_collision_exception']
   if 'no_collision_details' in kwargs:
      no_collision_details = kwargs['no_collision_details']
      del kwargs['no_collision_details']
   run = create(mod, releasegil=releasegil, **kwargs)
   iterate(mod, run=run, n_iter=n_iter, max_time=max_time,
           trajs_fileformstr=trajs_fileformstr, releasegil=releasegil)
   traj = gettraj(mod, run=run,
      no_collision_check=no_collision_check,
      no_collision_exception=no_collision_exception,
      no_collision_details=no_collision_details,
      releasegil=releasegil)
   destroy(mod, run=run, releasegil=releasegil)
   return traj

def execute(mod, releasegil=False, **kwargs):
    cmd = 'execute'
    mod.SendCommand( cmd, releasegil )

     
def playback(mod, time=None, releasegil=False, **kwargs):
    cmd = "playback"
    if time is not None:
        cmd += ' time %f' %time

    mod.SendCommand( cmd, releasegil )
    
def viewtsr(mod, index=None, releasegil=False, **kwargs):
    cmd = 'viewtsr' 

    if index is not None:
        cmd += " %d" % index
    
    mod.SendCommand( cmd, releasegil )

def removeconstraint(mod, index=None, releasegil=False, **kwargs):
    cmd = 'removeconstraint' 

    if index is not None:
        cmd += " %d" % index
    
    mod.SendCommand( cmd, releasegil )

