Cleanup and refactor
Docstring fixes