Cleaning up code, adding documentation