Simplify and restructure code.

I did a little exercise in restructuring your code :).

Feel free to take/leave whichever parts you find useful.