Refactoring and cleaning code