Code refactoring and small fixes