Colony » Colony Scout

#!/usr/bin/python

import wx

import wx.lib.intctrl 

import string

import subprocess

import signal

from collections import defaultdict

# ROS imports

import roslib; roslib.load_manifest("scoutsim")

import rospy

from scoutsim.srv import *

# Everyone likes some nice globals

scouts = defaultdict(str) # Associative array in the form [scout_name: current_behavior]

processes = defaultdict() # Associative array in the form [scout_name: process_handler]

#

# ROS functions

#

def AddToSimulator(name, x, y, theta):

        if name in scouts:

                return False

        scouts[name] = ""

        print "Calling ROS function to create scout "+name

        try:

                rospy.wait_for_service('/spawn');

                service = rospy.ServiceProxy('/spawn', Spawn);

                resp = service(x, y, theta, name)

                return True

        except rospy.ServiceException, e:

                return False

def GetBehaviors():

        return ["CW Circle", "CCW Circle", "Odometry", "Navigation Map", "Scheduler", "Warehouse", "Line Follow", "WL Test"]

def SetBehavior(name, behaviorLabel, behaviorButton=False, behavior="Pause"):

        if behavior == "Pause":

                if name not in processes:

                        return False

                if not behaviorButton.GetValue():

                        processes[name].send_signal(signal.SIGCONT)

                        print "Scout "+name+" was resumed."

                        behaviorButton.SetLabel("Pause")

                        behaviorLabel.SetLabel("  |  Current Behavior: "+scouts[name])

                else:

                        processes[name].send_signal(signal.SIGSTOP)

                        print "Scout "+name+" was paused."

                        behaviorButton.SetLabel("Resume")

                        behaviorLabel.SetLabel("  |  Current Behavior: Paused")

                return True

        if name in processes:

                processes[name].kill()

        processes[name] = subprocess.Popen("rosrun libscout libscout "+name+" "+str(GetBehaviors().index(behavior))+"", shell=True) 

        print "Scout "+name+" was set to behavior: "+behavior

        behaviorLabel.SetLabel("  |  Current Behavior: "+behavior)

        scouts[name] = behavior

def KillScout(name):

        try:

                rospy.wait_for_service('/kill');

                service = rospy.ServiceProxy('/kill', Kill);

                if name in processes:

                        processes[name].kill() # This is perhaps not working

                resp = service(name)

                del scouts[name]

                del processes[name]

                return True

        except rospy.ServiceException, e:

                return False

def PauseScout(name, behaviorLabel, behaviorButton):

        return SetBehavior(name, behaviorLabel, behaviorButton)

#

# End ROS functions

#

class GUI(wx.Frame):

        def __init__(self, parent, title):

                super(GUI, self).__init__(parent, title=title, style=wx.DEFAULT_FRAME_STYLE ^ wx.RESIZE_BORDER, size=(600, 600))

                self.InitUI()

                self.Show()     

        def InitUI(self):

                self.window = wx.ScrolledWindow(self, style=wx.VSCROLL)

                self.mainArea = wx.GridSizer(cols=1)

                sizer = wx.FlexGridSizer(rows=4, cols=8, hgap=5, vgap=5)

                # Labels

                blankText = wx.StaticText(self.window, label="")

                newScout = wx.StaticText(self.window, label="New Scout")

                newScoutName = wx.StaticText(self.window, label="Name:")

                startXTitle = wx.StaticText(self.window, label="X:")

                startYTitle = wx.StaticText(self.window, label="Y:")

                startThetaTitle = wx.StaticText(self.window, label=u'\u0398')

                # Inputs

                newScoutInput = wx.TextCtrl(self.window)

                startX = wx.lib.intctrl.IntCtrl(self.window, min=0)

                startY = wx.lib.intctrl.IntCtrl(self.window, min=0)

                startTheta = wx.lib.intctrl.IntCtrl(self.window, min=0, max=360)

                addButton = wx.Button(self.window, id=wx.ID_ADD)

                # Pretty Stuff

                hLine = wx.StaticLine(self.window, size=(600, 5))

                bottomHLine = wx.StaticLine(self.window, size=(600, 5))

                # Row 0

                sizer.Add(newScout)

                for i in range(7):

                        sizer.AddStretchSpacer(1)

                # Row 1

                sizer.AddMany([newScoutName, (newScoutInput, 0, wx.EXPAND), startXTitle, startX, startYTitle, startY, startThetaTitle, startTheta])

                # Row 2

                for i in range(7):

                        sizer.AddStretchSpacer(1)

                sizer.Add(addButton)

                # Row 3

                # Events

                addButton.Bind(wx.EVT_BUTTON, lambda event: self.NewScout(event, newScoutInput, startX, startY, startTheta))

                sizer.AddGrowableCol(idx=1)

                self.mainArea.Add(sizer, proportion=1, flag=wx.ALL|wx.EXPAND, border=10)

                self.window.SetSizer(self.mainArea)

                self.sizer = defaultdict()

        def NewScout(self, e, name, x, y, theta):

                if (name):

                        scoutCreated = self.AddScout(name.GetValue(), x.GetValue(), y.GetValue(), theta.GetValue())

                        if (scoutCreated):

                                wx.MessageBox("Scout '"+name.GetValue()+"' created successfully at ("+str(x.GetValue())+", "+str(y.GetValue())+") facing "+str(theta.GetValue())+u"\u00B0.", 'Scout Created', wx.OK | wx.ICON_INFORMATION)

                                name.SetValue("")

                                x.SetValue(0)

                                y.SetValue(0)

                                theta.SetValue(0)

                                return True

                wx.MessageBox('Error creating scout.  Please check that all information was entered properly and there are no duplicate scout names.', 'Info', wx.OK | wx.ICON_ERROR)

        def AddScoutBox(self, name):

                self.sizer[name] = wx.FlexGridSizer(rows=2, cols=4, hgap=5, vgap=5)

                # Labels

                scoutName = wx.StaticText(self.window, label="Scout: "+name)

                behaviorLabel = wx.StaticText(self.window, label="Behavior: ")

                currBehaviorLabel = wx.StaticText(self.window, label="  |  Current Behavior: Slacking off")

                # Inputs

                scoutChoices = wx.Choice(self.window, choices=GetBehaviors())

                pauseButton = wx.ToggleButton(self.window, label="Pause")

                runButton = wx.Button(self.window, label="Run")

                killButton = wx.Button(self.window, label="Kill")

                self.sizer[name].Add(scoutName)

                self.sizer[name].Add(currBehaviorLabel, wx.EXPAND|wx.ALIGN_RIGHT)

                self.sizer[name].AddStretchSpacer(1)

                self.sizer[name].Add(killButton, wx.ALIGN_RIGHT)

                self.sizer[name].Add(behaviorLabel)

                self.sizer[name].Add(scoutChoices, wx.EXPAND)

                self.sizer[name].Add(runButton)

                self.sizer[name].Add(pauseButton, wx.ALIGN_RIGHT)

                # Events

                pauseButton.Bind(wx.EVT_TOGGLEBUTTON, lambda event: PauseScout(name, currBehaviorLabel, pauseButton))

                runButton.Bind(wx.EVT_BUTTON, lambda event: SetBehavior(name, scoutChoices.GetStringSelection(), currBehaviorLabel))

                killButton.Bind(wx.EVT_BUTTON, lambda event: self.RemoveScout(name))

                self.mainArea.Add(self.sizer[name], proportion=1, flag=wx.ALL|wx.EXPAND, border=10)

                self.window.Layout()

                return True

        def AddScout(self, name, x, y, theta):

                if AddToSimulator(name, x, y, theta):

                        return self.AddScoutBox(name)

                wx.MessageBox("ROS call failed!", 'ROS Error', wx.OK | wx.ICON_ERROR)

                return False

        def RemoveScout(self, name):

                # I don't trust this bit. Remove() should work by itself, but it doesn't want to.

                # Hide()ing before Remove()ing does work for some reason though. It makes me think 

                # the sizer is still around somewhere, sneakily leaking memory.

                # Alas, such is life.

                print "Hidden scoutLayout for "+name+": "+str(self.mainArea.Hide(self.sizer[name], True))

                print "Removed scoutLayout for "+name+": "+str(self.mainArea.Remove(self.sizer[name]))+" (Maybe...)"

                del self.sizer[name]

                if KillScout(name):

                        self.window.Layout()

                        self.window.Refresh()

                        self.window.Update()

                return

if __name__ == '__main__':

        subprocess.Popen("rosrun scoutsim scoutsim_node race", shell=True)

        app = wx.App()

        GUI(None, title='Colony Scout Manager')

        app.MainLoop()