Bug fixes and a little clean up