Minor fixes and clean up