Code refactoring

Cleaned up some of the code