Few fixes and improvements