# Path: lfd/rapprentice/animate_traj.py
# def animate_traj(traj, robot, pause=True, step_viewer=1, restore=True, callback=None, execute_step_cond=None):
#     """make sure to set active DOFs beforehand"""
#     if restore: _saver = openravepy.RobotStateSaver(robot)
#     if step_viewer or pause: viewer = trajoptpy.GetViewer(robot.GetEnv())
#     for (i,dofs) in enumerate(traj):
#         sys.stdout.write("step %i/%i\r"%(i+1,len(traj)))
#         sys.stdout.flush()
#         if callback is not None: callback(i)
#         if execute_step_cond is not None and not execute_step_cond(i): continue
#         robot.SetActiveDOFValues(dofs)
#         if pause: viewer.Idle()
#         elif step_viewer!=0 and (i%step_viewer==0 or i==len(traj)-1): viewer.Step()
#     sys.stdout.write("\n")
# 
# Path: lfd/rapprentice/ropesim.py
# def transform(hmat, p):
# def in_grasp_region(robot, lr, pt):
#     def on_inner_side(pt, finger_lr):
# def retime_traj(robot, inds, traj, max_cart_vel=.02, max_finger_vel=.02, upsample_time=.1):
# def observe_cloud(pts, radius, upsample=0, upsample_rad=1):
#     def __init__(self, env, robot, rope_params):
#     def create(self, rope_pts):
#     def __del__(self):
#     def step(self):
#     def settle(self, max_steps=100, tol=.001, animate=False):
#     def observe_cloud(self, upsample=0, upsample_rad=1):
#     def raycast_cloud(self, T_w_k=None, obj=None, z=1., endpoints=0):
#     def grab_rope(self, lr):
#     def release_rope(self, lr):
#     def is_grabbing_rope(self, lr):
# class Simulation(object):
# 
# Path: lfd/rapprentice/math_utils.py
# def interp2d(x,xp,yp):
# def interp_mat(x, xp):
# def normalize(x):
# def normr(x):
# def normc(x):
# def norms(x,ax):
# def intround(x):
# def deriv(x):
# def linspace2d(start,end,n):
# def remove_duplicate_rows(mat):
# def invertHmat(hmat):
#     T = len(x)
#     R = hmat[:3,:3]
# 
# Path: lfd/rapprentice/plotting_openrave.py
# def draw_grid(env, f, mins, maxes, xres = .1, yres = .1, zres = .04, color = (1,1,0,1)):
# 
# Path: lfd/util/util.py
# def redprint(msg):
# def yellowprint(msg):
#     def parse_args(self, *args, **kw):
#     def __init__(self, adict):
#     def __init__(self):
#     def __enter__(self):
#     def __exit__(self, *_):
# class ArgumentParser(argparse.ArgumentParser):
# class Bunch(object):
# class suppress_stdout(object):


import bulletsimpy
import openravepy
import numpy as np
import re
import random
import settings
    import scipy.interpolate as si
from numpy import asarray
from lfd.rapprentice import animate_traj, ropesim, math_utils as mu, plotting_openrave
from lfd.util import util
        valid_inds = np.apply_along_axis(np.linalg.norm, 1, (ee_trajs['r'][:,:3,3] - ee_trajs['l'][:,:3,3])) < min_length - 0.02
        if not np.all(valid_inds):
            full_traj = (full_traj[0][valid_inds,:], full_traj[1])
            util.yellowprint("The grippers of the trajectory goes too far apart. Some steps of the trajectory are being removed.")
    return full_traj

def load_random_start_segment(demofile):
    start_keys = [k for k in demofile.keys() if k.startswith('demo') and k.endswith('00')]
    seg_name = random.choice(start_keys)
    return demofile[seg_name]['cloud_xyz'][:,:3]

def load_fake_data_segment(demofile, fake_data_segment, fake_data_transform):
    fake_seg = demofile[fake_data_segment]
    new_xyz = np.squeeze(fake_seg["cloud_xyz"])[:,:3]
    hmat = openravepy.matrixFromAxisAngle(fake_data_transform[3:6])
    hmat[:3,3] = fake_data_transform[0:3]
    new_xyz = new_xyz.dot(hmat[:3,:3].T) + hmat[:3,3][None,:]
    joint_names = asarray(fake_seg["joint_states"]["name"])
    joint_values = asarray(fake_seg["joint_states"]["position"][0])
    return new_xyz, joint_names, joint_values

def unif_resample(traj, max_diff, wt = None):        
    """
    Resample a trajectory so steps have same length in joint space    
    """
    tol = .005
    if wt is not None: 
        wt = np.atleast_2d(wt)
        traj = traj*wt
        

# Path: lfd/rapprentice/animate_traj.py
# def animate_traj(traj, robot, pause=True, step_viewer=1, restore=True, callback=None, execute_step_cond=None):
#     """make sure to set active DOFs beforehand"""
#     if restore: _saver = openravepy.RobotStateSaver(robot)
#     if step_viewer or pause: viewer = trajoptpy.GetViewer(robot.GetEnv())
#     for (i,dofs) in enumerate(traj):
#         sys.stdout.write("step %i/%i\r"%(i+1,len(traj)))
#         sys.stdout.flush()
#         if callback is not None: callback(i)
#         if execute_step_cond is not None and not execute_step_cond(i): continue
#         robot.SetActiveDOFValues(dofs)
#         if pause: viewer.Idle()
#         elif step_viewer!=0 and (i%step_viewer==0 or i==len(traj)-1): viewer.Step()
#     sys.stdout.write("\n")
# 
# Path: lfd/rapprentice/ropesim.py
# def transform(hmat, p):
# def in_grasp_region(robot, lr, pt):
#     def on_inner_side(pt, finger_lr):
# def retime_traj(robot, inds, traj, max_cart_vel=.02, max_finger_vel=.02, upsample_time=.1):
# def observe_cloud(pts, radius, upsample=0, upsample_rad=1):
#     def __init__(self, env, robot, rope_params):
#     def create(self, rope_pts):
#     def __del__(self):
#     def step(self):
#     def settle(self, max_steps=100, tol=.001, animate=False):
#     def observe_cloud(self, upsample=0, upsample_rad=1):
#     def raycast_cloud(self, T_w_k=None, obj=None, z=1., endpoints=0):
#     def grab_rope(self, lr):
#     def release_rope(self, lr):
#     def is_grabbing_rope(self, lr):
# class Simulation(object):
# 
# Path: lfd/rapprentice/math_utils.py
# def interp2d(x,xp,yp):
# def interp_mat(x, xp):
# def normalize(x):
# def normr(x):
# def normc(x):
# def norms(x,ax):
# def intround(x):
# def deriv(x):
# def linspace2d(start,end,n):
# def remove_duplicate_rows(mat):
# def invertHmat(hmat):
#     T = len(x)
#     R = hmat[:3,:3]
# 
# Path: lfd/rapprentice/plotting_openrave.py
# def draw_grid(env, f, mins, maxes, xres = .1, yres = .1, zres = .04, color = (1,1,0,1)):
# 
# Path: lfd/util/util.py
# def redprint(msg):
# def yellowprint(msg):
#     def parse_args(self, *args, **kw):
#     def __init__(self, adict):
#     def __init__(self):
#     def __enter__(self):
#     def __exit__(self, *_):
# class ArgumentParser(argparse.ArgumentParser):
# class Bunch(object):
# class suppress_stdout(object):


# Path: lfd/environment/sim_util.py
import bulletsimpy
import openravepy
import numpy as np
import re
import random
import settings
    import scipy.interpolate as si
from numpy import asarray
from lfd.rapprentice import animate_traj, ropesim, math_utils as mu, plotting_openrave
from lfd.util import util

        valid_inds = np.apply_along_axis(np.linalg.norm, 1, (ee_trajs['r'][:,:3,3] - ee_trajs['l'][:,:3,3])) < min_length - 0.02
        if not np.all(valid_inds):
            full_traj = (full_traj[0][valid_inds,:], full_traj[1])
            util.yellowprint("The grippers of the trajectory goes too far apart. Some steps of the trajectory are being removed.")
    return full_traj

def load_random_start_segment(demofile):
    start_keys = [k for k in demofile.keys() if k.startswith('demo') and k.endswith('00')]
    seg_name = random.choice(start_keys)
    return demofile[seg_name]['cloud_xyz'][:,:3]

def load_fake_data_segment(demofile, fake_data_segment, fake_data_transform):
    fake_seg = demofile[fake_data_segment]
    new_xyz = np.squeeze(fake_seg["cloud_xyz"])[:,:3]
    hmat = openravepy.matrixFromAxisAngle(fake_data_transform[3:6])
    hmat[:3,3] = fake_data_transform[0:3]
    new_xyz = new_xyz.dot(hmat[:3,:3].T) + hmat[:3,3][None,:]
    joint_names = asarray(fake_seg["joint_states"]["name"])
    joint_values = asarray(fake_seg["joint_states"]["position"][0])
    return new_xyz, joint_names, joint_values

def unif_resample(traj, max_diff, wt = None):        
    """
    Resample a trajectory so steps have same length in joint space    
    """
    tol = .005
    if wt is not None: 
        wt = np.atleast_2d(wt)
        traj = traj*wt
        

    dl = mu.norms(traj[1:] - traj[:-1],1)