Refactoring

Minor updates