Refactoring code v2