Code cleanup

Cleaned up the code a bit