Bug fixes and cleanup