#!/usr/bin/env python
# coding: utf-8

# In[1]:


#%load_ext wurlitzer
#^^^ the wurlitzer extension is used to capture C/C++ output to be displayed in the notebook
#^^^ this is very useful for debugging, but it doesn't work on windows

#Python 2/3 compatibility
from __future__ import print_function,division

import time
from klampt import *
from klampt import vis
from IPython.display import clear_output
from klampt.vis.ipython import Playback

world = WorldModel()
if world.loadFile("../data/tx90scenario0.xml"):
    sim = Simulator(world)
    vis.add("world",world)
    playback_widget = Playback(vis.scene())
    
    #If you'd like to show the print output from loading the file, comment out this line
    clear_output()
    vis.show()
    display(playback_widget)
    
    #NOTE: if you are going to add/modify items to the world in the same cell that it is created, you will
    #need to place all of those calls in a begin_rpc/end_rpc block
    vis.addText("HUD1","0",position=(1,1))
    vis.add("ghost",world.robot(0).getConfig(),color=(1,0,0,0.5))
else:
    print("There was a problem loading the world file")
    
#Controls:
#left mouse click to rotate the view
#right click or ctrl+click to pan the view
#mouse wheel or shift+click to zoom the view


# In[2]:


import random
from klampt.plan.cspace import *
from klampt.plan.robotcspace import RobotCSpace
from klampt.model.collide import WorldCollider
from klampt.model.trajectory import *

#first time the plan-to is called
next_move_time = 1.0
#if this is too large, Jupyter will complain...
planning_time_limit = 2.0

robot = world.robot(0)

def plan_to(qa,qb):
    collider = WorldCollider(world)
    cspace = RobotCSpace(robot,collider)
    cspace.eps = 1e-2
    optimizing = True
    MotionPlan.setOptions(type="sbl",connectionThreshold=2,perturbationRadius=2,shortcut=True)
    planner = MotionPlan(cspace)
    planner.setEndpoints(qa,qb)
    t0 = time.time()
    foundTime = None
    path = None
    vis.addText("HUD2","Planning...",position=(1,5))
    while time.time() - t0 < planning_time_limit:
        planner.planMore(1)
        if foundTime == None:
            path = planner.getPath()
            if path != None and len(path) > 0:
                foundTime = time.time()
                vis.addText("HUD2","Found first path in %gs"%(foundTime-t0,),pos=(1,5))
                if not optimizing:
                    break
    if foundTime != None:
        if optimizing:
            vis.addText("HUD2","Optimized for another %gs"%(time.time()-foundTime,),pos=(1,5))
        path = planner.getPath()
    else:
        path = None
    return path


def control_loop(t,controller):
    global next_move_time
    vis.addText("HUD1","%.2f"%(t,),pos=(1,1))
    if t+0.02 >= next_move_time:
        print("Planning on next time step for %fs..."%(planning_time_limit,))
    if t >= next_move_time:
        qmin,qmax = robot.getJointLimits()
        q = [random.uniform(a,b) for (a,b) in zip(qmin,qmax)]
        print("Calling plan to...",q)
        vis.setColor("ghost",1,1,0,0.5)
        vis.add("ghost",q)
        try:
            path = plan_to(controller.getCommandedConfig(),q)
        except Exception as e:
            print(e)
            path = None
        if path != None:
            vis.setColor("ghost",0,1,0,0.5)
            for i,q in enumerate(path):
                if i == 1:
                    controller.addMilestoneLinear(q)
                elif i >= 2:
                    controller.addMilestoneLinear(q)
        else:
            vis.setColor("ghost",1,0,0,0.5)
        #controller.setMilestone(q)
        next_move_time += 5.0
    pass


# In[3]:


dt = 0.02

def do_reset():
    global next_move_time
    sim.reset()
    sim.updateWorld()
    next_move_time = 2.0

def do_advance():
    global world,sim
    if world.numRobots() > 0:
        control_loop(sim.getTime(),sim.controller(0))  #call code in the above cell
    #print("Simulating...",dt)
    sim.simulate(dt)
    sim.updateWorld()
    #print("Done.")

#this binds the playback widget buttons
playback_widget.advance = do_advance
playback_widget.reset = do_reset
playback_widget.quiet = False


# In[ ]: