Closing issue #5 and refactor