Refactoring and code clean-up