bug fixes
fixes and code refactoring