Cleaned up code

Made the code easier to follow