clean up

Clean up the code and add some more documentation