Clean up docs and improve interface