Refactoring code (udpate)