Code refactoring and clean up