Some code cleanup
Refactored and cleaned up a little