code refactoring and clean up