Merge pull request #18 from fay-jai/refactor

Code clean up