Cleaned up code in Main Class