Cleaned up and modified some code
Merge pull request #1 from T145/master

Very simplistic refinement