Refactored some code
Merge pull request #1 from Sharlock93/refactor

Refactored some code