Major code refactoring and a few minor bug fixes