Minor fixes and code refactoring into different modules