Cleanup and added docs