Bug fixes and some code reorganization