Refactored some code and did more cleanup