Reorganization of code.

I saw your post on reddit, felt like killing some time, and refactored your code for you. Feel free to accept or reject as you like.