Various bug fixes and code clean up