Few bug fixes and improvements