Refactoring code
Fixes #1