Cleaned up main.py a little