Refactoring and small fixes