refactoring and code clean-up/fixed