simple_cli

#!/usr/bin/python3 -i
"""

Simple Cozmo CLI
================
A simple (and crude) Command Line Interface for Cozmo.

Available commands:

    (python code)
        Executable statements are directly passed to python's interpreter.
        The magic global variables ans holds the results of evaluating the
        last expression.
        
    exit    
        Quit the CLI; return to Python.

    monitor(robot[, Event])
        Monitor all event types in the dispatch table, or a specified type of Event
    
    unmonitor(robot[, Event])
        Turn off monitoring

    runfsm(module_name)
        Imports or reloads a state machine module and runs the
        state machine.

    tracefsm(trace_level)
        Sets the FSM tracing level (0-9). With no argument, returns
        the current level.

    tm message
        Sends 'message' as a text message to the currently running
        state machine.

    show active
        Shows active state nodes and transitions.

    show cam_viewer | path_viewer | worldmap_viewer
        Displays OpenCV viewer of the specified type.

    !cmd
        Runs 'cmd' in a shell and prints the result.
*********

NOTE: On Windows you must install the pyreadline package in
order to get readline.

Author:     David S. Touretzky, Carnegie Mellon University
=======

"""

# All that stuff we really need in order to get going.
import atexit 
import code
import datetime
import logging
import os
import platform
import re
import readline 
import rlcompleter 
import subprocess
import sys
import time
import traceback
from importlib import __import__, reload

try:
    from termcolor import cprint
except:
    def cprint(string,color=None):
        print(string)

try:
    import matplotlib
    import matplotlib.pyplot as plt
except:
    pass


import cozmo
from cozmo.util import *

from event_monitor import monitor, unmonitor

import cozmo_fsm
from cozmo_fsm import *
from cozmo_fsm.worldmap import ArucoMarkerObj, CustomMarkerObj, WallObj

# tab completion
readline.parse_and_bind('tab: complete') 
# history file
if 'HOME' in os.environ:  # Linux
    histfile = os.path.join(os.environ['HOME'], '.pythonhistory')
elif 'USERPROFILE' in os.environ:  # Windows
    histfile = os.path.join(os.environ['USERPROFILE'], '.pythonhistory')
else:
    histfile = '.pythonhistory'

try: 
    readline.read_history_file(histfile) 
except IOError: 
    pass 
atexit.register(readline.write_history_file, histfile) 
del rlcompleter

os_version = platform.system()
del platform

# Put current directory on search path.
if '.' not in sys.path:
    sys.path.append('.')

res = 0
ans = None

RUNNING = True


def simple_cli_callback(_wcube1, _wcube2, _wcube3, _wcharger):
    global wcube1, wcube2, wcube3, wcharger
    wcube1 = _wcube1
    wcube2 = _wcube2
    wcube3 = _wcube3
    wcharger = _wcharger

def runfsm(module_name, running_modules=dict()):
    """runfsm('modname') reloads that module and expects it to contain
    a class of the same name. It calls that class's constructor and then
    calls the instance's start() method."""

    global running_fsm
    running_fsm = cozmo_fsm.program.running_fsm
    if running_fsm:
        robot.stop_all_motors()
        running_fsm.stop()

    r_py = re.compile('.*\.py$')
    if r_py.match(module_name):
        print("\n'%s' is not a module name. Trying '%s' instead.\n" %
              (module_name, module_name[0:-3]))
        module_name = module_name[0:-3]

    found = False
    try:
        reload(running_modules[module_name])
        found = True
    except KeyError: pass
    except: raise
    if not found:
        try:
            running_modules[module_name] = __import__(module_name)
        except ImportError as e:
            print("Error loading %s: %s.  Check your search path.\n" %
                  (module_name,e))
            return
        except Exception as e:
            print('\n===> Error loading %s:' % module_name)
            raise

    py_filepath = running_modules[module_name].__file__
    fsm_filepath = py_filepath[0:-2] + 'fsm'
    try:
        py_time = datetime.datetime.fromtimestamp(os.path.getmtime(py_filepath))
        fsm_time = datetime.datetime.fromtimestamp(os.path.getmtime(fsm_filepath))
        if py_time < fsm_time:
            cprint('Warning: %s.py is older than %s.fsm. Should you run genfsm?' %
                   (module_name,module_name), color="yellow")
    except: pass

    # The parent node class's constructor must match the module name.
    the_module = running_modules[module_name]
    the_class = the_module.__getattribute__(module_name) \
                if module_name in dir(the_module) else None
    if not isinstance(the_class,type) or not issubclass(the_class,StateNode):
        cprint("Module %s does not contain a StateNode class named %s.\n" %
              (module_name, module_name), color="red")
        return       
    the_module.robot = robot
    the_module.world = robot.world
    the_module.charger = robot.world.charger
    the_module.cube1 = robot.world.light_cubes[cozmo.objects.LightCube1Id]
    the_module.cube2 = robot.world.light_cubes[cozmo.objects.LightCube2Id]
    the_module.cube3 = robot.world.light_cubes[cozmo.objects.LightCube3Id]
    # Class's __init__ method will call setup, which can reference the above variables.
    running_fsm = cozmo_fsm.program.running_fsm = the_class()
    running_fsm.simple_cli_callback = simple_cli_callback
    cli_globals = globals()
    cli_globals['running_fsm'] = running_fsm
    robot.loop.call_soon(running_fsm.start)
    return running_fsm


def text_message(msg):
    running_fsm = cozmo_fsm.program.running_fsm
    if not running_fsm or not running_fsm.running:
        print('No state machine running. Use runfsm(module_name) to start a state machine.')
        return
    try:
        running_fsm.robot.erouter.post(TextMsgEvent(msg))
    except KeyboardInterrupt: raise
    except Exception as e:
        traceback.print_exc()
        print()


def start_stuff(args):
    spec = args[0] if len(args) > 0 else ""
    if spec == 'perched':
        try:
            cams = [int(x) for x in args[1:]]
        except:
            print('Usage: start perched [camera_number...]')
            return
        robot.world.perched.start_perched_camera_thread(cams)
    elif spec == 'server':
        robot.world.server.start_server_thread()
    elif spec == 'client':
        if len(args) != 2:
            print('Usage: start client IP_address')
            return
        robot.world.client.start_client_thread(args[1])
    elif spec == 'shared_map':
        robot.world.client.use_shared_map()
        print('Now using shared map.')
    else:
        print("""Usage:
  start perched
  start server
  start client [IP_Address]
  start shared_map
""")

def show_stuff(args):
    global running_fsm
    running_fsm = cozmo_fsm.program.running_fsm
    spec = args[0] if len(args) > 0 else ""
    if spec == 'active':
        if not running_fsm:
            print('No state machine present.')
        elif not running_fsm.running:
            print("State machine '%s' is not running." % running_fsm.name)
        else:
            show_active(running_fsm,0)
    elif spec == "kine":
        show_kine(args[1:])
    elif spec == 'cam_viewer' or spec=='viewer':
        if running_fsm:
            running_fsm.stop()
        running_fsm = StateMachineProgram(cam_viewer=True,
                                          simple_cli_callback=simple_cli_callback)
        running_fsm.set_name("CamViewer")
        running_fsm.simple_cli_callback = simple_cli_callback
        cozmo_fsm.program.running_fsm = running_fsm
        robot.loop.call_soon(running_fsm.start)
    elif spec == "crosshairs":
        if running_fsm:
            running_fsm.viewer_crosshairs = not running_fsm.viewer_crosshairs
    elif spec == "particle_viewer":
        if not robot.world.particle_viewer:
            robot.world.particle_viewer = ParticleViewer(robot)
            robot.world.particle_viewer.start()
    elif spec == "path_viewer":
        if not robot.world.path_viewer:
            robot.world.path_viewer = PathViewer(robot,world.rrt)
            robot.world.path_viewer.start()
    elif spec == "worldmap_viewer":
        if not robot.world.worldmap_viewer:
            robot.world.worldmap_viewer = WorldMapViewer(robot)
            robot.world.worldmap_viewer.start()
    elif spec == "pose":
        show_pose()
    elif spec == "landmarks":
        show_landmarks()
    elif spec == "objects":
        show_objects()
    elif spec == "particle":
        show_particle(args[1:])
    elif spec == "camera":
        show_camera(args[1:])
    else:
        print("""Invalid option. Try one of:
  show viewer
  show crosshairs
  show worldmap_viewer
  show particle_viewer
  show path_viewer
  show active
  show kine [joint]
  show pose
  show landmarks
  show objects
  show particle [n]
  show camera n
  """)


def show_active(node,depth):
    if node.running: print('  '*depth, node)
    for child in node.children.values():
        show_active(child, depth+1)
    for trans in node.transitions:
        if trans.running: print('  '*(depth+1), trans)

def show_kine(args):
    if len(args) == 0:
        show_kine_tree(0, robot.kine.joints['base'])
        print()
    elif len(args) == 1:
        show_kine_joint(args[0])
    else:
        print('Usage:  show kine [joint]')

def show_kine_tree(level, joint):
    qstring = ''
    if joint.type != 'fixed':
        if isinstance(joint.q, (float,int)):
            qval = ('%9.5g' % joint.q).strip()
            if joint.type == 'revolute':
                qval = qval + (' (%.1f deg.)' % (joint.q*180/pi))
        else:
            qval = '(' + (', '.join([('%9.5g' % v).strip() for v in joint.q])) + ')'
        qstring = ' q=' + qval
    print('  '*level, joint.name, ': ', joint.type, qstring, sep='')
    for child in joint.children:
        show_kine_tree(level+1, child)

def show_kine_joint(name):
    if name not in robot.kine.joints:
        print("'"+repr(name)+"' is not the name of a joint.  Try 'show kine'.")
        return
    joint = robot.kine.joints[name]
    fmt = '%10s'

    def formatq(type,val):
        if type == 'revolute':
            if val == inf:
                return 'inf'
            elif val == -inf:
                return '-inf'
            jrad = ('%9.5g' % val).strip() + ' radians'
            jdeg = '(' + ('%9.5g' % (val * 180/pi)).strip() + ' degrees)' if val != 0 else ''
            return jrad + ' ' + jdeg
        elif type == 'prismatic':
            return ('%9.5g' % val).strip() + ' mm'
        elif type == 'fixed':
            return ''
        elif type == 'world':
            if val is None:
                return ''
            else:
                return '(' + (', '.join(['%9.5g' % x for x in val])) + ')'
        else:
            raise ValueError(type)
        
    print(fmt % 'Name:', name)
    print(fmt % 'Type:', joint.type)
    print(fmt % 'Parent:', joint.parent.name if joint.parent else '')
    print(fmt % 'Descr.:', joint.description)
    print(fmt % 'q:', formatq(joint.type, joint.q))
    print(fmt % 'qmin:', formatq(joint.type, joint.qmin))
    print(fmt % 'qmax:', formatq(joint.type, joint.qmax))
    print(fmt % 'DH d:', formatq('prismatic',joint.d))
    print(fmt % 'DH theta:', formatq('revolute',joint.theta))
    print(fmt % 'DH alpha:', formatq('revolute',joint.alpha))
    print(fmt % 'DH r:', formatq('prismatic',joint.r))
    print(fmt % 'Link in base frame:')
    tprint(robot.kine.link_to_base(name))
    print()
    
def show_pose():
    print('robot.pose is:   %6.1f %6.1f @ %6.1f deg.' %
          (robot.pose.position.x, robot.pose.position.y, robot.pose_angle.degrees))
    print('particle filter: %6.1f %6.1f @ %6.1f deg.' %
          (*robot.world.particle_filter.pose[0:2],
           robot.world.particle_filter.pose[2]*180/pi))
    print()

def show_landmarks():
    landmarks = robot.world.particle_filter.sensor_model.landmarks
    print('The particle filter has %d landmark%s:' %
          (len(landmarks), '' if (len(landmarks) == 1) else 's'))
    show_landmarks_helper(landmarks)

def show_landmarks_helper(landmarks):
    for (key,value) in landmarks.items():
        x = value[0][0,0]
        y = value[0][1,0]
        theta = value[1] * 180/pi
        if isinstance(key, int):
            print('  Aruco marker %d' % key, end='')
        elif isinstance(key, LightCube):
            print('  Cube %d' % key.cube_id, end='')
        else:
            print('  %r' % key, end='')
        print(' at (%.1f, %.1f) @ %.1f deg' % (x, y, theta))
    print()
    

def show_objects():
    objs = robot.world.world_map.objects
    print('%d object%s in the world map:' %
          (len(objs), '' if len(objs) == 1 else 's'))
    basics = [charger, cube1, cube2, cube3]
    ordered_keys = []
    for key in basics:
        if key in objs:
            ordered_keys.append(key)
    customs = []
    arucos = []
    walls = []
    misc = []
    for (key,value) in objs.items():
        if key in basics:
            pass
        elif isinstance(value, CustomMarkerObj):
            customs.append(key)
        elif isinstance(value, ArucoMarkerObj):
            arucos.append(key)
        elif isinstance(value, WallObj):
            walls.append(key)
        else:
            misc.append(key)
    arucos.sort()
    walls.sort()
    ordered_keys = ordered_keys + customs + arucos + walls + misc
    for key in ordered_keys:
        print('  ', objs[key])
    print()

def show_particle(args):
    if len(args) == 0:
        particle_number = '(best=%d)' % robot.world.particle_filter.best_particle_index
        particle = robot.world.particle_filter.best_particle
    elif len(args) > 1:
        print('Usage:  show particle [number]')
        return
    else:
        try:
            particle_number = int(args[0])
            particle = robot.world.particle_filter.particles[particle_number]
        except ValueError:
            print('Usage:  show particle [number]')
            return
        except IndexError:
            print('Particle number must be between 0 and',
                  len(robot.world.particle_filter.particles)-1)
            return
    print ('Particle %s:  x=%6.1f  y=%6.1f  theta=%6.1f deg   log wt=%f [%.25f]' %
           (particle_number, particle.x, particle.y, particle.theta*180/pi,
            particle.log_weight, particle.weight))
    if len(particle.landmarks) > 0:
        print('Landmarks:')
        show_landmarks_helper(particle.landmarks)
    else:
        print()

def show_camera(args):
    if len(args) != 1:
        print('Usage:  show camera n, where n is a camera number, typically 0 or 1.')
        return
    try:
        cam = int(args[0])
    except ValueError:
        show_camera()
    robot.world.perched.check_camera(cam)

def do_shell_command(cmd):
    try:
        subprocess.call(cmd, shell=True)
    except Exception as e:
        print(e)


def run(sdk_conn):
    global robot
    robot = sdk_conn.wait_for_robot()
    if "TK_VIEWER" in sys.argv:
        time.sleep(1.5) # allow time for Tk to set up the viewer window
    try:
        robot.erouter = cozmo_fsm.evbase.EventRouter()
        robot.erouter.robot = robot
        cozmo_fsm.evbase.robot_for_loading = robot
    except: pass
    cli_loop(robot)


def cli_loop(robot):
    global ans, RUNNING

    cli_globals = globals()
    cli_globals['world'] = robot.world
    cli_globals['light_cubes'] = world.light_cubes
    cli_globals['cube1'] = light_cubes[cozmo.objects.LightCube1Id]
    cli_globals['cube2'] = light_cubes[cozmo.objects.LightCube2Id]
    cli_globals['cube3'] = light_cubes[cozmo.objects.LightCube3Id]
    cli_globals['charger'] = robot.world.charger
    cli_globals['ans'] = None

    running_fsm = cozmo_fsm.program.running_fsm = \
        StateMachineProgram(cam_viewer=False, simple_cli_callback=simple_cli_callback)
    cli_globals['running_fsm'] = running_fsm
    running_fsm.start()

    cli_loop._console = code.InteractiveConsole()
    while True:
        if RUNNING == False:
            return
        cli_loop._line = ''
        while cli_loop._line == '':
            readline.write_history_file(histfile)
            try:
                if os_version == 'Darwin':   # Tkinter breaks console on Macs
                    print('C> ', end='')
                    cli_loop._line = sys.stdin.readline().strip()
                else:
                    cli_loop._line = cli_loop._console.raw_input('C> ').strip()
            except KeyboardInterrupt:
                process_interrupt()
                continue
            except EOFError:
                print("EOF.\nType 'exit' to exit.\n")
                continue
            try:
                robot.kine.get_pose()
            except: pass
        if cli_loop._line[0] == '!':
            do_shell_command(cli_loop._line[1:])
            continue
        elif cli_loop._line[0:3] == 'tm ' or cli_loop._line == 'tm':
            text_message(cli_loop._line[3:])
            continue
        elif cli_loop._line[0:5] == 'show ' or cli_loop._line == 'show':
            show_args = cli_loop._line[5:].split(' ')
            show_stuff(show_args)
            continue
        elif cli_loop._line[0:6] == 'start ' or cli_loop._line == 'start':
            start_args = cli_loop._line[6:].split(' ')
            start_stuff(start_args)
            continue
        cli_loop._do_await = False
        if cli_loop._line[0:7] == 'import ' or cli_loop._line[0:5] == 'from '  or \
               cli_loop._line[0:7] == 'global ' or cli_loop._line[0:4] == 'del '   or \
               cli_loop._line[0:4] == 'for ' or \
               cli_loop._line[0:4] == 'def '    or cli_loop._line[0:6] == 'async ' :
            # Can't use assignment to capture a return value, so None.
            ans = None
        elif cli_loop._line[0:6] == 'await ':
            cli_loop._do_await = True
            cli_loop._line = 'ans=' + cli_loop._line[6:]
        elif cli_loop._line[0:5] == 'exit':
            # Clean up
            try:
                world_viewer.exited = True
            except: pass
            if running_fsm:
                running_fsm.stop()
            RUNNING=False
        else:
            cli_loop._line = 'ans=' + cli_loop._line
        try:
            cli_globals['charger'] = robot.world.charger  # charger may have appeared
            exec(cli_loop._line, cli_globals)
            if cli_loop._do_await:
                print("Can't use await outside of an async def.")
                ans = None # ans = await ans
            if not ans is None:
                print(ans,end='\n\n')
        except KeyboardInterrupt:
            print('Keyboard interrupt!')
        except SystemExit:
            print('Type exit() again to exit Python.')
            RUNNING = False
        except Exception:
            traceback.print_exc()
            print()

VERBOSE = False  # True if we want all the log messages


def suppress_filter(log_record):
    message = log_record.msg if isinstance(log_record.msg, str) else repr(log_record.msg)
    if VERBOSE:
        return True
    if log_record.levelno == logging.ERROR and \
       message.startswith("Received a custom object type:"):
        return False
    if message.startswith("Defined: ") and \
       len(log_record.args) > 0 and \
           isinstance(log_record.args[0], cozmo.objects.CustomObject):
        return False
    if message.startswith("Invalidating pose for") or \
       message.startswith("Robot delocalized") or \
       message.startswith("ObjectPowerLevel event") or \
       message.startswith("ObjectConnectionState event") or \
       message.startswith("Object connected") or \
       message.startswith("Object disconnected"):
        return False
    return True

logging_is_setup = False


def start_connection():
    global logging_is_setup, VERBOSE
    if not logging_is_setup:
        cozmo.setup_basic_logging()
        logging_is_setup = True
    if "VERBOSE" in sys.argv:
        VERBOSE = True
    cozmo.logger.addFilter(suppress_filter)
    cozmo.robot.Robot.drive_off_charger_on_connect = False
    connection_error = None

    try:
        if len(sys.argv) >= 2:
            if sys.argv[1] == "TK_VIEWER":
                cozmo.connect_with_tkviewer(run)
            else:
                print("\nUnrecognized argument '%s'. Use TK_VIEWER instead.\n" % sys.argv[1])
                cozmo.connect(run)
        else:
            cozmo.connect(run)
    except cozmo.exceptions.SDKVersionMismatch as e:
        print('\n\n***** SDK version mismatch:',e,'\n\n')
        connection_error = e
    except cozmo.ConnectionError as e:
        connection_error = e
    if connection_error is not None:
        sys.exit("A connection error occurred: %s" % connection_error)


def process_interrupt():
    robot.stop_all_motors()
    running_fsm = cozmo_fsm.program.running_fsm
    if running_fsm and running_fsm.running:
        print('\nKeyboardInterrupt: stopping', running_fsm.name)
        running_fsm.stop()
    else:
        print("\nKeyboardInterrupt. Type 'exit' to exit.")


if __name__ == '__main__':
    start_connection()