Cleanup and doc comments