Code refactoring, clean up and bug fixing