Clean up code
More documentation