Refactoring code to make it cleaner