Finish code refactoring

Fixes #5 and #9