simple refactoring and a few bug fixes