Cleaning & refactoring code