Refactoring code to make it look good