Make changes based on code reviews and update docs