Clean up the code a bit and add some more comments