Refactoring code, correction bug