Refactored some of the code.
Merge branch 'hotfix/refactor'