Cleaned up the code a bit for the Main.java file