Important fixes to refactoring