Misc. bug fixes and clean up