#!/usr/bin/env python
import os
import sys

import matplotlib.pylab as plt

from .builder import create_trajectory
from .util import print_error_message, print_usage
from .util import print_files_in_path, get_trajectories_in_path


def plot_trajectory(name):
    STEPS = 600
    DELTA = 1 if name != 'linear' else 0.1
    trajectory = create_trajectory(name, STEPS)

    x = [trajectory.get_position_at(i * DELTA).x for i in range(STEPS)]
    y = [trajectory.get_position_at(i * DELTA).y for i in range(STEPS)]

    trajectory_fig, trajectory_plot = plt.subplots(1, 1)
    trajectory_plot.plot(x, y, label='trajectory', lw=3)
    trajectory_plot.set_title(name.title() + ' Trajectory', fontsize=20)
    trajectory_plot.set_xlabel(r'$x{\rm[m]}$', fontsize=18)
    trajectory_plot.set_ylabel(r'$y{\rm[m]}$', fontsize=18)
    trajectory_plot.legend(loc=0)
    trajectory_plot.grid()
    plt.show()


if __name__ == '__main__':
    if len(sys.argv) != 2:
        print_error_message()
        sys.exit(1)

    arg = sys.argv[1]

    if arg == '--help':
        print_usage()
        sys.exit(0)

    path = os.sep.join(__file__.split(os.sep)[:-1])

    if arg == '--list':
        print_files_in_path(path)
        sys.exit(0)

    if arg not in get_trajectories_in_path(path):
        print('Error: trajectory does not exist!')
        sys.exit(2)

    plot_trajectory(arg)

#!/usr/bin/env python
from __future__ import print_function
import matplotlib.pylab as plt
import os

from .builder import create_trajectory

def new_line(f):
    def wrapper(*args, **kwargs):
        print()
        f(*args, **kwargs)
        print()
    return wrapper


def plot_trajectory(name):
    STEPS = 600
    DELTA = 1 if name != 'linear' else 0.1
    trajectory = create_trajectory(name, STEPS)

    x = [trajectory.get_position_at(i * DELTA).x for i in range(STEPS)]
    y = [trajectory.get_position_at(i * DELTA).y for i in range(STEPS)]

    trajectory_fig, trajectory_plot = plt.subplots(1, 1)
    trajectory_plot.plot(x, y, label='trajectory', lw=3)
    trajectory_plot.set_title(name.title() + ' Trajectory', fontsize=20)
    trajectory_plot.set_xlabel(r'$x{\rm[m]}$', fontsize=18)
    trajectory_plot.set_ylabel(r'$y{\rm[m]}$', fontsize=18)
    trajectory_plot.legend(loc=0)
    trajectory_plot.grid()
    plt.show()


class TerminateProgramException(Exception):
    pass


class Menu:
    MENU = """Menu:
1. List Trajectories
2. Plot Trajectory
3. Quit
        """

    def __init__(self):
        self.trajectories = self._get_trajectories()
        self.choices = {
            '1': self.list_trajectories,
            '2': self.plot_trajectory,
            '3': self.quit
        }

    @staticmethod
    def _get_trajectories():
        path = os.sep.join(__file__.split(os.sep)[:-1])
        trajectories = [file.replace('_trajectory.py', '')
                        for file in os.listdir(path)
                        if '_trajectory' in file and '.pyc' not in file]

        return {str(i + 1): trajectory
                for i, trajectory in enumerate(trajectories)}

    def run(self):
        while True:
            print(self.MENU)
            choice = raw_input('Enter an option: ')

            try:
                self.choices[choice]()
            except KeyError:
                print('Error: "{}" is not a valid option.'.format(choice))
            except TerminateProgramException:
                break

    @new_line
    def list_trajectories(self):
        print('Trajectories:')
        for k, v in sorted(self.trajectories.items()):
            print('{}. {}'.format(k, v))

    @new_line
    def plot_trajectory(self):
        while True:
            self.list_trajectories()
            choice = raw_input('What trajectory do you want plot? ')
            try:
                trajectory = self.trajectories[choice]
            except KeyError:
                print('Invalid option.')
            else:
                plot_trajectory(trajectory)
                break

    def quit(self):
        raise TerminateProgramException()


if __name__ == '__main__':
    Menu().run()

Implement menu for trajectory module