Cleanup and some bug fixes