General cleaning up of code and refactoring