Clean up the code
Merge pull request #11 from s-mage/patch-1

Clean code a bit