from py_trees.composites import Sequence, Parallel
from py_trees.trees import BehaviourTree
from py_trees.common import ParallelPolicy, Status
from setup import robot, timestep
from arm_controller import MoveArm
from detector import JarFinder
from go_back import GoBack
from gripper import OpenGripper, CloseGripper
from navigator import ApproachJar, RotateRobot
from positions import grasp_pos, lift_pos, retract_pos, place_pos, raise_torso, rotate_pos
from sleep import Sleep
from constants import JAR_OFFSETS, PG_DIST

# This is the sequence of picking up the jar, placing it, then turning back to
# the counter. It was placed in a function given its repetitiveness.
def generate_repeat_sequence(i):
    id = (i+1, ) # The number to show in the names.
    seq = Sequence("Moving jar %d"%id, memory=True)
    seq.add_child(JarFinder("Detect jar %d"%id, "honey jar" if i==1 else "jam jar", JAR_OFFSETS[i]))
    if i: # Move back to get some space, but only for the 2nd and 3rd jars.
        seq.add_child(GoBack("Move back for pick up %d"%id, distance=.45))
    # Pick up sequence.
    seq.add_child(MoveArm("Grasping pose %d"%id, grasp_pos))
    seq.add_child(ApproachJar("Approach jar %d"%id))
    seq.add_child(CloseGripper("Grasp jar %d"%id))
    seq.add_child(GoBack("Move back after pick up %d"%id, distance=PG_DIST[i]))
    seq.add_child(MoveArm("Lift arm %d"%id, lift_pos))
    seq.add_child(MoveArm("Retract the arm after pick up %d"%id, retract_pos))
    seq.add_child(Sleep(name="Pause to kill the inertia"))
    # Face the table to place the jars on the table.
    seq.add_child(RotateRobot("Rotate towards table %d"%id, -2.5))
    seq.add_child(Sleep(name="Pause to kill the inertia"))
    # Place the jar on the table.
    seq.add_child(MoveArm("Releasing pose %d"%id, place_pos))
    seq.add_child(Sleep(name="Pause to kill the inertia"))
    seq.add_child(OpenGripper("Release the jar %d"%id))
    seq.add_child(MoveArm("Raise the torso %d"%id, raise_torso))
    if i == 2: # There is no need to continue after this for the 3rd jar.
        return seq # There is no need to continue after this for the 3rd jar.
    # Return to the kitchen counter to get the next jar.
    seq.add_child(Sleep(name="Pause to kill the inertia"))
    seq.add_child(GoBack("Create some space from the table %d"%id, distance=.4 if i else .0))
    seq.add_child(MoveArm("Retract the arm after placement %d"%id, rotate_pos))
    seq.add_child(Sleep(name="Pause to kill the inertia"))
    seq.add_child(RotateRobot("Rotate towards the counter %d"%id, .3))
    seq.add_child(Sleep(name="Pause to kill the inertia"))
    return seq

tree = Sequence("Root", memory=True) # The main sequence.
tree.add_child(
    Parallel("Initialising", policy=ParallelPolicy.SuccessOnAll(), children=[
        Sleep(name="Wait for world to stabilise", steps=120),
        MoveArm("Initialise to the grasp position", grasp_pos),
        MoveArm("Correct head angle", {"head_2_joint": -.3}), # Better camera angle.
    ]))
for i in range(3):
    tree.add_child(generate_repeat_sequence(i)) # Move the jars one by one.
tree = BehaviourTree(tree) # BehaviourTree allows better control.
tree.setup()

while robot.step(timestep) != -1:
    tree.tick()
    if tree.root.status in (Status.SUCCESS, Status.FAILURE):
        print("Tree finished executing with status %s" % tree.root.status)
        break