Refactoring and bug fixed