Code refactoring. Minor bug fixes (#4, #5)