Various bug fixes and improvements