Bug fixes and code clean up